File size: 18,090 Bytes
2f4b4ff |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 |
---
language: en
license: gpl-3.0
tags:
- 3d-vision
- medical-imaging
- deep-learning
- pytorch
- fcd-detection
- brain-mri
model-index:
- name: MobileNetV3-Small-3D-FCD
results:
- task:
name: Image Classification
type: image-classification
dataset:
name: FCD Multi-Site Cohort (Bonn/DeepFCD)
type: nifti_medical_scans
metrics:
- type: AUROC
value: 0.9829
- type: Accuracy
value: 0.9378
- type: F1-Score
value: 0.9382
---
# 🧠 3D Deep Learning Model for Focal Cortical Dysplasia (FCD) Detection
This model is an optimized **3D deep-learning solution** for the automatic detection of **Focal Cortical Dysplasia (FCD)** in volumetric brain MRI data.
---
## Model Description
### Model Architecture
The core model is a custom, optimized **MobileNetV3-Small–style 3D Convolutional Neural Network (CNN)**, adapted for volumetric (3D) medical data.
**Key Features:**
* **3D Inverted Residual Blocks (IRB3D)**
* **3D Squeeze-and-Excitation (SE3D)**: Channel-wise attention mechanism.
* **Multi-View Fusion**: Integrates features from axial, coronal, and sagittal views.
* **Stochastic Depth**: Regularization technique to prevent overfitting.
* **Total Parameters:** 800,692 (Trainable)
### Intended Use
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).
### Out-of-Scope Use
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.
## Training Data and Evaluation
### Training Data
The model was trained on a comprehensive dataset combining **Bonn NIfTI data** and **DeepFCD HDF5 patches**, totaling **2,170 subjects/samples**.
**Preprocessing:**
* **ComBat Harmonization** was applied to correct for batch effects.
* **TorchIO** was used for 3D-specific data augmentation.
### Training Details
| Parameter | Value |
| :--- | :--- |
| **Framework** | PyTorch (with `torch.cuda.amp`) |
| **Optimizer** | AdamW |
| **Learning Rate** | 2e-4 | Managed by a **CosineWarmup Scheduler** for stable convergence. |
| **Batches/Steps** | `BATCH_SIZE = 16`, `GRAD_ACCUM_STEPS = 4` | Effective batch size is 64. |
| **Epochs/Stopping** | Up to 100 epochs | Training utilizes early stopping based on validation accuracy. |
| **Data Split** | Subject-Level Split (Ensured no data leakage between subjects) |
| **Augmentation** | TorchIO | Applies 3D-specific data augmentation (e.g., random flips, affine transformations, noise) for enhanced generalization. |
### Evaluation Results
Performance was measured on a held-out **Test Set** using the best validation checkpoint (Epoch 32).
| Metric | Score |
| :--- | :--- |
| **AUROC** | 0.9829 |
| **Accuracy** | **0.9378** |
| **F1-Score** | 0.9382 |
### Interpretability (Grad-CAM Saliency)
**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.
### Feature Embeddings
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.
## Example usage
1. **Dependencies:** Ensure you have the necessary libraries installed:
```bash
pip install torch numpy nibabel scipy
```
2. **Input Files:** The original pipeline expects two files: a T1-weighted image (`T1w`) and a FLAIR image, as it uses 2 input channels.
3. **Save the script:** Save this code as `predict.py`.
```python
import os
import sys
import argparse
import numpy as np
import nibabel as nib
from scipy import ndimage
import torch
import torch.nn as nn
import torch.nn.functional as F
from pathlib import Path
IMG_SIZE = (128, 128, 128)
NUM_CLASSES = 2
IN_CHANNELS = 2
def conv_bn_act(inp, oup, k, s, act=nn.Hardswish):
return nn.Sequential(
nn.Conv3d(inp, oup, k, s, k // 2, bias=False),
nn.BatchNorm3d(oup),
act(inplace=True)
)
class SE3D(nn.Module):
def __init__(self, ch, r=4):
super().__init__()
self.se = nn.Sequential(
nn.AdaptiveAvgPool3d(1),
nn.Conv3d(ch, ch // r, 1),
nn.ReLU(inplace=True),
nn.Conv3d(ch // r, ch, 1),
nn.Hardsigmoid(inplace=True)
)
def forward(self, x):
return x * self.se(x)
class StochasticDepth(nn.Module):
def __init__(self, p=0.0):
super().__init__()
self.p = p
def forward(self, x):
return x
class IRB3D(nn.Module):
def __init__(self, inp, oup, k, s, exp, se=True, nl=nn.Hardswish, sd_p=0.0):
super().__init__()
self.use_res = (s == 1 and inp == oup)
hid = inp * exp
layers = []
if exp != 1:
layers.append(conv_bn_act(inp, hid, 1, 1, nl))
layers.extend([
nn.Conv3d(hid, hid, k, s, k // 2, groups=hid, bias=False),
nn.BatchNorm3d(hid),
nl(inplace=True)
])
if se:
layers.append(SE3D(hid))
layers.extend([
nn.Conv3d(hid, oup, 1, 1, 0, bias=False),
nn.BatchNorm3d(oup)
])
self.conv = nn.Sequential(*layers)
self.sd = StochasticDepth(sd_p) if sd_p > 0 else nn.Identity()
def forward(self, x):
out = self.conv(x)
if self.use_res:
out = self.sd(out) + x
return out
class MultiViewFusion(nn.Module):
def __init__(self, ch):
super().__init__()
self.axial = nn.Conv3d(ch, ch, (3, 1, 1), padding=(1, 0, 0), groups=ch)
self.coronal = nn.Conv3d(ch, ch, (1, 3, 1), padding=(0, 1, 0), groups=ch)
self.sagittal = nn.Conv3d(ch, ch, (1, 1, 3), padding=(0, 0, 1), groups=ch)
self.fusion = nn.Conv3d(ch * 3, ch, 1)
def forward(self, x):
ax = self.axial(x)
cor = self.coronal(x)
sag = self.sagittal(x)
fused = torch.cat([ax, cor, sag], dim=1)
return self.fusion(fused)
class MobileNetV3Small3D(nn.Module):
def __init__(self, num_classes=NUM_CLASSES, in_ch=IN_CHANNELS, use_checkpointing=False):
super().__init__()
self.stem = conv_bn_act(in_ch, 16, 3, 2)
configs = [
(16, 16, 3, 2, 1, True, nn.ReLU), (16, 24, 3, 2, 4, False, nn.ReLU),
(24, 24, 3, 1, 3, False, nn.ReLU), (24, 40, 5, 2, 3, True, nn.Hardswish),
(40, 40, 5, 1, 3, True, nn.Hardswish), (40, 48, 5, 1, 3, True, nn.Hardswish),
(48, 96, 5, 2, 6, True, nn.Hardswish), (96, 96, 5, 1, 6, True, nn.Hardswish),
]
sd_probs = np.linspace(0, 0.2, len(configs))
self.blocks = nn.ModuleList([
IRB3D(inp, oup, k, s, exp, se, nl, sd_p)
for (inp, oup, k, s, exp, se, nl), sd_p in zip(configs, sd_probs)
])
self.fusion = MultiViewFusion(96)
self.conv_head = conv_bn_act(96, 576, 1, 1)
self.pool = nn.AdaptiveAvgPool3d(1)
self.classifier = nn.Sequential(
nn.Linear(576, 256),
nn.Hardswish(inplace=True),
nn.Dropout(0.2),
nn.Linear(256, num_classes)
)
def forward(self, x):
x = self.stem(x)
for blk in self.blocks:
x = blk(x)
x = self.fusion(x)
x = self.conv_head(x)
x = self.pool(x)
x = x.flatten(1)
x = self.classifier(x)
return x
def load_nifti(path):
try:
img = nib.load(str(path))
return img.get_fdata().astype(np.float32)
except Exception as e:
print(f"Failed to load {path}: {e}")
return None
def normalize_intensity(img):
if img.max() == img.min():
return np.zeros_like(img, dtype=np.float32)
nonzero = img[img > 0]
if len(nonzero) == 0:
return img.astype(np.float32)
p1, p99 = np.percentile(nonzero, [1, 99])
img = np.clip(img, p1, p99)
img = (img - img.min()) / (img.max() - img.min() + 1e-8)
return img.astype(np.float32)
def resize3d(img, target, is_label=False):
if img.size == 0:
return np.zeros(target, dtype=np.float32)
zoom_factors = [t / s for s, t in zip(img.shape, target)]
order = 0 if is_label else 1
return ndimage.zoom(img, zoom_factors, order=order).astype(np.float32)
def preprocess_nifti_pair(t1_path, flair_path, target_size=IMG_SIZE):
t1 = load_nifti(t1_path)
flair = load_nifti(flair_path)
if t1 is None or flair is None:
raise FileNotFoundError("One or both input NIfTI files could not be loaded.")
t1 = normalize_intensity(t1)
flair = normalize_intensity(flair)
t1 = resize3d(t1, target_size)
flair = resize3d(flair, target_size)
img = np.stack([t1, flair], axis=0).astype(np.float32)
return torch.from_numpy(img).unsqueeze(0)
def predict(model_path, t1_file_path, flair_file_path):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")
model = MobileNetV3Small3D(num_classes=NUM_CLASSES, in_ch=IN_CHANNELS)
try:
checkpoint = torch.load(model_path, map_location=device)
if 'model' in checkpoint:
model.load_state_dict(checkpoint['model'])
else:
model.load_state_dict(checkpoint)
model.to(device)
model.eval()
print(f"Successfully loaded model from {model_path}")
except Exception as e:
print(f"Error loading model weights from {model_path}: {e}")
print("Please ensure your .pth file contains the correct state_dict for MobileNetV3Small3D.")
sys.exit(1)
try:
input_tensor = preprocess_nifti_pair(t1_file_path, flair_file_path)
input_tensor = input_tensor.to(device)
except Exception as e:
print(f"Error during file preprocessing: {e}")
sys.exit(1)
with torch.no_grad():
logits = model(input_tensor)
probabilities = F.softmax(logits, dim=1).squeeze(0)
predicted_class = probabilities.argmax().item()
print("\n--- Prediction Results ---")
print(f"Input T1w file: {t1_file_path}")
print(f"Input FLAIR file: {flair_file_path}")
print(f"Probabilities (Class 0, Class 1): {probabilities.cpu().numpy()}")
print(f"Predicted Class Index: {predicted_class}")
if NUM_CLASSES == 2:
class_labels = {0: "Control", 1: "FCD"}
predicted_label = class_labels.get(predicted_class, f"Class {predicted_class}")
print(f"Predicted Label: {predicted_label}")
else:
print("Model has more than 2 classes. Prediction is based on class index.")
def main():
parser = argparse.ArgumentParser(description="Run 3D CNN prediction using a trained .pth model.")
parser.add_argument("model_path", type=str, help="Path to the MobileNetV3Small3D model file (.pth).")
parser.add_argument("t1_file", type=str, help="Path to the T1w NIfTI file for prediction.")
parser.add_argument("flair_file", type=str, help="Path to the FLAIR NIfTI file for prediction.")
example_t1 = "./path/to/sub-01_T1w.nii.gz"
example_flair = "./path/to/sub-01_FLAIR.nii.gz"
example_model = "./out/models/best_model.pth"
if len(sys.argv) == 1:
print("--- Usage Example ---")
print(f"python {Path(sys.argv[0]).name} {example_model} {example_t1} {example_flair}")
print("\n--- Note ---")
print("This script requires a paired T1w and FLAIR NIfTI file as input, as the original model expects 2 channels.")
sys.exit(1)
args = parser.parse_args()
predict(args.model_path, args.t1_file, args.flair_file)
if __name__ == "__main__":
import warnings
warnings.filterwarnings("ignore")
main()
```
4. **Execution:** Run the script from your terminal:
```bash
python predict.py /path/to/your/model.pth /path/to/your/T1w.nii.gz /path/to/your/FLAIR.nii.gz
```
### Visual gallery
### Pipeline architecture
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.
### Multi-view slices
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.
### Training curves
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.
### Embeddings and harmonization
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.
### Cohort composition
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.
### Prediction dashboard
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.
### Asymmetry examples
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.
### Grad-CAM saliency montages (3D and overlays; for 10 samples)
**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.
## ⚖️ Ethics and Limitations
### Clinical Risk and Human Oversight
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.
### Bias and Generalization
The model's performance is intrinsically linked to the composition of its training data.
* **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.
* **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.
### Input Modality Validation
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.
|