kssrikar4 commited on
Commit
2f4b4ff
·
verified ·
1 Parent(s): 03447fc

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +456 -459
README.md CHANGED
@@ -1,459 +1,456 @@
1
- ---
2
- language: en
3
- license: gpl-3.0
4
- tags:
5
- - 3d-vision
6
- - medical-imaging
7
- - deep-learning
8
- - pytorch
9
- - fcd-detection
10
- - brain-mri
11
- model-index:
12
- - name: MobileNetV3-Small-3D-FCD
13
- results:
14
- - task:
15
- name: Image Classification
16
- type: image-classification
17
- dataset:
18
- name: FCD Multi-Site Cohort (Bonn/DeepFCD)
19
- type: nifti_medical_scans
20
- metrics:
21
- - type: AUROC
22
- value: 0.9691
23
- - type: AUPRC
24
- value: 0.9755
25
- - type: Accuracy
26
- value: 0.9263
27
- - type: F1-Score
28
- value: 0.9288
29
- ---
30
-
31
- # 🧠 3D Deep Learning Model for Focal Cortical Dysplasia (FCD) Detection
32
-
33
- This model is an optimized **3D deep-learning solution** for the automatic detection of **Focal Cortical Dysplasia (FCD)** in volumetric brain MRI data.
34
-
35
- ---
36
-
37
- ## Model Description
38
-
39
- ### Model Architecture
40
- The core model is a custom, optimized **MobileNetV3-Small–style 3D Convolutional Neural Network (CNN)**, adapted for volumetric (3D) medical data.
41
-
42
- **Key Features:**
43
- * **3D Inverted Residual Blocks (IRB3D)**
44
- * **3D Squeeze-and-Excitation (SE3D)**: Channel-wise attention mechanism.
45
- * **Multi-View Fusion**: Integrates features from axial, coronal, and sagittal views.
46
- * **Stochastic Depth**: Regularization technique to prevent overfitting.
47
- * **Total Parameters:** 800,692 (Trainable)
48
-
49
- ### Intended Use
50
- The model is intended for **computational neuroscience and clinical research** to aid in the **preliminary identification and localization** of FCD lesions in 3D brain scans (NIfTI format). It functions as a robust **binary classifier** (FCD present/absent).
51
-
52
- ### Out-of-Scope Use
53
- This model is **not a standalone diagnostic tool**. Predictions must be validated by qualified medical professionals. It is only validated for use with T1-weighted MRI data similar to the training cohort.
54
-
55
- ## Training Data and Evaluation
56
-
57
- ### Training Data
58
- The model was trained on a comprehensive dataset combining **Bonn NIfTI data** and **DeepFCD HDF5 patches**, totaling **2,170 subjects/samples**.
59
-
60
- **Preprocessing:**
61
- * **ComBat Harmonization** was applied to correct for batch effects.
62
- * **TorchIO** was used for 3D-specific data augmentation.
63
-
64
- ### Training Details
65
- | Parameter | Value |
66
- | :--- | :--- |
67
- | **Framework** | PyTorch (with `torch.cuda.amp`) |
68
- | **Optimizer** | AdamW |
69
- | **Learning Rate** | 2e-4 | Managed by a **CosineWarmup Scheduler** for stable convergence. |
70
- | **Batches/Steps** | `BATCH_SIZE = 16`, `GRAD_ACCUM_STEPS = 4` | Effective batch size is 64. |
71
- | **Epochs/Stopping** | Up to 100 epochs | Training utilizes early stopping based on validation accuracy. |
72
- | **Data Split** | Subject-Level Split (Ensured no data leakage between subjects) |
73
- | **Augmentation** | TorchIO | Applies 3D-specific data augmentation (e.g., random flips, affine transformations, noise) for enhanced generalization. |
74
-
75
- ### Evaluation Results
76
- Performance was measured on a held-out **Test Set** using the best validation checkpoint (Epoch 16).
77
-
78
- | Metric | Score |
79
- | :--- | :--- |
80
- | **AUROC** | **0.9691** |
81
- | **AUPRC** | **0.9755** |
82
- | **Accuracy** | 0.9263 |
83
- | **F1-Score** | 0.9288 |
84
-
85
- ### Interpretability (Grad-CAM Saliency)
86
- **Grad-CAM** (Gradient-weighted Class Activation Mapping) is implemented to generate saliency maps that visualize the regions of the 3D volume most critical to the model's prediction. This provides clinical researchers with crucial visual evidence for the model's decision-making process.
87
-
88
- ### Feature Embeddings
89
- Post-training, high-dimensional features are extracted and visualized using **UMAP** and **t-SNE** to confirm data quality, validate the effectiveness of the ComBat harmonization, and visualize data clustering.
90
-
91
- ## Example usage
92
-
93
- 1. **Dependencies:** Ensure you have the necessary libraries installed:
94
- ```bash
95
- pip install torch numpy nibabel scipy
96
- ```
97
- 2. **Input Files:** The original pipeline expects two files: a T1-weighted image (`T1w`) and a FLAIR image, as it uses 2 input channels.
98
- 3. **Save the script:** Save this code as `predict.py`.
99
-
100
- ```python
101
- import os
102
- import sys
103
- import argparse
104
- import numpy as np
105
- import nibabel as nib
106
- from scipy import ndimage
107
-
108
- import torch
109
- import torch.nn as nn
110
- import torch.nn.functional as F
111
- from pathlib import Path
112
-
113
-
114
- IMG_SIZE = (128, 128, 128)
115
- NUM_CLASSES = 2
116
- IN_CHANNELS = 2
117
-
118
- def conv_bn_act(inp, oup, k, s, act=nn.Hardswish):
119
- return nn.Sequential(
120
- nn.Conv3d(inp, oup, k, s, k // 2, bias=False),
121
- nn.BatchNorm3d(oup),
122
- act(inplace=True)
123
- )
124
-
125
- class SE3D(nn.Module):
126
- def __init__(self, ch, r=4):
127
- super().__init__()
128
- self.se = nn.Sequential(
129
- nn.AdaptiveAvgPool3d(1),
130
- nn.Conv3d(ch, ch // r, 1),
131
- nn.ReLU(inplace=True),
132
- nn.Conv3d(ch // r, ch, 1),
133
- nn.Hardsigmoid(inplace=True)
134
- )
135
-
136
- def forward(self, x):
137
- return x * self.se(x)
138
-
139
- class StochasticDepth(nn.Module):
140
- def __init__(self, p=0.0):
141
- super().__init__()
142
- self.p = p
143
-
144
- def forward(self, x):
145
- return x
146
-
147
- class IRB3D(nn.Module):
148
- def __init__(self, inp, oup, k, s, exp, se=True, nl=nn.Hardswish, sd_p=0.0):
149
- super().__init__()
150
- self.use_res = (s == 1 and inp == oup)
151
- hid = inp * exp
152
-
153
- layers = []
154
- if exp != 1:
155
- layers.append(conv_bn_act(inp, hid, 1, 1, nl))
156
- layers.extend([
157
- nn.Conv3d(hid, hid, k, s, k // 2, groups=hid, bias=False),
158
- nn.BatchNorm3d(hid),
159
- nl(inplace=True)
160
- ])
161
- if se:
162
- layers.append(SE3D(hid))
163
- layers.extend([
164
- nn.Conv3d(hid, oup, 1, 1, 0, bias=False),
165
- nn.BatchNorm3d(oup)
166
- ])
167
-
168
- self.conv = nn.Sequential(*layers)
169
- self.sd = StochasticDepth(sd_p) if sd_p > 0 else nn.Identity()
170
-
171
- def forward(self, x):
172
- out = self.conv(x)
173
- if self.use_res:
174
- out = self.sd(out) + x
175
- return out
176
-
177
- class MultiViewFusion(nn.Module):
178
- def __init__(self, ch):
179
- super().__init__()
180
- self.axial = nn.Conv3d(ch, ch, (3, 1, 1), padding=(1, 0, 0), groups=ch)
181
- self.coronal = nn.Conv3d(ch, ch, (1, 3, 1), padding=(0, 1, 0), groups=ch)
182
- self.sagittal = nn.Conv3d(ch, ch, (1, 1, 3), padding=(0, 0, 1), groups=ch)
183
- self.fusion = nn.Conv3d(ch * 3, ch, 1)
184
-
185
- def forward(self, x):
186
- ax = self.axial(x)
187
- cor = self.coronal(x)
188
- sag = self.sagittal(x)
189
- fused = torch.cat([ax, cor, sag], dim=1)
190
- return self.fusion(fused)
191
-
192
- class MobileNetV3Small3D(nn.Module):
193
- def __init__(self, num_classes=NUM_CLASSES, in_ch=IN_CHANNELS, use_checkpointing=False):
194
- super().__init__()
195
- self.stem = conv_bn_act(in_ch, 16, 3, 2)
196
- configs = [
197
- (16, 16, 3, 2, 1, True, nn.ReLU), (16, 24, 3, 2, 4, False, nn.ReLU),
198
- (24, 24, 3, 1, 3, False, nn.ReLU), (24, 40, 5, 2, 3, True, nn.Hardswish),
199
- (40, 40, 5, 1, 3, True, nn.Hardswish), (40, 48, 5, 1, 3, True, nn.Hardswish),
200
- (48, 96, 5, 2, 6, True, nn.Hardswish), (96, 96, 5, 1, 6, True, nn.Hardswish),
201
- ]
202
- sd_probs = np.linspace(0, 0.2, len(configs))
203
- self.blocks = nn.ModuleList([
204
- IRB3D(inp, oup, k, s, exp, se, nl, sd_p)
205
- for (inp, oup, k, s, exp, se, nl), sd_p in zip(configs, sd_probs)
206
- ])
207
- self.fusion = MultiViewFusion(96)
208
- self.conv_head = conv_bn_act(96, 576, 1, 1)
209
- self.pool = nn.AdaptiveAvgPool3d(1)
210
- self.classifier = nn.Sequential(
211
- nn.Linear(576, 256),
212
- nn.Hardswish(inplace=True),
213
- nn.Dropout(0.2),
214
- nn.Linear(256, num_classes)
215
- )
216
-
217
- def forward(self, x):
218
- x = self.stem(x)
219
- for blk in self.blocks:
220
- x = blk(x)
221
- x = self.fusion(x)
222
- x = self.conv_head(x)
223
- x = self.pool(x)
224
- x = x.flatten(1)
225
- x = self.classifier(x)
226
- return x
227
-
228
-
229
- def load_nifti(path):
230
- try:
231
- img = nib.load(str(path))
232
- return img.get_fdata().astype(np.float32)
233
- except Exception as e:
234
- print(f"Failed to load {path}: {e}")
235
- return None
236
-
237
- def normalize_intensity(img):
238
- if img.max() == img.min():
239
- return np.zeros_like(img, dtype=np.float32)
240
-
241
- nonzero = img[img > 0]
242
- if len(nonzero) == 0:
243
- return img.astype(np.float32)
244
-
245
- p1, p99 = np.percentile(nonzero, [1, 99])
246
- img = np.clip(img, p1, p99)
247
- img = (img - img.min()) / (img.max() - img.min() + 1e-8)
248
- return img.astype(np.float32)
249
-
250
- def resize3d(img, target, is_label=False):
251
- if img.size == 0:
252
- return np.zeros(target, dtype=np.float32)
253
-
254
- zoom_factors = [t / s for s, t in zip(img.shape, target)]
255
- order = 0 if is_label else 1
256
- return ndimage.zoom(img, zoom_factors, order=order).astype(np.float32)
257
-
258
- def preprocess_nifti_pair(t1_path, flair_path, target_size=IMG_SIZE):
259
- t1 = load_nifti(t1_path)
260
- flair = load_nifti(flair_path)
261
-
262
- if t1 is None or flair is None:
263
- raise FileNotFoundError("One or both input NIfTI files could not be loaded.")
264
-
265
- t1 = normalize_intensity(t1)
266
- flair = normalize_intensity(flair)
267
-
268
- t1 = resize3d(t1, target_size)
269
- flair = resize3d(flair, target_size)
270
-
271
- img = np.stack([t1, flair], axis=0).astype(np.float32)
272
-
273
- return torch.from_numpy(img).unsqueeze(0)
274
-
275
-
276
- def predict(model_path, t1_file_path, flair_file_path):
277
-
278
- device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
279
- print(f"Using device: {device}")
280
-
281
- model = MobileNetV3Small3D(num_classes=NUM_CLASSES, in_ch=IN_CHANNELS)
282
-
283
- try:
284
- checkpoint = torch.load(model_path, map_location=device)
285
- if 'model' in checkpoint:
286
- model.load_state_dict(checkpoint['model'])
287
- else:
288
- model.load_state_dict(checkpoint)
289
-
290
- model.to(device)
291
- model.eval()
292
- print(f"Successfully loaded model from {model_path}")
293
- except Exception as e:
294
- print(f"Error loading model weights from {model_path}: {e}")
295
- print("Please ensure your .pth file contains the correct state_dict for MobileNetV3Small3D.")
296
- sys.exit(1)
297
-
298
- try:
299
- input_tensor = preprocess_nifti_pair(t1_file_path, flair_file_path)
300
- input_tensor = input_tensor.to(device)
301
- except Exception as e:
302
- print(f"Error during file preprocessing: {e}")
303
- sys.exit(1)
304
-
305
- with torch.no_grad():
306
- logits = model(input_tensor)
307
- probabilities = F.softmax(logits, dim=1).squeeze(0)
308
-
309
- predicted_class = probabilities.argmax().item()
310
-
311
- print("\n--- Prediction Results ---")
312
- print(f"Input T1w file: {t1_file_path}")
313
- print(f"Input FLAIR file: {flair_file_path}")
314
- print(f"Probabilities (Class 0, Class 1): {probabilities.cpu().numpy()}")
315
- print(f"Predicted Class Index: {predicted_class}")
316
-
317
- if NUM_CLASSES == 2:
318
- class_labels = {0: "Control", 1: "FCD"}
319
- predicted_label = class_labels.get(predicted_class, f"Class {predicted_class}")
320
- print(f"Predicted Label: {predicted_label}")
321
- else:
322
- print("Model has more than 2 classes. Prediction is based on class index.")
323
-
324
-
325
- def main():
326
- parser = argparse.ArgumentParser(description="Run 3D CNN prediction using a trained .pth model.")
327
- parser.add_argument("model_path", type=str, help="Path to the MobileNetV3Small3D model file (.pth).")
328
- parser.add_argument("t1_file", type=str, help="Path to the T1w NIfTI file for prediction.")
329
- parser.add_argument("flair_file", type=str, help="Path to the FLAIR NIfTI file for prediction.")
330
-
331
- example_t1 = "./path/to/sub-01_T1w.nii.gz"
332
- example_flair = "./path/to/sub-01_FLAIR.nii.gz"
333
- example_model = "./out/models/best_model.pth"
334
-
335
- if len(sys.argv) == 1:
336
- print("--- Usage Example ---")
337
- print(f"python {Path(sys.argv[0]).name} {example_model} {example_t1} {example_flair}")
338
- print("\n--- Note ---")
339
- print("This script requires a paired T1w and FLAIR NIfTI file as input, as the original model expects 2 channels.")
340
- sys.exit(1)
341
-
342
- args = parser.parse_args()
343
-
344
- predict(args.model_path, args.t1_file, args.flair_file)
345
-
346
- if __name__ == "__main__":
347
- import warnings
348
- warnings.filterwarnings("ignore")
349
- main()
350
- ```
351
- 4. **Execution:** Run the script from your terminal:
352
- ```bash
353
- python predict.py /path/to/your/model.pth /path/to/your/T1w.nii.gz /path/to/your/FLAIR.nii.gz
354
- ```
355
-
356
- ### Visual gallery
357
-
358
- ### Pipeline architecture
359
- This image shows the **MobileNetV3Small3D architecture** used for classification, detailing the various blocks and the overall flow from the dual-channel input to the final classification layer. It serves as a structural blueprint of the model.
360
-
361
- ![Architecture](vis/architecture.png)
362
-
363
- ### Multi-view slices
364
- This visualization displays **sample cross-sections (axial, coronal, sagittal views)** of the preprocessed T1w and FLAIR MR images. It provides a visual check of the input data quality and the multi-view nature of the medical imaging inputs.
365
-
366
- ![Multiview Sample](vis/multiview_sample.png)
367
-
368
- ### Training curves
369
- These plots track the model's performance metrics, specifically **loss and accuracy/AUC/F1-score**, over epochs for both the training and validation sets. They are essential for confirming convergence, detecting overfitting (a large gap between training and validation performance), and determining the optimal training duration.
370
-
371
- ![Training Curves](vis/training_curves.png)
372
-
373
- ### Embeddings and harmonization
374
- These figures demonstrate **ComBat harmonization**, a technique used to mitigate scanner/site-specific batch effects in the features extracted by the model. The **UMAP/t-SNE** plots show the high-dimensional feature vectors reduced to 2D, illustrating how well the data from different sites have been mixed (**harmonized**) for robust classification, compared to the unharmonized features.
375
-
376
- ![Harmonization](vis/harmonization.png)
377
-
378
- ![UMAP/t-SNE (harmonized)](vis/harmonized_embeddings.png)
379
-
380
- ### Cohort composition
381
- This section provides visualizations of the **data split** (training, validation, test) and the **distribution of demographic or clinical variables** (e.g., site, age, sex, diagnosis) across these subsets. It verifies that the splits were performed correctly at the subject level and that the cohorts are balanced and representative.
382
-
383
- ![Cohort Splits](vis/cohort_splits.png)
384
-
385
- ![Cohort Grid](vis/cohort_grid.png)
386
-
387
- ### Prediction dashboard
388
- This single visualization summarizes the final **model performance on the test set**. It typically includes the **confusion matrix**, **ROC curve**, **precision-recall curve**, and key metrics (like AUC, F1-score, accuracy) in one comprehensive figure for easy evaluation of the model's discriminative ability.
389
-
390
- ![Prediction Dashboard](vis/prediction_dashboard.png)
391
-
392
- ### Asymmetry examples
393
- These images showcase examples of **asymmetry maps** or **lesion segmentation visualizations**. They visually confirm the presence of the suspected Focal Cortical Dysplasia (FCD) lesion and provide a qualitative comparison between the ground truth and the model's (or pipeline's) inherent ability to highlight the pathological region based on structural differences.
394
-
395
- ![Asymmetry 0](vis/asymmetry_0.png)
396
-
397
- ![Asymmetry 1](vis/asymmetry_1.png)
398
-
399
- ![Asymmetry 2](vis/asymmetry_2.png)
400
-
401
- ### Grad-CAM saliency montages (3D and overlays; for 10 samples)
402
-
403
- **Grad-CAM** (Gradient-weighted Class Activation Mapping) is an eXplainable AI (XAI) technique. These montages show heatmaps that indicate **which regions of the 3D input volume were most influential** in the model's final classification decision for various samples. The $3D$ view and the $Overlay$ view confirm that the model correctly focused its attention on the actual FCD lesion areas when making its prediction.
404
- ![Grad-CAM 3D 0](saliency/sample_0_cam_3d.png)
405
-
406
- ![Grad-CAM Overlay 0](saliency/gradcam_overlay_0.png)
407
-
408
- ![Grad-CAM 3D 1](saliency/sample_1_cam_3d.png)
409
-
410
- ![Grad-CAM Overlay 1](saliency/gradcam_overlay_1.png)
411
-
412
- ![Grad-CAM 3D 2](saliency/sample_2_cam_3d.png)
413
-
414
- ![Grad-CAM Overlay 2](saliency/gradcam_overlay_2.png)
415
-
416
- ![Grad-CAM 3D 3](saliency/sample_3_cam_3d.png)
417
-
418
- ![Grad-CAM Overlay 3](saliency/gradcam_overlay_3.png)
419
-
420
- ![Grad-CAM 3D 4](saliency/sample_4_cam_3d.png)
421
-
422
- ![Grad-CAM Overlay 4](saliency/gradcam_overlay_4.png)
423
-
424
- ![Grad-CAM 3D 5](saliency/sample_5_cam_3d.png)
425
-
426
- ![Grad-CAM Overlay 5](saliency/gradcam_overlay_5.png)
427
-
428
- ![Grad-CAM 3D 6](saliency/sample_6_cam_3d.png)
429
-
430
- ![Grad-CAM Overlay 6](saliency/gradcam_overlay_6.png)
431
-
432
- ![Grad-CAM 3D 7](saliency/sample_7_cam_3d.png)
433
-
434
- ![Grad-CAM Overlay 7](saliency/gradcam_overlay_7.png)
435
-
436
- ![Grad-CAM 3D 8](saliency/sample_8_cam_3d.png)
437
-
438
- ![Grad-CAM Overlay 8](saliency/gradcam_overlay_8.png)
439
-
440
- ![Grad-CAM 3D 9](saliency/sample_9_cam_3d.png)
441
-
442
- ![Grad-CAM Overlay 9](saliency/gradcam_overlay_9.png)
443
-
444
- ## ⚖️ Ethics and Limitations
445
-
446
- ### Clinical Risk and Human Oversight
447
-
448
- This model is a **research tool** and not a certified medical device. The primary ethical consideration is the **risk of misapplication** in a diagnostic setting. Predictions **must not** be used for direct patient management without rigorous validation and final interpretation by qualified human medical professionals.
449
-
450
- ### Bias and Generalization
451
-
452
- The model's performance is intrinsically linked to the composition of its training data.
453
-
454
- * **Geographic/Scanner Bias:** The model may exhibit performance degradation when applied to data acquired from scanner models or acquisition protocols significantly different from the multi-site Bonn/DeepFCD cohort, despite ComBat harmonization.
455
- * **Cohort Limitations:** Performance on demographics (e.g., age ranges, FCD types) underrepresented in the training data is not guaranteed. Researchers must perform validation on their specific target cohort.
456
-
457
- ### Input Modality Validation
458
-
459
- The model is rigorously validated only for use with **T1-weighted MRI** volumes that have undergone specific harmonization and preprocessing steps. Any deviation from this input protocol (e.g., using T2 or FLAIR data) is a significant limitation and will compromise prediction reliability.
 
1
+ ---
2
+ language: en
3
+ license: gpl-3.0
4
+ tags:
5
+ - 3d-vision
6
+ - medical-imaging
7
+ - deep-learning
8
+ - pytorch
9
+ - fcd-detection
10
+ - brain-mri
11
+ model-index:
12
+ - name: MobileNetV3-Small-3D-FCD
13
+ results:
14
+ - task:
15
+ name: Image Classification
16
+ type: image-classification
17
+ dataset:
18
+ name: FCD Multi-Site Cohort (Bonn/DeepFCD)
19
+ type: nifti_medical_scans
20
+ metrics:
21
+ - type: AUROC
22
+ value: 0.9829
23
+ - type: Accuracy
24
+ value: 0.9378
25
+ - type: F1-Score
26
+ value: 0.9382
27
+ ---
28
+
29
+ # 🧠 3D Deep Learning Model for Focal Cortical Dysplasia (FCD) Detection
30
+
31
+ This model is an optimized **3D deep-learning solution** for the automatic detection of **Focal Cortical Dysplasia (FCD)** in volumetric brain MRI data.
32
+
33
+ ---
34
+
35
+ ## Model Description
36
+
37
+ ### Model Architecture
38
+ The core model is a custom, optimized **MobileNetV3-Small–style 3D Convolutional Neural Network (CNN)**, adapted for volumetric (3D) medical data.
39
+
40
+ **Key Features:**
41
+ * **3D Inverted Residual Blocks (IRB3D)**
42
+ * **3D Squeeze-and-Excitation (SE3D)**: Channel-wise attention mechanism.
43
+ * **Multi-View Fusion**: Integrates features from axial, coronal, and sagittal views.
44
+ * **Stochastic Depth**: Regularization technique to prevent overfitting.
45
+ * **Total Parameters:** 800,692 (Trainable)
46
+
47
+ ### Intended Use
48
+ The model is intended for **computational neuroscience and clinical research** to aid in the **preliminary identification and localization** of FCD lesions in 3D brain scans (NIfTI format). It functions as a robust **binary classifier** (FCD present/absent).
49
+
50
+ ### Out-of-Scope Use
51
+ This model is **not a standalone diagnostic tool**. Predictions must be validated by qualified medical professionals. It is only validated for use with T1-weighted MRI data similar to the training cohort.
52
+
53
+ ## Training Data and Evaluation
54
+
55
+ ### Training Data
56
+ The model was trained on a comprehensive dataset combining **Bonn NIfTI data** and **DeepFCD HDF5 patches**, totaling **2,170 subjects/samples**.
57
+
58
+ **Preprocessing:**
59
+ * **ComBat Harmonization** was applied to correct for batch effects.
60
+ * **TorchIO** was used for 3D-specific data augmentation.
61
+
62
+ ### Training Details
63
+ | Parameter | Value |
64
+ | :--- | :--- |
65
+ | **Framework** | PyTorch (with `torch.cuda.amp`) |
66
+ | **Optimizer** | AdamW |
67
+ | **Learning Rate** | 2e-4 | Managed by a **CosineWarmup Scheduler** for stable convergence. |
68
+ | **Batches/Steps** | `BATCH_SIZE = 16`, `GRAD_ACCUM_STEPS = 4` | Effective batch size is 64. |
69
+ | **Epochs/Stopping** | Up to 100 epochs | Training utilizes early stopping based on validation accuracy. |
70
+ | **Data Split** | Subject-Level Split (Ensured no data leakage between subjects) |
71
+ | **Augmentation** | TorchIO | Applies 3D-specific data augmentation (e.g., random flips, affine transformations, noise) for enhanced generalization. |
72
+
73
+ ### Evaluation Results
74
+ Performance was measured on a held-out **Test Set** using the best validation checkpoint (Epoch 32).
75
+
76
+ | Metric | Score |
77
+ | :--- | :--- |
78
+ | **AUROC** | 0.9829 |
79
+ | **Accuracy** | **0.9378** |
80
+ | **F1-Score** | 0.9382 |
81
+
82
+ ### Interpretability (Grad-CAM Saliency)
83
+ **Grad-CAM** (Gradient-weighted Class Activation Mapping) is implemented to generate saliency maps that visualize the regions of the 3D volume most critical to the model's prediction. This provides clinical researchers with crucial visual evidence for the model's decision-making process.
84
+
85
+ ### Feature Embeddings
86
+ Post-training, high-dimensional features are extracted and visualized using **UMAP** and **t-SNE** to confirm data quality, validate the effectiveness of the ComBat harmonization, and visualize data clustering.
87
+
88
+ ## Example usage
89
+
90
+ 1. **Dependencies:** Ensure you have the necessary libraries installed:
91
+ ```bash
92
+ pip install torch numpy nibabel scipy
93
+ ```
94
+ 2. **Input Files:** The original pipeline expects two files: a T1-weighted image (`T1w`) and a FLAIR image, as it uses 2 input channels.
95
+ 3. **Save the script:** Save this code as `predict.py`.
96
+
97
+ ```python
98
+ import os
99
+ import sys
100
+ import argparse
101
+ import numpy as np
102
+ import nibabel as nib
103
+ from scipy import ndimage
104
+
105
+ import torch
106
+ import torch.nn as nn
107
+ import torch.nn.functional as F
108
+ from pathlib import Path
109
+
110
+
111
+ IMG_SIZE = (128, 128, 128)
112
+ NUM_CLASSES = 2
113
+ IN_CHANNELS = 2
114
+
115
+ def conv_bn_act(inp, oup, k, s, act=nn.Hardswish):
116
+ return nn.Sequential(
117
+ nn.Conv3d(inp, oup, k, s, k // 2, bias=False),
118
+ nn.BatchNorm3d(oup),
119
+ act(inplace=True)
120
+ )
121
+
122
+ class SE3D(nn.Module):
123
+ def __init__(self, ch, r=4):
124
+ super().__init__()
125
+ self.se = nn.Sequential(
126
+ nn.AdaptiveAvgPool3d(1),
127
+ nn.Conv3d(ch, ch // r, 1),
128
+ nn.ReLU(inplace=True),
129
+ nn.Conv3d(ch // r, ch, 1),
130
+ nn.Hardsigmoid(inplace=True)
131
+ )
132
+
133
+ def forward(self, x):
134
+ return x * self.se(x)
135
+
136
+ class StochasticDepth(nn.Module):
137
+ def __init__(self, p=0.0):
138
+ super().__init__()
139
+ self.p = p
140
+
141
+ def forward(self, x):
142
+ return x
143
+
144
+ class IRB3D(nn.Module):
145
+ def __init__(self, inp, oup, k, s, exp, se=True, nl=nn.Hardswish, sd_p=0.0):
146
+ super().__init__()
147
+ self.use_res = (s == 1 and inp == oup)
148
+ hid = inp * exp
149
+
150
+ layers = []
151
+ if exp != 1:
152
+ layers.append(conv_bn_act(inp, hid, 1, 1, nl))
153
+ layers.extend([
154
+ nn.Conv3d(hid, hid, k, s, k // 2, groups=hid, bias=False),
155
+ nn.BatchNorm3d(hid),
156
+ nl(inplace=True)
157
+ ])
158
+ if se:
159
+ layers.append(SE3D(hid))
160
+ layers.extend([
161
+ nn.Conv3d(hid, oup, 1, 1, 0, bias=False),
162
+ nn.BatchNorm3d(oup)
163
+ ])
164
+
165
+ self.conv = nn.Sequential(*layers)
166
+ self.sd = StochasticDepth(sd_p) if sd_p > 0 else nn.Identity()
167
+
168
+ def forward(self, x):
169
+ out = self.conv(x)
170
+ if self.use_res:
171
+ out = self.sd(out) + x
172
+ return out
173
+
174
+ class MultiViewFusion(nn.Module):
175
+ def __init__(self, ch):
176
+ super().__init__()
177
+ self.axial = nn.Conv3d(ch, ch, (3, 1, 1), padding=(1, 0, 0), groups=ch)
178
+ self.coronal = nn.Conv3d(ch, ch, (1, 3, 1), padding=(0, 1, 0), groups=ch)
179
+ self.sagittal = nn.Conv3d(ch, ch, (1, 1, 3), padding=(0, 0, 1), groups=ch)
180
+ self.fusion = nn.Conv3d(ch * 3, ch, 1)
181
+
182
+ def forward(self, x):
183
+ ax = self.axial(x)
184
+ cor = self.coronal(x)
185
+ sag = self.sagittal(x)
186
+ fused = torch.cat([ax, cor, sag], dim=1)
187
+ return self.fusion(fused)
188
+
189
+ class MobileNetV3Small3D(nn.Module):
190
+ def __init__(self, num_classes=NUM_CLASSES, in_ch=IN_CHANNELS, use_checkpointing=False):
191
+ super().__init__()
192
+ self.stem = conv_bn_act(in_ch, 16, 3, 2)
193
+ configs = [
194
+ (16, 16, 3, 2, 1, True, nn.ReLU), (16, 24, 3, 2, 4, False, nn.ReLU),
195
+ (24, 24, 3, 1, 3, False, nn.ReLU), (24, 40, 5, 2, 3, True, nn.Hardswish),
196
+ (40, 40, 5, 1, 3, True, nn.Hardswish), (40, 48, 5, 1, 3, True, nn.Hardswish),
197
+ (48, 96, 5, 2, 6, True, nn.Hardswish), (96, 96, 5, 1, 6, True, nn.Hardswish),
198
+ ]
199
+ sd_probs = np.linspace(0, 0.2, len(configs))
200
+ self.blocks = nn.ModuleList([
201
+ IRB3D(inp, oup, k, s, exp, se, nl, sd_p)
202
+ for (inp, oup, k, s, exp, se, nl), sd_p in zip(configs, sd_probs)
203
+ ])
204
+ self.fusion = MultiViewFusion(96)
205
+ self.conv_head = conv_bn_act(96, 576, 1, 1)
206
+ self.pool = nn.AdaptiveAvgPool3d(1)
207
+ self.classifier = nn.Sequential(
208
+ nn.Linear(576, 256),
209
+ nn.Hardswish(inplace=True),
210
+ nn.Dropout(0.2),
211
+ nn.Linear(256, num_classes)
212
+ )
213
+
214
+ def forward(self, x):
215
+ x = self.stem(x)
216
+ for blk in self.blocks:
217
+ x = blk(x)
218
+ x = self.fusion(x)
219
+ x = self.conv_head(x)
220
+ x = self.pool(x)
221
+ x = x.flatten(1)
222
+ x = self.classifier(x)
223
+ return x
224
+
225
+
226
+ def load_nifti(path):
227
+ try:
228
+ img = nib.load(str(path))
229
+ return img.get_fdata().astype(np.float32)
230
+ except Exception as e:
231
+ print(f"Failed to load {path}: {e}")
232
+ return None
233
+
234
+ def normalize_intensity(img):
235
+ if img.max() == img.min():
236
+ return np.zeros_like(img, dtype=np.float32)
237
+
238
+ nonzero = img[img > 0]
239
+ if len(nonzero) == 0:
240
+ return img.astype(np.float32)
241
+
242
+ p1, p99 = np.percentile(nonzero, [1, 99])
243
+ img = np.clip(img, p1, p99)
244
+ img = (img - img.min()) / (img.max() - img.min() + 1e-8)
245
+ return img.astype(np.float32)
246
+
247
+ def resize3d(img, target, is_label=False):
248
+ if img.size == 0:
249
+ return np.zeros(target, dtype=np.float32)
250
+
251
+ zoom_factors = [t / s for s, t in zip(img.shape, target)]
252
+ order = 0 if is_label else 1
253
+ return ndimage.zoom(img, zoom_factors, order=order).astype(np.float32)
254
+
255
+ def preprocess_nifti_pair(t1_path, flair_path, target_size=IMG_SIZE):
256
+ t1 = load_nifti(t1_path)
257
+ flair = load_nifti(flair_path)
258
+
259
+ if t1 is None or flair is None:
260
+ raise FileNotFoundError("One or both input NIfTI files could not be loaded.")
261
+
262
+ t1 = normalize_intensity(t1)
263
+ flair = normalize_intensity(flair)
264
+
265
+ t1 = resize3d(t1, target_size)
266
+ flair = resize3d(flair, target_size)
267
+
268
+ img = np.stack([t1, flair], axis=0).astype(np.float32)
269
+
270
+ return torch.from_numpy(img).unsqueeze(0)
271
+
272
+
273
+ def predict(model_path, t1_file_path, flair_file_path):
274
+
275
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
276
+ print(f"Using device: {device}")
277
+
278
+ model = MobileNetV3Small3D(num_classes=NUM_CLASSES, in_ch=IN_CHANNELS)
279
+
280
+ try:
281
+ checkpoint = torch.load(model_path, map_location=device)
282
+ if 'model' in checkpoint:
283
+ model.load_state_dict(checkpoint['model'])
284
+ else:
285
+ model.load_state_dict(checkpoint)
286
+
287
+ model.to(device)
288
+ model.eval()
289
+ print(f"Successfully loaded model from {model_path}")
290
+ except Exception as e:
291
+ print(f"Error loading model weights from {model_path}: {e}")
292
+ print("Please ensure your .pth file contains the correct state_dict for MobileNetV3Small3D.")
293
+ sys.exit(1)
294
+
295
+ try:
296
+ input_tensor = preprocess_nifti_pair(t1_file_path, flair_file_path)
297
+ input_tensor = input_tensor.to(device)
298
+ except Exception as e:
299
+ print(f"Error during file preprocessing: {e}")
300
+ sys.exit(1)
301
+
302
+ with torch.no_grad():
303
+ logits = model(input_tensor)
304
+ probabilities = F.softmax(logits, dim=1).squeeze(0)
305
+
306
+ predicted_class = probabilities.argmax().item()
307
+
308
+ print("\n--- Prediction Results ---")
309
+ print(f"Input T1w file: {t1_file_path}")
310
+ print(f"Input FLAIR file: {flair_file_path}")
311
+ print(f"Probabilities (Class 0, Class 1): {probabilities.cpu().numpy()}")
312
+ print(f"Predicted Class Index: {predicted_class}")
313
+
314
+ if NUM_CLASSES == 2:
315
+ class_labels = {0: "Control", 1: "FCD"}
316
+ predicted_label = class_labels.get(predicted_class, f"Class {predicted_class}")
317
+ print(f"Predicted Label: {predicted_label}")
318
+ else:
319
+ print("Model has more than 2 classes. Prediction is based on class index.")
320
+
321
+
322
+ def main():
323
+ parser = argparse.ArgumentParser(description="Run 3D CNN prediction using a trained .pth model.")
324
+ parser.add_argument("model_path", type=str, help="Path to the MobileNetV3Small3D model file (.pth).")
325
+ parser.add_argument("t1_file", type=str, help="Path to the T1w NIfTI file for prediction.")
326
+ parser.add_argument("flair_file", type=str, help="Path to the FLAIR NIfTI file for prediction.")
327
+
328
+ example_t1 = "./path/to/sub-01_T1w.nii.gz"
329
+ example_flair = "./path/to/sub-01_FLAIR.nii.gz"
330
+ example_model = "./out/models/best_model.pth"
331
+
332
+ if len(sys.argv) == 1:
333
+ print("--- Usage Example ---")
334
+ print(f"python {Path(sys.argv[0]).name} {example_model} {example_t1} {example_flair}")
335
+ print("\n--- Note ---")
336
+ print("This script requires a paired T1w and FLAIR NIfTI file as input, as the original model expects 2 channels.")
337
+ sys.exit(1)
338
+
339
+ args = parser.parse_args()
340
+
341
+ predict(args.model_path, args.t1_file, args.flair_file)
342
+
343
+ if __name__ == "__main__":
344
+ import warnings
345
+ warnings.filterwarnings("ignore")
346
+ main()
347
+ ```
348
+ 4. **Execution:** Run the script from your terminal:
349
+ ```bash
350
+ python predict.py /path/to/your/model.pth /path/to/your/T1w.nii.gz /path/to/your/FLAIR.nii.gz
351
+ ```
352
+
353
+ ### Visual gallery
354
+
355
+ ### Pipeline architecture
356
+ This image shows the **MobileNetV3Small3D architecture** used for classification, detailing the various blocks and the overall flow from the dual-channel input to the final classification layer. It serves as a structural blueprint of the model.
357
+
358
+ ![Architecture](vis/architecture.png)
359
+
360
+ ### Multi-view slices
361
+ This visualization displays **sample cross-sections (axial, coronal, sagittal views)** of the preprocessed T1w and FLAIR MR images. It provides a visual check of the input data quality and the multi-view nature of the medical imaging inputs.
362
+
363
+ ![Multiview Sample](vis/multiview_sample.png)
364
+
365
+ ### Training curves
366
+ These plots track the model's performance metrics, specifically **loss and accuracy/AUC/F1-score**, over epochs for both the training and validation sets. They are essential for confirming convergence, detecting overfitting (a large gap between training and validation performance), and determining the optimal training duration.
367
+
368
+ ![Training Curves](vis/training_curves.png)
369
+
370
+ ### Embeddings and harmonization
371
+ These figures demonstrate **ComBat harmonization**, a technique used to mitigate scanner/site-specific batch effects in the features extracted by the model. The **UMAP/t-SNE** plots show the high-dimensional feature vectors reduced to 2D, illustrating how well the data from different sites have been mixed (**harmonized**) for robust classification, compared to the unharmonized features.
372
+
373
+ ![Harmonization](vis/harmonization.png)
374
+
375
+ ![UMAP/t-SNE (harmonized)](vis/harmonized_embeddings.png)
376
+
377
+ ### Cohort composition
378
+ This section provides visualizations of the **data split** (training, validation, test) and the **distribution of demographic or clinical variables** (e.g., site, age, sex, diagnosis) across these subsets. It verifies that the splits were performed correctly at the subject level and that the cohorts are balanced and representative.
379
+
380
+ ![Cohort Splits](vis/cohort_splits.png)
381
+
382
+ ![Cohort Grid](vis/cohort_grid.png)
383
+
384
+ ### Prediction dashboard
385
+ This single visualization summarizes the final **model performance on the test set**. It typically includes the **confusion matrix**, **ROC curve**, **precision-recall curve**, and key metrics (like AUC, F1-score, accuracy) in one comprehensive figure for easy evaluation of the model's discriminative ability.
386
+
387
+ ![Prediction Dashboard](vis/prediction_dashboard.png)
388
+
389
+ ### Asymmetry examples
390
+ These images showcase examples of **asymmetry maps** or **lesion segmentation visualizations**. They visually confirm the presence of the suspected Focal Cortical Dysplasia (FCD) lesion and provide a qualitative comparison between the ground truth and the model's (or pipeline's) inherent ability to highlight the pathological region based on structural differences.
391
+
392
+ ![Asymmetry 0](vis/asymmetry_0.png)
393
+
394
+ ![Asymmetry 1](vis/asymmetry_1.png)
395
+
396
+ ![Asymmetry 2](vis/asymmetry_2.png)
397
+
398
+ ### Grad-CAM saliency montages (3D and overlays; for 10 samples)
399
+
400
+ **Grad-CAM** (Gradient-weighted Class Activation Mapping) is an eXplainable AI (XAI) technique. These montages show heatmaps that indicate **which regions of the 3D input volume were most influential** in the model's final classification decision for various samples. The $3D$ view and the $Overlay$ view confirm that the model correctly focused its attention on the actual FCD lesion areas when making its prediction.
401
+ ![Grad-CAM 3D 0](saliency/sample_0_cam_3d.png)
402
+
403
+ ![Grad-CAM Overlay 0](saliency/gradcam_overlay_0.png)
404
+
405
+ ![Grad-CAM 3D 1](saliency/sample_1_cam_3d.png)
406
+
407
+ ![Grad-CAM Overlay 1](saliency/gradcam_overlay_1.png)
408
+
409
+ ![Grad-CAM 3D 2](saliency/sample_2_cam_3d.png)
410
+
411
+ ![Grad-CAM Overlay 2](saliency/gradcam_overlay_2.png)
412
+
413
+ ![Grad-CAM 3D 3](saliency/sample_3_cam_3d.png)
414
+
415
+ ![Grad-CAM Overlay 3](saliency/gradcam_overlay_3.png)
416
+
417
+ ![Grad-CAM 3D 4](saliency/sample_4_cam_3d.png)
418
+
419
+ ![Grad-CAM Overlay 4](saliency/gradcam_overlay_4.png)
420
+
421
+ ![Grad-CAM 3D 5](saliency/sample_5_cam_3d.png)
422
+
423
+ ![Grad-CAM Overlay 5](saliency/gradcam_overlay_5.png)
424
+
425
+ ![Grad-CAM 3D 6](saliency/sample_6_cam_3d.png)
426
+
427
+ ![Grad-CAM Overlay 6](saliency/gradcam_overlay_6.png)
428
+
429
+ ![Grad-CAM 3D 7](saliency/sample_7_cam_3d.png)
430
+
431
+ ![Grad-CAM Overlay 7](saliency/gradcam_overlay_7.png)
432
+
433
+ ![Grad-CAM 3D 8](saliency/sample_8_cam_3d.png)
434
+
435
+ ![Grad-CAM Overlay 8](saliency/gradcam_overlay_8.png)
436
+
437
+ ![Grad-CAM 3D 9](saliency/sample_9_cam_3d.png)
438
+
439
+ ![Grad-CAM Overlay 9](saliency/gradcam_overlay_9.png)
440
+
441
+ ## ⚖️ Ethics and Limitations
442
+
443
+ ### Clinical Risk and Human Oversight
444
+
445
+ This model is a **research tool** and not a certified medical device. The primary ethical consideration is the **risk of misapplication** in a diagnostic setting. Predictions **must not** be used for direct patient management without rigorous validation and final interpretation by qualified human medical professionals.
446
+
447
+ ### Bias and Generalization
448
+
449
+ The model's performance is intrinsically linked to the composition of its training data.
450
+
451
+ * **Geographic/Scanner Bias:** The model may exhibit performance degradation when applied to data acquired from scanner models or acquisition protocols significantly different from the multi-site Bonn/DeepFCD cohort, despite ComBat harmonization.
452
+ * **Cohort Limitations:** Performance on demographics (e.g., age ranges, FCD types) underrepresented in the training data is not guaranteed. Researchers must perform validation on their specific target cohort.
453
+
454
+ ### Input Modality Validation
455
+
456
+ The model is rigorously validated only for use with **T1-weighted MRI** volumes that have undergone specific harmonization and preprocessing steps. Any deviation from this input protocol (e.g., using T2 or FLAIR data) is a significant limitation and will compromise prediction reliability.