README.md

---
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.

![Architecture](vis/architecture.png)

### 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.

![Multiview Sample](vis/multiview_sample.png)

### 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.

![Training Curves](vis/training_curves.png)

### 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.

![Harmonization](vis/harmonization.png)  

![UMAP/t-SNE (harmonized)](vis/harmonized_embeddings.png)

### 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.

![Cohort Splits](vis/cohort_splits.png)  

![Cohort Grid](vis/cohort_grid.png)

### 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.

![Prediction Dashboard](vis/prediction_dashboard.png)

### 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.

![Asymmetry 0](vis/asymmetry_0.png)  

![Asymmetry 1](vis/asymmetry_1.png)  

![Asymmetry 2](vis/asymmetry_2.png)

### 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.
![Grad-CAM 3D 0](saliency/sample_0_cam_3d.png)  

![Grad-CAM Overlay 0](saliency/gradcam_overlay_0.png)  

![Grad-CAM 3D 1](saliency/sample_1_cam_3d.png)  

![Grad-CAM Overlay 1](saliency/gradcam_overlay_1.png)  

![Grad-CAM 3D 2](saliency/sample_2_cam_3d.png)  

![Grad-CAM Overlay 2](saliency/gradcam_overlay_2.png)  

![Grad-CAM 3D 3](saliency/sample_3_cam_3d.png)  

![Grad-CAM Overlay 3](saliency/gradcam_overlay_3.png)  

![Grad-CAM 3D 4](saliency/sample_4_cam_3d.png)  

![Grad-CAM Overlay 4](saliency/gradcam_overlay_4.png)  

![Grad-CAM 3D 5](saliency/sample_5_cam_3d.png)  

![Grad-CAM Overlay 5](saliency/gradcam_overlay_5.png)  

![Grad-CAM 3D 6](saliency/sample_6_cam_3d.png)  

![Grad-CAM Overlay 6](saliency/gradcam_overlay_6.png)  

![Grad-CAM 3D 7](saliency/sample_7_cam_3d.png)  

![Grad-CAM Overlay 7](saliency/gradcam_overlay_7.png)  

![Grad-CAM 3D 8](saliency/sample_8_cam_3d.png)  

![Grad-CAM Overlay 8](saliency/gradcam_overlay_8.png)  

![Grad-CAM 3D 9](saliency/sample_9_cam_3d.png)  

![Grad-CAM Overlay 9](saliency/gradcam_overlay_9.png)

## ⚖️ 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.