kssrikar4/FCD_Detection

🧠 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
Batches/Steps	BATCH_SIZE = 16, GRAD_ACCUM_STEPS = 4
Epochs/Stopping	Up to 100 epochs
Data Split	Subject-Level Split (Ensured no data leakage between subjects)
Augmentation	TorchIO

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:

   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.

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:

   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.