Swinunetr Pretrained Brats2025

---

tags:

• medical-imaging
• brain-tumor-segmentation
• brats
• swinunetr
• pytorch
• monai

---

SwinUNETR Brain Tumor Segmentation (BraTS)

This repository provides a trained SwinUNETR model for brain tumor segmentation on multi-modal MRI scans from the BraTS dataset.

The model takes 4 MRI modalities as input and predicts 3 tumor regions.

Channel	Tumor Region
0	Enhancing Tumor (ET)
1	Tumor Core (TC)
2	Whole Tumor (WT)

The model is implemented using PyTorch and MONAI.

---

Model Architecture

The model is based on SwinUNETR, a transformer-based 3D segmentation architecture.

Configuration:

• Input channels: 4
• Output channels: 3
• Feature size: 48
• Architecture: Swin Transformer + UNet decoder
• Inference method: Sliding Window Inference

---

Cross-Validation Performance

The model was trained using 5-fold cross-validation for 200 epochs per fold.
Metrics reported below correspond to the final model checkpoint, not the best validation epoch.

Epochs Trained	Fold	Best Epoch	Dice TC	Dice WT	Dice ET	Avg Dice
200	0	199	0.92289	0.93232	0.88946	0.91489
200	1	169	0.90106	0.92192	0.87619	0.89973
200	2	169	0.91470	0.93339	0.87918	0.90909
200	3	139	0.89212	0.92091	0.86040	0.89115
200	4	169	0.90400	0.92522	0.88650	0.90525
Average			0.90695	0.92675	0.87835	0.90402

Notes

• TC: Tumor Core
• WT: Whole Tumor
• ET: Enhancing Tumor
• The last row reports the average across folds.
• Metrics are computed using the final trained model checkpoint rather than the best validation epoch.

---

Requirements

Install dependencies:

pip install torch monai nibabel numpy

Recommended environment:

• Python >= 3.9
• PyTorch >= 2.0
• MONAI >= 1.3
• CUDA recommended for GPU inference

---

Expected Input Format

The model expects 4 MRI modalities stacked as channels.

shape = (4, H, W, D)

Modalities order:

[T1, T1ce, T2, FLAIR]

Images should be:

• skull stripped
• co-registered
• normalized (handled by MONAI transforms below)

---

Model Definition

import torch
import torch.nn as nn
from monai.networks.nets import SwinUNETR

class MONAI_Swin_UNet_Generator(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = SwinUNETR(
            in_channels=4,
            out_channels=3,
            feature_size=48,
            drop_rate=0.0,
            attn_drop_rate=0.0,
            dropout_path_rate=0.0,
            use_checkpoint=True,
        )

    def forward(self, x):
        return self.net(x)

model = MONAI_Swin_UNet_Generator()

---

Load Trained Weights

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

checkpoint = torch.load(
    "model_best_fold_0.pth",
    map_location=device
)

model.load_state_dict(checkpoint["state_dict"])
model.to(device)
model.eval()

---

Custom Label Transform

BraTS segmentation labels are converted into tumor region masks.

Original BraTS labels:

Label	Meaning
0	Background
1	Necrotic / Non-enhancing tumor
2	Edema
3	Enhancing tumor

These are converted into three region masks:

Channel	Region
0	Enhancing Tumor (ET)
1	Tumor Core (TC)
2	Whole Tumor (WT)

import torch
from monai.transforms import MapTransform

class ConvertToBraTSRegionsd(MapTransform):
    def __init__(self, keys):
        super().__init__(keys)
    
    def __call__(self, data):
        for key in self.keys:
            seg = data[key].clone()

            if seg.shape[0] == 1:
                seg = seg[0]

            mask_ET = (seg == 3)
            mask_TC = (seg == 1) | (seg == 3)
            mask_WT = (seg == 1) | (seg == 2) | (seg == 3)

            data[key] = torch.stack(
                [mask_ET, mask_TC, mask_WT],
                dim=0
            ).float()

        return data

---

Data Preprocessing

from monai import transforms

basic_transform = transforms.Compose([
    transforms.LoadImaged(keys=["image", "label"]),
    ConvertToBraTSRegionsd(keys=["label"]),
    transforms.NormalizeIntensityd(
        keys="image",
        nonzero=True,
        channel_wise=True
    ),
])

---

Dataset and DataLoader

from monai import data

test_ds = data.Dataset(
    data=test_files,
    transform=basic_transform
)

test_loader = data.DataLoader(
    test_ds,
    batch_size=1,
    shuffle=False,
    num_workers=1,
    pin_memory=True,
)

---

Sliding Window Inference

3D MRI volumes are processed using MONAI sliding window inference.

from monai.inferers import sliding_window_inference
from functools import partial

roi = (128,128,128)

model_inferer = partial(
    sliding_window_inference,
    roi_size=roi,
    sw_batch_size=1,
    predictor=model,
    overlap=0.7,
)

---

Run Inference

import numpy as np

with torch.no_grad():

    for batch_data in test_loader:

        image = batch_data["image"].to(device)

        prob = torch.sigmoid(model_inferer(image))

        seg = prob[0].cpu().numpy()

        seg = (seg > 0.5).astype(np.int8)

Output shape:

(3, H, W, D)

Each channel corresponds to one tumor region.

---

Quick Example (Single Volume)

image = torch.randn(1,4,128,128,128).to(device)

with torch.no_grad():
    prob = torch.sigmoid(model_inferer(image))
    seg = (prob > 0.5).int()

---

Output

The model produces three segmentation maps:

Channel	Region
0	Enhancing Tumor (ET)
1	Tumor Core (TC)
2	Whole Tumor (WT)

Output shape:

(3, H, W, D)

---

Citation

Original Paper,

Hatamizadeh et al.
Swin UNETR: Swin Transformers for Semantic Segmentation of Brain Tumors in MRI Images