Mask-Guided Multi-Channel SwinUNETR for Breast MRI Classification

This repository hosts the trained weights and inference pipeline for the Mask-Guided Multi-Channel SwinUNETR framework, which achieved second place in the multi-center ODELIA Challenge 2025 (held in conjunction with MICCAI 2025 and the Deep-Breath workshop). The model is designed to perform robust breast cancer classification from Dynamic Contrast-Enhanced (DCE) MRI scans.

For a comprehensive description of the methodology, evaluation, and clinical motivation, please refer to the preprint:
👉 Read the Paper (arXiv:2508.20621)

Model Description

• Developed by: Smriti Joshi, Lidia Garrucho, Richard Osuala, Oliver Diaz, Karim Lekadir (Universitat de Barcelona, Computer Vision Center, ICREA)
• Model Type: 3D Transformer-based classification framework with a SwinUNETR backbone.
• Input Channels: 4-channel maximum intensity projections (MIPs) derived from DCE-MRI:
  1. First post-contrast phase
  2. Subtraction 1 (first post-contrast minus pre-contrast)
  3. Subtraction 2 (second post-contrast minus pre-contrast)
  4. Last subtraction (last post-contrast minus pre-contrast)
• Output Classes: no lesion (normal), benign lesion, or malignant lesion (predicted for left and right breasts independently).
• Core Methodology:
  • Breast Region Masking: Preprocessing integrates a breast-segmentation-guided classification strategy to exclude background noise (e.g., chest wall, air), focusing feature representation strictly on breast tissue.
  • Enhancement Kinetics: The 4-channel input captures both anatomical enhancement and temporal dynamics of contrast uptake.
  • Ensemble Scheme: To manage severe class imbalance (predominance of normal cases), the final model ensembles predictions from two model variants: one trained with natural class weighting and another with inverse-frequency class weighting.

---

Setup & Dependencies

To execute the inference code, create a virtual environment and install the required libraries.

Using venv

# Create a virtual environment
python -m venv venv

# Activate the environment
# On Linux/Mac:
source venv/bin/activate
# On Windows:
# venv\Scripts\activate

# Install dependencies
pip install torch torchvision numpy huggingface_hub torchio matplotlib transformers monai albumentations

Using Conda

# Create a conda environment
conda create -n odelia_swinunetr python=3.10
conda activate odelia_swinunetr

# Install dependencies
pip install torch torchvision numpy huggingface_hub torchio matplotlib transformers monai albumentations

---

Inference & Getting Started

To classify a breast MRI study, download the corresponding network components from this repository:

from huggingface_hub import hf_hub_download

# Replace with your repository path once uploaded
repo_id = "your-username/mask-guided-swinunetr-odelia"

# Download essential model scripts and weights
hf_hub_download(repo_id=repo_id, filename="models.py", local_dir="./")
hf_hub_download(repo_id=repo_id, filename="preprocessing.py", local_dir="./")
hf_hub_download(repo_id=repo_id, filename="weights/swinunetr_natural.pt", local_dir="./")
hf_hub_download(repo_id=repo_id, filename="weights/swinunetr_weighted.pt", local_dir="./")

Inference Code Example

import torch
import numpy as np
from preprocessing import preprocess_study # Standardized TorchIO resizing & MIP extraction
from models import SwinUNETRClassifier # SwinUNETR classification backbone with lightweight linear head

# 1. Load and preprocess the raw multi-center breast DCE-MRI study (.nii.gz)
# This includes: resampling to 0.7 x 0.7 x 3 mm^3, cropping to breast region, splitting left/right halves,
# applying predicted breast masks, and projecting along the z-axis (MIP)
mips_left, mips_right = preprocess_study("path/to/patient_DCE_MRI.nii.gz")

# 2. Convert to tensors and normalize using dataset statistics
# Channels: (Phase 1, Subtraction 1, Subtraction 2, Last Subtraction)
channel_means = np.array([0.2074, 0.1290, 0.1396, 0.1470])
channel_stds = np.array([0.2110, 0.1629, 0.1620, 0.1626])

def prepare_tensor(mip_data):
    # mip_data shape: [4, H, W] (typically cropped to 256 x 256)
    normalized = (mip_data - channel_means[:, None, None]) / channel_stds[:, None, None]
    return torch.tensor(normalized, dtype=torch.float32).unsqueeze(0) # Batch size 1

input_left = prepare_tensor(mips_left).cuda()
input_right = prepare_tensor(mips_right).cuda()

# 3. Initialize models and load weights
# Model A: Trained with natural class weights
model_natural = SwinUNETRClassifier(img_size=(256, 256), in_channels=4, out_classes=3).cuda()
model_natural.load_state_dict(torch.load("weights/swinunetr_natural.pt"))
model_natural.eval()

# Model B: Trained with inverse-frequency class weights
model_weighted = SwinUNETRClassifier(img_size=(256, 256), in_channels=4, out_classes=3).cuda()
model_weighted.load_state_dict(torch.load("weights/swinunetr_weighted.pt"))
model_weighted.eval()

# 4. Perform Inference and Ensemble
classes = ["no lesion", "benign", "malignant"]

with torch.no_grad():
    # Left breast forward pass
    preds_nat_l = torch.softmax(model_natural(input_left), dim=1)
    preds_wt_l = torch.softmax(model_weighted(input_left), dim=1)
    ensemble_probs_l = (preds_nat_l + preds_wt_l) / 2.0
    
    # Right breast forward pass
    preds_nat_r = torch.softmax(model_natural(input_right), dim=1)
    preds_wt_r = torch.softmax(model_weighted(input_right), dim=1)
    ensemble_probs_r = (preds_nat_r + preds_wt_r) / 2.0

print(f"Left Breast Prediction: {classes[torch.argmax(ensemble_probs_l).item()]} (Probabilities: {ensemble_probs_l.cpu().numpy()})")
print(f"Right Breast Prediction: {classes[torch.argmax(ensemble_probs_r).item()]} (Probabilities: {ensemble_probs_r.cpu().numpy()})")

---

Training Data & Procedure

Segmentation Network

To generate precise masks, a SwinUNETR segmentation model was trained using 136 total cases from two open datasets:

1. Duke-Breast-Cancer-MRI (36 cases with manually annotated breast masks)
2. Breast-Cancer-DCE-MRI (100 cases from Yunnan Cancer Hospital with whole-breast segmentations) Trained for 100 epochs, batch size 20, learning rate 1e-4, using a cosine annealing schedule.

Classification Network

The classification backbone was trained on the ODELIA Challenge dataset:

• Total Cases: 511 multi-center breast MRI studies (each containing a T2-weighted and dynamic DCE-MRI series).
• Heterogeneity: Scans were acquired across six European centers (CAM, MHA, RSH, RUMC, UKA, UMCU) on both 1.5 T and 3 T scanners from multiple vendors.
• Class Distribution: Strongly imbalanced (66.8% No Lesion, 13.4% Benign, 19.8% Malignant).
• Optimization: Trained for 300 epochs, batch size 10, learning rate 1e-4, AdamW, and cosine annealing schedule with 5-epoch warm-up.
• Augmentation: Robust augmentation via Albumentations including flips, affine rotations, elastic deformations, grid/optical distortions, brightness/contrast scaling, and coarse dropout (missing signal simulation).

---

Evaluation Results

Evaluation was performed using five-fold cross-validation stratified by patient-level lesion labels, preventing patient-level leakage across folds.

Cross-Validation Performance (Averages)

• Natural Class Weighting (Model A): Tends to favor the majority "no lesion" class, achieving a higher overall micro-AUC but suffering from lowered sensitivity on benign and malignant lesions.
• Inverse-Frequency Weighting (Model B): Dramatically improves benign detection rates but increases false positives in the normal class.
• Ensemble (Final Model): Strikes an optimal balance, reducing prediction variance and stabilizing performance across the heterogeneous multi-center scans.

Held-Out Test Set Performance

The ensembled models across the five training folds achieved the following performance on the independent challenge test set:

Metric	Score
Micro-AUC	0.8610
Sensitivity (at 90% Specificity)	0.6201
Specificity (at 90% Sensitivity)	0.5678
Overall Score	0.6830

---

Citation

If you use this model or code in your research, please cite the corresponding paper:

@article{joshi2025mask,
  title={Mask-Guided Multi-Channel SwinUNETR Framework for Robust MRI Classification},
  author={Joshi, Smriti and Garrucho, Lidia and Osuala, Richard and Diaz, Oliver and Lekadir, Karim},
  journal={arXiv preprint arXiv:2508.20621},
  year={2025}
}

For questions regarding model architecture, training, or permissions, please contact Smriti Joshi (smriti.joshi@ub.edu).