radiologist-core 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- radiologist_core-0.1.0/LICENSE +21 -0
- radiologist_core-0.1.0/PKG-INFO +268 -0
- radiologist_core-0.1.0/README.md +220 -0
- radiologist_core-0.1.0/pyproject.toml +90 -0
- radiologist_core-0.1.0/pyproject.toml.orig +79 -0
- radiologist_core-0.1.0/src/radiologist/core/__init__.py +43 -0
- radiologist_core-0.1.0/src/radiologist/core/callbacks/__init__.py +35 -0
- radiologist_core-0.1.0/src/radiologist/core/callbacks/attribution.py +577 -0
- radiologist_core-0.1.0/src/radiologist/core/callbacks/best_metric.py +90 -0
- radiologist_core-0.1.0/src/radiologist/core/callbacks/onnx_export.py +98 -0
- radiologist_core-0.1.0/src/radiologist/core/callbacks/wandb_summary.py +69 -0
- radiologist_core-0.1.0/src/radiologist/core/configs/__init__.py +23 -0
- radiologist_core-0.1.0/src/radiologist/core/configs/callbacks/attribution.yaml +8 -0
- radiologist_core-0.1.0/src/radiologist/core/configs/callbacks/best_metric.yaml +5 -0
- radiologist_core-0.1.0/src/radiologist/core/configs/callbacks/default.yaml +8 -0
- radiologist_core-0.1.0/src/radiologist/core/configs/callbacks/lr_monitor.yaml +4 -0
- radiologist_core-0.1.0/src/radiologist/core/configs/callbacks/model_checkpoint.yaml +16 -0
- radiologist_core-0.1.0/src/radiologist/core/configs/callbacks/onnx_export.yaml +7 -0
- radiologist_core-0.1.0/src/radiologist/core/configs/callbacks/wandb_summary.yaml +5 -0
- radiologist_core-0.1.0/src/radiologist/core/configs/datamodule/default.yaml +56 -0
- radiologist_core-0.1.0/src/radiologist/core/configs/debug/barebones.yaml +18 -0
- radiologist_core-0.1.0/src/radiologist/core/configs/debug/base.yaml +31 -0
- radiologist_core-0.1.0/src/radiologist/core/configs/debug/fast_dev_run.yaml +13 -0
- radiologist_core-0.1.0/src/radiologist/core/configs/debug/limit.yaml +23 -0
- radiologist_core-0.1.0/src/radiologist/core/configs/debug/overfit.yaml +20 -0
- radiologist_core-0.1.0/src/radiologist/core/configs/eval.yaml +24 -0
- radiologist_core-0.1.0/src/radiologist/core/configs/extras.yaml +5 -0
- radiologist_core-0.1.0/src/radiologist/core/configs/hydra/default.yaml +9 -0
- radiologist_core-0.1.0/src/radiologist/core/configs/loggers/wandb.yaml +17 -0
- radiologist_core-0.1.0/src/radiologist/core/configs/module/loss/focal_loss.yaml +8 -0
- radiologist_core-0.1.0/src/radiologist/core/configs/module/metric/fbeta_score.yaml +16 -0
- radiologist_core-0.1.0/src/radiologist/core/configs/module/optimizer/adamw.yaml +8 -0
- radiologist_core-0.1.0/src/radiologist/core/configs/module/resnet50.yaml +7 -0
- radiologist_core-0.1.0/src/radiologist/core/configs/module/scheduler/sequential.yaml +8 -0
- radiologist_core-0.1.0/src/radiologist/core/configs/paths.yaml +7 -0
- radiologist_core-0.1.0/src/radiologist/core/configs/strategy/auto.yaml +2 -0
- radiologist_core-0.1.0/src/radiologist/core/configs/train.yaml +28 -0
- radiologist_core-0.1.0/src/radiologist/core/configs/trainer.yaml +14 -0
- radiologist_core-0.1.0/src/radiologist/core/data/__init__.py +27 -0
- radiologist_core-0.1.0/src/radiologist/core/data/datamodule.py +342 -0
- radiologist_core-0.1.0/src/radiologist/core/data/dtypes.py +34 -0
- radiologist_core-0.1.0/src/radiologist/core/data/shards.py +125 -0
- radiologist_core-0.1.0/src/radiologist/core/losses.py +100 -0
- radiologist_core-0.1.0/src/radiologist/core/module.py +333 -0
- radiologist_core-0.1.0/src/radiologist/core/registry/__init__.py +27 -0
- radiologist_core-0.1.0/src/radiologist/core/registry/export.py +190 -0
- radiologist_core-0.1.0/src/radiologist/core/resume.py +91 -0
- radiologist_core-0.1.0/src/radiologist/core/train.py +180 -0
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 @CedrickArmel
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
|
@@ -0,0 +1,268 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: radiologist-core
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Modeling library for the radiologist pipeline
|
|
5
|
+
Keywords: chest-xray,medical-imaging,machine-learning,pytorch
|
|
6
|
+
Author: Cédrick-Armel YEBOUET
|
|
7
|
+
License-Expression: MIT
|
|
8
|
+
License-File: LICENSE
|
|
9
|
+
Classifier: Development Status :: 3 - Alpha
|
|
10
|
+
Classifier: Intended Audience :: Science/Research
|
|
11
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
12
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
13
|
+
Requires-Dist: pillow>=12.2.0
|
|
14
|
+
Requires-Dist: radiologist-etl>=0.1.0
|
|
15
|
+
Requires-Dist: radiologist-registry>=0.1.0
|
|
16
|
+
Requires-Dist: radiologist-utils>=0.1.0
|
|
17
|
+
Requires-Dist: torch>=2.12.0
|
|
18
|
+
Requires-Dist: torchvision>=0.27.0
|
|
19
|
+
Requires-Dist: webdataset>=1.0.0
|
|
20
|
+
Requires-Dist: lightning>=2.6.0
|
|
21
|
+
Requires-Dist: torchmetrics>=1.9.0
|
|
22
|
+
Requires-Dist: hydra-core>=1.3.0
|
|
23
|
+
Requires-Dist: omegaconf>=2.3.0
|
|
24
|
+
Requires-Dist: fsspec>=2026.4.0
|
|
25
|
+
Requires-Dist: captum>=0.9.0 ; extra == 'all'
|
|
26
|
+
Requires-Dist: gcsfs>=2026.4.0 ; extra == 'all'
|
|
27
|
+
Requires-Dist: onnx>=1.21.0 ; extra == 'all'
|
|
28
|
+
Requires-Dist: onnxruntime~=1.23.2 ; extra == 'all'
|
|
29
|
+
Requires-Dist: onnxscript>=0.7.0 ; extra == 'all'
|
|
30
|
+
Requires-Dist: wandb>=0.27.0 ; extra == 'all'
|
|
31
|
+
Requires-Dist: gcsfs>=2026.4.0 ; extra == 'gcs'
|
|
32
|
+
Requires-Dist: onnx>=1.21.0 ; extra == 'onnx-export'
|
|
33
|
+
Requires-Dist: onnxruntime~=1.23.2 ; extra == 'onnx-export'
|
|
34
|
+
Requires-Dist: onnxscript>=0.7.0 ; extra == 'onnx-export'
|
|
35
|
+
Requires-Dist: wandb>=0.27.0 ; extra == 'wandb'
|
|
36
|
+
Requires-Dist: captum>=0.9.0 ; extra == 'xai'
|
|
37
|
+
Requires-Python: >=3.10
|
|
38
|
+
Project-URL: Homepage, https://github.com/CedrickArmel/radiologist
|
|
39
|
+
Project-URL: Repository, https://github.com/CedrickArmel/radiologist
|
|
40
|
+
Project-URL: Documentation, https://cedrickarmel.github.io/radiologist/
|
|
41
|
+
Project-URL: Issues, https://github.com/CedrickArmel/radiologist/issues
|
|
42
|
+
Provides-Extra: all
|
|
43
|
+
Provides-Extra: gcs
|
|
44
|
+
Provides-Extra: onnx-export
|
|
45
|
+
Provides-Extra: wandb
|
|
46
|
+
Provides-Extra: xai
|
|
47
|
+
Description-Content-Type: text/markdown
|
|
48
|
+
|
|
49
|
+
# radiologist-core
|
|
50
|
+
|
|
51
|
+
[](https://github.com/CedrickArmel/radiologist/actions/workflows/ci.yml)
|
|
52
|
+
[](https://codecov.io/gh/CedrickArmel/radiologist)
|
|
53
|
+
[](https://pypi.org/project/radiologist-core/)
|
|
54
|
+

|
|
55
|
+
|
|
56
|
+
The ML engine. Provides the Lightning training loop, streaming data module, focal loss, GradCAM attribution, and W&B model registry integration for chest X-ray classification.
|
|
57
|
+
|
|
58
|
+
## Business context
|
|
59
|
+
|
|
60
|
+
Once clean, labelled WebDataset shards exist (produced by `radiologist-etl`), this package trains a classifier capable of distinguishing:
|
|
61
|
+
|
|
62
|
+
- **Healthy** lung (no finding)
|
|
63
|
+
- **Viral** pneumonia / COVID (infectious infiltrates)
|
|
64
|
+
- **Opacity** (other lung opacities)
|
|
65
|
+
|
|
66
|
+
The trained model is exported as two ONNX artefacts — one for deterministic inference with visual explanations (GradCAM), one for uncertainty-aware inference (MC-Dropout) — and linked to a W&B Model Registry collection for downstream consumption by an inference or serving layer.
|
|
67
|
+
|
|
68
|
+
## Key capabilities
|
|
69
|
+
|
|
70
|
+
| Capability | What it enables |
|
|
71
|
+
|---|---|
|
|
72
|
+
| Class-balanced streaming | Trains on imbalanced clinical data without manual oversampling scripts |
|
|
73
|
+
| Prior-calibrated bias init | Logit outputs are calibrated to class prevalence from the first step |
|
|
74
|
+
| Focal Loss | Focuses learning on hard, misclassified examples rather than easy majority samples |
|
|
75
|
+
| GradCAM attribution | Produces heatmaps showing which lung region drove the prediction |
|
|
76
|
+
| MC-Dropout ONNX export | Enables uncertainty estimation at inference time (multiple stochastic passes) |
|
|
77
|
+
| HPO-friendly best-metric tracking | Optuna / W&B Sweeps receive the best epoch score, not the last |
|
|
78
|
+
|
|
79
|
+
## Package layout
|
|
80
|
+
|
|
81
|
+
```
|
|
82
|
+
radiologist-core/src/radiologist/core/
|
|
83
|
+
├── train.py # Hydra entry point and train() function
|
|
84
|
+
├── module.py # LModule (LightningModule)
|
|
85
|
+
├── losses.py # FocalLoss
|
|
86
|
+
├── data/
|
|
87
|
+
│ ├── datamodule.py # WebDatasetDataModule
|
|
88
|
+
│ └── shards.py # shard discovery helpers
|
|
89
|
+
├── callbacks/
|
|
90
|
+
│ ├── attribution.py # GradCAM + Integrated Gradients
|
|
91
|
+
│ ├── best_metric.py # best-epoch metric tracking
|
|
92
|
+
│ └── wandb_summary.py
|
|
93
|
+
├── registry/
|
|
94
|
+
│ └── export.py # export_onnx (checkpoint → deterministic + MC-Dropout ONNX)
|
|
95
|
+
└── configs/ # Hydra config tree
|
|
96
|
+
```
|
|
97
|
+
|
|
98
|
+
## Training a model
|
|
99
|
+
|
|
100
|
+
### Prerequisites
|
|
101
|
+
|
|
102
|
+
- Shards on GCS (or local) produced by `radiologist-etl`
|
|
103
|
+
- W&B API key in the environment (`WANDB_API_KEY`)
|
|
104
|
+
|
|
105
|
+
### Quick start
|
|
106
|
+
|
|
107
|
+
```bash
|
|
108
|
+
cd radiologist-core
|
|
109
|
+
uv run --active python -m radiologist.core.train
|
|
110
|
+
```
|
|
111
|
+
|
|
112
|
+
The defaults load from `src/radiologist/core/configs/train.yaml`. Override on the command line:
|
|
113
|
+
|
|
114
|
+
```bash
|
|
115
|
+
uv run --active python -m radiologist.core.train \
|
|
116
|
+
trainer.max_epochs=30 \
|
|
117
|
+
datamodule.shard_root=gs://my-bucket/shards/ \
|
|
118
|
+
datamodule.split_manifest_uri=gs://my-bucket/manifests/manifest-abc123.jsonl \
|
|
119
|
+
seed=42
|
|
120
|
+
```
|
|
121
|
+
|
|
122
|
+
### Evaluation only
|
|
123
|
+
|
|
124
|
+
```bash
|
|
125
|
+
uv run --active python -m radiologist.core.train \
|
|
126
|
+
--config-name eval \
|
|
127
|
+
ckpt_path=/path/to/checkpoint.ckpt
|
|
128
|
+
```
|
|
129
|
+
|
|
130
|
+
## Core components
|
|
131
|
+
|
|
132
|
+
### `LModule`
|
|
133
|
+
|
|
134
|
+
A backbone-agnostic `LightningModule`. Pass any `nn.Module` as `net`.
|
|
135
|
+
|
|
136
|
+
```python
|
|
137
|
+
from radiologist.core import LModule
|
|
138
|
+
from torchvision.models import resnet50
|
|
139
|
+
|
|
140
|
+
module = LModule(
|
|
141
|
+
net=resnet50(num_classes=3),
|
|
142
|
+
loss=FocalLoss(gamma=2, alpha=1),
|
|
143
|
+
metric=MulticlassFBetaScore(num_classes=3, beta=1.0, average="macro"),
|
|
144
|
+
optimizer=partial(AdamW, lr=1e-3, weight_decay=1e-2),
|
|
145
|
+
scheduler=partial(sequential_scheduler, ...),
|
|
146
|
+
trainable_layers=None, # None = full re-init; list of dot-paths = fine-tune
|
|
147
|
+
priors=None, # overridden from datamodule at setup time
|
|
148
|
+
)
|
|
149
|
+
```
|
|
150
|
+
|
|
151
|
+
On `setup('fit')`, `LModule` does one of two things depending on `trainable_layers`:
|
|
152
|
+
|
|
153
|
+
- **`None`** — reinitialise all weights (Kaiming normal on Conv, Xavier on Linear). Use when training from scratch.
|
|
154
|
+
- **`["layer4", "fc"]`** — freeze all parameters, then selectively unfreeze by dot-path. Use when fine-tuning a pretrained backbone.
|
|
155
|
+
|
|
156
|
+
In both cases, if class priors are available (from the datamodule), the final `nn.Linear` bias is initialised to `−log(priors)`, giving calibrated starting logits.
|
|
157
|
+
|
|
158
|
+
### `WebDatasetDataModule`
|
|
159
|
+
|
|
160
|
+
Streams images from WebDataset tar shards. Handles class-balanced sampling automatically.
|
|
161
|
+
|
|
162
|
+
```python
|
|
163
|
+
from radiologist.core import WebDatasetDataModule
|
|
164
|
+
|
|
165
|
+
dm = WebDatasetDataModule(
|
|
166
|
+
shard_root="gs://bucket/shards/",
|
|
167
|
+
split_manifest_uri="gs://bucket/manifests/manifest-abc123.jsonl",
|
|
168
|
+
label_map={"normal": "healthy", "pneumonia": "viral", "COVID": "viral"},
|
|
169
|
+
batch_size=32,
|
|
170
|
+
)
|
|
171
|
+
```
|
|
172
|
+
|
|
173
|
+
The `label_map` collapses raw ETL folder names (e.g. `normal`, `COVID`) into model class names (e.g. `healthy`, `viral`). This decouples the dataset's folder structure from the model's output space.
|
|
174
|
+
|
|
175
|
+
Training dataloader: one `wds.WebDataset` pipeline per class with `resampled=True` (infinite streaming), combined via `wds.RandomMix` weighted by inverse class frequency, then unbatch → shuffle → rebatch for global shuffling within an epoch. Validation and test dataloaders are flat sequential pipelines with no resampling.
|
|
176
|
+
|
|
177
|
+
### `FocalLoss`
|
|
178
|
+
|
|
179
|
+
```python
|
|
180
|
+
from radiologist.core import FocalLoss
|
|
181
|
+
|
|
182
|
+
loss = FocalLoss(gamma=2.0, alpha=1.0, use_softmax=True, reduction="mean")
|
|
183
|
+
```
|
|
184
|
+
|
|
185
|
+
Applies softmax to logits, computes `alpha * (1 − pt)^gamma * −log(pt)`. Supports optional one-hot conversion and `mean | sum | none` reductions.
|
|
186
|
+
|
|
187
|
+
### Callbacks
|
|
188
|
+
|
|
189
|
+
#### `AttributionCallback`
|
|
190
|
+
|
|
191
|
+
Computes Layer GradCAM and Integrated Gradients every `every_n_val_epochs` validation epochs and on all test batches. Saves normalised PNGs to `{log_dir}/attributions/` and logs them to W&B. Skipped gracefully when `captum` is not installed.
|
|
192
|
+
|
|
193
|
+
```yaml
|
|
194
|
+
# configs/callbacks/default.yaml
|
|
195
|
+
attribution:
|
|
196
|
+
_target_: radiologist.core.callbacks.AttributionCallback
|
|
197
|
+
target_layer: layer4.1.conv2
|
|
198
|
+
every_n_val_epochs: 5
|
|
199
|
+
n_test_batches: 4
|
|
200
|
+
n_samples_per_batch: 4
|
|
201
|
+
```
|
|
202
|
+
|
|
203
|
+
#### `BestMetricCallback`
|
|
204
|
+
|
|
205
|
+
Writes `best_{monitor}` to `trainer.callback_metrics` after every validation epoch. When training finishes, HPO frameworks (Optuna, W&B Sweeps) read the best epoch score rather than the last.
|
|
206
|
+
|
|
207
|
+
#### `WandbDefineSummaryCallback`
|
|
208
|
+
|
|
209
|
+
Calls `wandb.run.define_metric` at fit start so the W&B run summary panel highlights the best validation score automatically.
|
|
210
|
+
|
|
211
|
+
## Exporting and promoting a trained model
|
|
212
|
+
|
|
213
|
+
`radiologist.core.registry.export_onnx` turns a Lightning checkpoint into two local ONNX files — it has no W&B interaction:
|
|
214
|
+
|
|
215
|
+
```python
|
|
216
|
+
from radiologist.core.registry import export_onnx
|
|
217
|
+
|
|
218
|
+
result = export_onnx(
|
|
219
|
+
ckpt_path="/path/to/checkpoint.ckpt",
|
|
220
|
+
run_id="wandb-run-id",
|
|
221
|
+
input_shape=(1, 1, 224, 224),
|
|
222
|
+
classes=["healthy", "viral", "opacity"],
|
|
223
|
+
cam_target_layer="layer4.1.conv2",
|
|
224
|
+
out_dir="/tmp/onnx-export",
|
|
225
|
+
)
|
|
226
|
+
```
|
|
227
|
+
|
|
228
|
+
This produces:
|
|
229
|
+
|
|
230
|
+
1. **Deterministic** — `_CamWrapper` forward hook returns `(logits, activation)`. Useful for inference with visual explanation.
|
|
231
|
+
2. **MC-Dropout** — `nn.Dropout` layers left in training mode (`TrainingMode.PRESERVE`, no constant folding). Run multiple forward passes and aggregate for uncertainty estimation.
|
|
232
|
+
|
|
233
|
+
Uploading the exported ONNX files as W&B artifacts and linking them into a registry collection is handled by `radiologist-registry` (`WandbRegistry.log_model_artifacts()` then `WandbRegistry.promote()`, or the `radiologist-registry push` / `promote` CLI). See [`radiologist-registry/README.md`](../radiologist-registry/README.md) for the full flow.
|
|
234
|
+
|
|
235
|
+
## Configuration reference
|
|
236
|
+
|
|
237
|
+
The Hydra config tree lives at `src/radiologist/core/configs/`. Key files:
|
|
238
|
+
|
|
239
|
+
| File | Purpose |
|
|
240
|
+
|---|---|
|
|
241
|
+
| `train.yaml` | Root config; wires all sub-configs |
|
|
242
|
+
| `eval.yaml` | Evaluation-only mode (`train: false`, `ckpt_path: ???`) |
|
|
243
|
+
| `trainer.yaml` | Lightning Trainer (precision, gradient clipping, deterministic) |
|
|
244
|
+
| `datamodule/default.yaml` | Shard URIs, label map, transforms, normalisation |
|
|
245
|
+
| `module/resnet50.yaml` | ResNet-50 backbone, `num_classes` interpolated from datamodule |
|
|
246
|
+
| `module/loss/focal_loss.yaml` | γ=2, α=1, softmax, mean |
|
|
247
|
+
| `module/optimizer/adamw.yaml` | AdamW, lr=1e-3, weight_decay=1e-2 |
|
|
248
|
+
| `module/scheduler/sequential.yaml` | Linear warmup 500 steps → cosine annealing 10 000 steps |
|
|
249
|
+
| `module/metric/fbeta_score.yaml` | Macro F1 (`MulticlassFBetaScore`, β=1) |
|
|
250
|
+
| `callbacks/default.yaml` | BestMetric, WandbSummary, Attribution, ModelCheckpoint, LRMonitor |
|
|
251
|
+
|
|
252
|
+
## Optional extras
|
|
253
|
+
|
|
254
|
+
Install the `onnx-export` extra for ONNX export and W&B registry features:
|
|
255
|
+
|
|
256
|
+
```bash
|
|
257
|
+
uv add --active "radiologist-core[onnx-export]"
|
|
258
|
+
```
|
|
259
|
+
|
|
260
|
+
Adds: `onnx`, `onnxruntime`, `onnxscript`.
|
|
261
|
+
|
|
262
|
+
## Dependencies
|
|
263
|
+
|
|
264
|
+
Core: `radiologist-utils`, `torch`, `torchvision`, `lightning`, `webdataset`, `torchmetrics`, `wandb`, `hydra-core`.
|
|
265
|
+
|
|
266
|
+
Optional (`onnx-export`): `onnx`, `onnxruntime`, `onnxscript`.
|
|
267
|
+
|
|
268
|
+
Optional (`attribution`): `captum`.
|
|
@@ -0,0 +1,220 @@
|
|
|
1
|
+
# radiologist-core
|
|
2
|
+
|
|
3
|
+
[](https://github.com/CedrickArmel/radiologist/actions/workflows/ci.yml)
|
|
4
|
+
[](https://codecov.io/gh/CedrickArmel/radiologist)
|
|
5
|
+
[](https://pypi.org/project/radiologist-core/)
|
|
6
|
+

|
|
7
|
+
|
|
8
|
+
The ML engine. Provides the Lightning training loop, streaming data module, focal loss, GradCAM attribution, and W&B model registry integration for chest X-ray classification.
|
|
9
|
+
|
|
10
|
+
## Business context
|
|
11
|
+
|
|
12
|
+
Once clean, labelled WebDataset shards exist (produced by `radiologist-etl`), this package trains a classifier capable of distinguishing:
|
|
13
|
+
|
|
14
|
+
- **Healthy** lung (no finding)
|
|
15
|
+
- **Viral** pneumonia / COVID (infectious infiltrates)
|
|
16
|
+
- **Opacity** (other lung opacities)
|
|
17
|
+
|
|
18
|
+
The trained model is exported as two ONNX artefacts — one for deterministic inference with visual explanations (GradCAM), one for uncertainty-aware inference (MC-Dropout) — and linked to a W&B Model Registry collection for downstream consumption by an inference or serving layer.
|
|
19
|
+
|
|
20
|
+
## Key capabilities
|
|
21
|
+
|
|
22
|
+
| Capability | What it enables |
|
|
23
|
+
|---|---|
|
|
24
|
+
| Class-balanced streaming | Trains on imbalanced clinical data without manual oversampling scripts |
|
|
25
|
+
| Prior-calibrated bias init | Logit outputs are calibrated to class prevalence from the first step |
|
|
26
|
+
| Focal Loss | Focuses learning on hard, misclassified examples rather than easy majority samples |
|
|
27
|
+
| GradCAM attribution | Produces heatmaps showing which lung region drove the prediction |
|
|
28
|
+
| MC-Dropout ONNX export | Enables uncertainty estimation at inference time (multiple stochastic passes) |
|
|
29
|
+
| HPO-friendly best-metric tracking | Optuna / W&B Sweeps receive the best epoch score, not the last |
|
|
30
|
+
|
|
31
|
+
## Package layout
|
|
32
|
+
|
|
33
|
+
```
|
|
34
|
+
radiologist-core/src/radiologist/core/
|
|
35
|
+
├── train.py # Hydra entry point and train() function
|
|
36
|
+
├── module.py # LModule (LightningModule)
|
|
37
|
+
├── losses.py # FocalLoss
|
|
38
|
+
├── data/
|
|
39
|
+
│ ├── datamodule.py # WebDatasetDataModule
|
|
40
|
+
│ └── shards.py # shard discovery helpers
|
|
41
|
+
├── callbacks/
|
|
42
|
+
│ ├── attribution.py # GradCAM + Integrated Gradients
|
|
43
|
+
│ ├── best_metric.py # best-epoch metric tracking
|
|
44
|
+
│ └── wandb_summary.py
|
|
45
|
+
├── registry/
|
|
46
|
+
│ └── export.py # export_onnx (checkpoint → deterministic + MC-Dropout ONNX)
|
|
47
|
+
└── configs/ # Hydra config tree
|
|
48
|
+
```
|
|
49
|
+
|
|
50
|
+
## Training a model
|
|
51
|
+
|
|
52
|
+
### Prerequisites
|
|
53
|
+
|
|
54
|
+
- Shards on GCS (or local) produced by `radiologist-etl`
|
|
55
|
+
- W&B API key in the environment (`WANDB_API_KEY`)
|
|
56
|
+
|
|
57
|
+
### Quick start
|
|
58
|
+
|
|
59
|
+
```bash
|
|
60
|
+
cd radiologist-core
|
|
61
|
+
uv run --active python -m radiologist.core.train
|
|
62
|
+
```
|
|
63
|
+
|
|
64
|
+
The defaults load from `src/radiologist/core/configs/train.yaml`. Override on the command line:
|
|
65
|
+
|
|
66
|
+
```bash
|
|
67
|
+
uv run --active python -m radiologist.core.train \
|
|
68
|
+
trainer.max_epochs=30 \
|
|
69
|
+
datamodule.shard_root=gs://my-bucket/shards/ \
|
|
70
|
+
datamodule.split_manifest_uri=gs://my-bucket/manifests/manifest-abc123.jsonl \
|
|
71
|
+
seed=42
|
|
72
|
+
```
|
|
73
|
+
|
|
74
|
+
### Evaluation only
|
|
75
|
+
|
|
76
|
+
```bash
|
|
77
|
+
uv run --active python -m radiologist.core.train \
|
|
78
|
+
--config-name eval \
|
|
79
|
+
ckpt_path=/path/to/checkpoint.ckpt
|
|
80
|
+
```
|
|
81
|
+
|
|
82
|
+
## Core components
|
|
83
|
+
|
|
84
|
+
### `LModule`
|
|
85
|
+
|
|
86
|
+
A backbone-agnostic `LightningModule`. Pass any `nn.Module` as `net`.
|
|
87
|
+
|
|
88
|
+
```python
|
|
89
|
+
from radiologist.core import LModule
|
|
90
|
+
from torchvision.models import resnet50
|
|
91
|
+
|
|
92
|
+
module = LModule(
|
|
93
|
+
net=resnet50(num_classes=3),
|
|
94
|
+
loss=FocalLoss(gamma=2, alpha=1),
|
|
95
|
+
metric=MulticlassFBetaScore(num_classes=3, beta=1.0, average="macro"),
|
|
96
|
+
optimizer=partial(AdamW, lr=1e-3, weight_decay=1e-2),
|
|
97
|
+
scheduler=partial(sequential_scheduler, ...),
|
|
98
|
+
trainable_layers=None, # None = full re-init; list of dot-paths = fine-tune
|
|
99
|
+
priors=None, # overridden from datamodule at setup time
|
|
100
|
+
)
|
|
101
|
+
```
|
|
102
|
+
|
|
103
|
+
On `setup('fit')`, `LModule` does one of two things depending on `trainable_layers`:
|
|
104
|
+
|
|
105
|
+
- **`None`** — reinitialise all weights (Kaiming normal on Conv, Xavier on Linear). Use when training from scratch.
|
|
106
|
+
- **`["layer4", "fc"]`** — freeze all parameters, then selectively unfreeze by dot-path. Use when fine-tuning a pretrained backbone.
|
|
107
|
+
|
|
108
|
+
In both cases, if class priors are available (from the datamodule), the final `nn.Linear` bias is initialised to `−log(priors)`, giving calibrated starting logits.
|
|
109
|
+
|
|
110
|
+
### `WebDatasetDataModule`
|
|
111
|
+
|
|
112
|
+
Streams images from WebDataset tar shards. Handles class-balanced sampling automatically.
|
|
113
|
+
|
|
114
|
+
```python
|
|
115
|
+
from radiologist.core import WebDatasetDataModule
|
|
116
|
+
|
|
117
|
+
dm = WebDatasetDataModule(
|
|
118
|
+
shard_root="gs://bucket/shards/",
|
|
119
|
+
split_manifest_uri="gs://bucket/manifests/manifest-abc123.jsonl",
|
|
120
|
+
label_map={"normal": "healthy", "pneumonia": "viral", "COVID": "viral"},
|
|
121
|
+
batch_size=32,
|
|
122
|
+
)
|
|
123
|
+
```
|
|
124
|
+
|
|
125
|
+
The `label_map` collapses raw ETL folder names (e.g. `normal`, `COVID`) into model class names (e.g. `healthy`, `viral`). This decouples the dataset's folder structure from the model's output space.
|
|
126
|
+
|
|
127
|
+
Training dataloader: one `wds.WebDataset` pipeline per class with `resampled=True` (infinite streaming), combined via `wds.RandomMix` weighted by inverse class frequency, then unbatch → shuffle → rebatch for global shuffling within an epoch. Validation and test dataloaders are flat sequential pipelines with no resampling.
|
|
128
|
+
|
|
129
|
+
### `FocalLoss`
|
|
130
|
+
|
|
131
|
+
```python
|
|
132
|
+
from radiologist.core import FocalLoss
|
|
133
|
+
|
|
134
|
+
loss = FocalLoss(gamma=2.0, alpha=1.0, use_softmax=True, reduction="mean")
|
|
135
|
+
```
|
|
136
|
+
|
|
137
|
+
Applies softmax to logits, computes `alpha * (1 − pt)^gamma * −log(pt)`. Supports optional one-hot conversion and `mean | sum | none` reductions.
|
|
138
|
+
|
|
139
|
+
### Callbacks
|
|
140
|
+
|
|
141
|
+
#### `AttributionCallback`
|
|
142
|
+
|
|
143
|
+
Computes Layer GradCAM and Integrated Gradients every `every_n_val_epochs` validation epochs and on all test batches. Saves normalised PNGs to `{log_dir}/attributions/` and logs them to W&B. Skipped gracefully when `captum` is not installed.
|
|
144
|
+
|
|
145
|
+
```yaml
|
|
146
|
+
# configs/callbacks/default.yaml
|
|
147
|
+
attribution:
|
|
148
|
+
_target_: radiologist.core.callbacks.AttributionCallback
|
|
149
|
+
target_layer: layer4.1.conv2
|
|
150
|
+
every_n_val_epochs: 5
|
|
151
|
+
n_test_batches: 4
|
|
152
|
+
n_samples_per_batch: 4
|
|
153
|
+
```
|
|
154
|
+
|
|
155
|
+
#### `BestMetricCallback`
|
|
156
|
+
|
|
157
|
+
Writes `best_{monitor}` to `trainer.callback_metrics` after every validation epoch. When training finishes, HPO frameworks (Optuna, W&B Sweeps) read the best epoch score rather than the last.
|
|
158
|
+
|
|
159
|
+
#### `WandbDefineSummaryCallback`
|
|
160
|
+
|
|
161
|
+
Calls `wandb.run.define_metric` at fit start so the W&B run summary panel highlights the best validation score automatically.
|
|
162
|
+
|
|
163
|
+
## Exporting and promoting a trained model
|
|
164
|
+
|
|
165
|
+
`radiologist.core.registry.export_onnx` turns a Lightning checkpoint into two local ONNX files — it has no W&B interaction:
|
|
166
|
+
|
|
167
|
+
```python
|
|
168
|
+
from radiologist.core.registry import export_onnx
|
|
169
|
+
|
|
170
|
+
result = export_onnx(
|
|
171
|
+
ckpt_path="/path/to/checkpoint.ckpt",
|
|
172
|
+
run_id="wandb-run-id",
|
|
173
|
+
input_shape=(1, 1, 224, 224),
|
|
174
|
+
classes=["healthy", "viral", "opacity"],
|
|
175
|
+
cam_target_layer="layer4.1.conv2",
|
|
176
|
+
out_dir="/tmp/onnx-export",
|
|
177
|
+
)
|
|
178
|
+
```
|
|
179
|
+
|
|
180
|
+
This produces:
|
|
181
|
+
|
|
182
|
+
1. **Deterministic** — `_CamWrapper` forward hook returns `(logits, activation)`. Useful for inference with visual explanation.
|
|
183
|
+
2. **MC-Dropout** — `nn.Dropout` layers left in training mode (`TrainingMode.PRESERVE`, no constant folding). Run multiple forward passes and aggregate for uncertainty estimation.
|
|
184
|
+
|
|
185
|
+
Uploading the exported ONNX files as W&B artifacts and linking them into a registry collection is handled by `radiologist-registry` (`WandbRegistry.log_model_artifacts()` then `WandbRegistry.promote()`, or the `radiologist-registry push` / `promote` CLI). See [`radiologist-registry/README.md`](../radiologist-registry/README.md) for the full flow.
|
|
186
|
+
|
|
187
|
+
## Configuration reference
|
|
188
|
+
|
|
189
|
+
The Hydra config tree lives at `src/radiologist/core/configs/`. Key files:
|
|
190
|
+
|
|
191
|
+
| File | Purpose |
|
|
192
|
+
|---|---|
|
|
193
|
+
| `train.yaml` | Root config; wires all sub-configs |
|
|
194
|
+
| `eval.yaml` | Evaluation-only mode (`train: false`, `ckpt_path: ???`) |
|
|
195
|
+
| `trainer.yaml` | Lightning Trainer (precision, gradient clipping, deterministic) |
|
|
196
|
+
| `datamodule/default.yaml` | Shard URIs, label map, transforms, normalisation |
|
|
197
|
+
| `module/resnet50.yaml` | ResNet-50 backbone, `num_classes` interpolated from datamodule |
|
|
198
|
+
| `module/loss/focal_loss.yaml` | γ=2, α=1, softmax, mean |
|
|
199
|
+
| `module/optimizer/adamw.yaml` | AdamW, lr=1e-3, weight_decay=1e-2 |
|
|
200
|
+
| `module/scheduler/sequential.yaml` | Linear warmup 500 steps → cosine annealing 10 000 steps |
|
|
201
|
+
| `module/metric/fbeta_score.yaml` | Macro F1 (`MulticlassFBetaScore`, β=1) |
|
|
202
|
+
| `callbacks/default.yaml` | BestMetric, WandbSummary, Attribution, ModelCheckpoint, LRMonitor |
|
|
203
|
+
|
|
204
|
+
## Optional extras
|
|
205
|
+
|
|
206
|
+
Install the `onnx-export` extra for ONNX export and W&B registry features:
|
|
207
|
+
|
|
208
|
+
```bash
|
|
209
|
+
uv add --active "radiologist-core[onnx-export]"
|
|
210
|
+
```
|
|
211
|
+
|
|
212
|
+
Adds: `onnx`, `onnxruntime`, `onnxscript`.
|
|
213
|
+
|
|
214
|
+
## Dependencies
|
|
215
|
+
|
|
216
|
+
Core: `radiologist-utils`, `torch`, `torchvision`, `lightning`, `webdataset`, `torchmetrics`, `wandb`, `hydra-core`.
|
|
217
|
+
|
|
218
|
+
Optional (`onnx-export`): `onnx`, `onnxruntime`, `onnxscript`.
|
|
219
|
+
|
|
220
|
+
Optional (`attribution`): `captum`.
|
|
@@ -0,0 +1,90 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
build-backend = "uv_build"
|
|
3
|
+
requires = ["uv_build>=0.9.28,<0.10.0"]
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
dependencies = [
|
|
7
|
+
"Pillow>=12.2.0",
|
|
8
|
+
"radiologist-etl>=0.1.0",
|
|
9
|
+
"radiologist-registry>=0.1.0",
|
|
10
|
+
"radiologist-utils>=0.1.0",
|
|
11
|
+
"torch>=2.12.0",
|
|
12
|
+
"torchvision>=0.27.0",
|
|
13
|
+
"webdataset>=1.0.0",
|
|
14
|
+
"lightning>=2.6.0",
|
|
15
|
+
"torchmetrics>=1.9.0",
|
|
16
|
+
"hydra-core>=1.3.0",
|
|
17
|
+
"omegaconf>=2.3.0",
|
|
18
|
+
"fsspec>=2026.4.0",
|
|
19
|
+
]
|
|
20
|
+
name = "radiologist-core"
|
|
21
|
+
version = "0.1.0"
|
|
22
|
+
description = "Modeling library for the radiologist pipeline"
|
|
23
|
+
readme = "README.md"
|
|
24
|
+
license = "MIT"
|
|
25
|
+
license-files = ["LICENSE"]
|
|
26
|
+
keywords = [
|
|
27
|
+
"chest-xray",
|
|
28
|
+
"medical-imaging",
|
|
29
|
+
"machine-learning",
|
|
30
|
+
"pytorch",
|
|
31
|
+
]
|
|
32
|
+
classifiers = [
|
|
33
|
+
"Development Status :: 3 - Alpha",
|
|
34
|
+
"Intended Audience :: Science/Research",
|
|
35
|
+
"Programming Language :: Python :: 3.10",
|
|
36
|
+
"Topic :: Scientific/Engineering :: Artificial Intelligence",
|
|
37
|
+
]
|
|
38
|
+
requires-python = ">=3.10"
|
|
39
|
+
|
|
40
|
+
[[project.authors]]
|
|
41
|
+
name = "Cédrick-Armel YEBOUET"
|
|
42
|
+
|
|
43
|
+
[project.scripts]
|
|
44
|
+
radiologist-core = "radiologist.core.train:main"
|
|
45
|
+
|
|
46
|
+
[project.optional-dependencies]
|
|
47
|
+
all = [
|
|
48
|
+
"captum>=0.9.0",
|
|
49
|
+
"gcsfs>=2026.4.0",
|
|
50
|
+
"onnx>=1.21.0",
|
|
51
|
+
"onnxruntime~=1.23.2",
|
|
52
|
+
"onnxscript>=0.7.0",
|
|
53
|
+
"wandb>=0.27.0",
|
|
54
|
+
]
|
|
55
|
+
gcs = ["gcsfs>=2026.4.0"]
|
|
56
|
+
onnx-export = [
|
|
57
|
+
"onnx>=1.21.0",
|
|
58
|
+
"onnxruntime~=1.23.2",
|
|
59
|
+
"onnxscript>=0.7.0",
|
|
60
|
+
]
|
|
61
|
+
wandb = ["wandb>=0.27.0"]
|
|
62
|
+
xai = ["captum>=0.9.0"]
|
|
63
|
+
|
|
64
|
+
[project.urls]
|
|
65
|
+
Homepage = "https://github.com/CedrickArmel/radiologist"
|
|
66
|
+
Repository = "https://github.com/CedrickArmel/radiologist"
|
|
67
|
+
Documentation = "https://cedrickarmel.github.io/radiologist/"
|
|
68
|
+
Issues = "https://github.com/CedrickArmel/radiologist/issues"
|
|
69
|
+
|
|
70
|
+
[tool.commitizen]
|
|
71
|
+
name = "cz_conventional_commits"
|
|
72
|
+
version_provider = "uv"
|
|
73
|
+
version_scheme = "pep440"
|
|
74
|
+
tag_format = "${version}-radiologist-core"
|
|
75
|
+
ignored_tag_formats = ["${version}-radiologist-*"]
|
|
76
|
+
update_changelog_on_bump = true
|
|
77
|
+
bump_message = "bump: radiologist-core $current_version → $new_version"
|
|
78
|
+
|
|
79
|
+
[tool.uv.build-backend]
|
|
80
|
+
namespace = true
|
|
81
|
+
module-name = "radiologist.core"
|
|
82
|
+
|
|
83
|
+
[tool.uv.sources.radiologist-etl]
|
|
84
|
+
workspace = true
|
|
85
|
+
|
|
86
|
+
[tool.uv.sources.radiologist-registry]
|
|
87
|
+
workspace = true
|
|
88
|
+
|
|
89
|
+
[tool.uv.sources.radiologist-utils]
|
|
90
|
+
workspace = true
|
|
@@ -0,0 +1,79 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
build-backend = "uv_build"
|
|
3
|
+
requires = ["uv_build>=0.9.28,<0.10.0"]
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
dependencies = [
|
|
7
|
+
"Pillow>=12.2.0",
|
|
8
|
+
"radiologist-etl>=0.1.0",
|
|
9
|
+
"radiologist-registry>=0.1.0",
|
|
10
|
+
"radiologist-utils>=0.1.0",
|
|
11
|
+
"torch>=2.12.0",
|
|
12
|
+
"torchvision>=0.27.0",
|
|
13
|
+
"webdataset>=1.0.0",
|
|
14
|
+
"lightning>=2.6.0",
|
|
15
|
+
"torchmetrics>=1.9.0",
|
|
16
|
+
"hydra-core>=1.3.0",
|
|
17
|
+
"omegaconf>=2.3.0",
|
|
18
|
+
"fsspec>=2026.4.0",
|
|
19
|
+
]
|
|
20
|
+
name = "radiologist-core"
|
|
21
|
+
version = "0.1.0"
|
|
22
|
+
description = "Modeling library for the radiologist pipeline"
|
|
23
|
+
readme = "README.md"
|
|
24
|
+
license = "MIT"
|
|
25
|
+
license-files = ["LICENSE"]
|
|
26
|
+
authors = [{ name = "Cédrick-Armel YEBOUET" }]
|
|
27
|
+
keywords = ["chest-xray", "medical-imaging", "machine-learning", "pytorch"]
|
|
28
|
+
classifiers = [
|
|
29
|
+
"Development Status :: 3 - Alpha",
|
|
30
|
+
"Intended Audience :: Science/Research",
|
|
31
|
+
"Programming Language :: Python :: 3.10",
|
|
32
|
+
"Topic :: Scientific/Engineering :: Artificial Intelligence",
|
|
33
|
+
]
|
|
34
|
+
requires-python = ">=3.10"
|
|
35
|
+
|
|
36
|
+
[project.scripts]
|
|
37
|
+
radiologist-core = "radiologist.core.train:main"
|
|
38
|
+
|
|
39
|
+
[project.optional-dependencies]
|
|
40
|
+
all = [
|
|
41
|
+
"captum>=0.9.0",
|
|
42
|
+
"gcsfs>=2026.4.0",
|
|
43
|
+
"onnx>=1.21.0",
|
|
44
|
+
"onnxruntime~=1.23.2",
|
|
45
|
+
"onnxscript>=0.7.0",
|
|
46
|
+
"wandb>=0.27.0",
|
|
47
|
+
]
|
|
48
|
+
gcs = ["gcsfs>=2026.4.0"]
|
|
49
|
+
onnx-export = [
|
|
50
|
+
"onnx>=1.21.0",
|
|
51
|
+
"onnxruntime~=1.23.2",
|
|
52
|
+
"onnxscript>=0.7.0",
|
|
53
|
+
]
|
|
54
|
+
wandb = ["wandb>=0.27.0"]
|
|
55
|
+
xai = ["captum>=0.9.0"]
|
|
56
|
+
|
|
57
|
+
[project.urls]
|
|
58
|
+
Homepage = "https://github.com/CedrickArmel/radiologist"
|
|
59
|
+
Repository = "https://github.com/CedrickArmel/radiologist"
|
|
60
|
+
Documentation = "https://cedrickarmel.github.io/radiologist/"
|
|
61
|
+
Issues = "https://github.com/CedrickArmel/radiologist/issues"
|
|
62
|
+
|
|
63
|
+
[tool.commitizen]
|
|
64
|
+
name = "cz_conventional_commits"
|
|
65
|
+
version_provider = "uv"
|
|
66
|
+
version_scheme = "pep440"
|
|
67
|
+
tag_format = "${version}-radiologist-core"
|
|
68
|
+
ignored_tag_formats = ["${version}-radiologist-*"]
|
|
69
|
+
update_changelog_on_bump = true
|
|
70
|
+
bump_message = "bump: radiologist-core $current_version → $new_version"
|
|
71
|
+
|
|
72
|
+
[tool.uv.build-backend]
|
|
73
|
+
namespace = true
|
|
74
|
+
module-name = "radiologist.core"
|
|
75
|
+
|
|
76
|
+
[tool.uv.sources]
|
|
77
|
+
radiologist-etl = { workspace = true }
|
|
78
|
+
radiologist-registry = { workspace = true }
|
|
79
|
+
radiologist-utils = { workspace = true }
|