SwinUNETR 3D Brain Tumor Segmentation

Full-stack medical imaging system for 3D brain tumor segmentation. Trains a SwinUNETR v2 model on BraTS 2021 MRI data, serves it via FastAPI on AWS, visualises results in a WebGL viewer with live surgical metrics and gesture control, and runs LangGraph AI agents for clinical interpretation, explainability, and PDF report generation.

Results

Trained on BraTS 2021 (1,251 cases, fold 0) — 100 epochs on RunPod H100 SXM.

Metric	Val (best ckpt)	Test (188 cases)
Mean Dice	0.9052	0.8961
Tumor Core (TC)	—	0.8787
Whole Tumor (WT)	—	0.9282
Enhancing Tumor (ET)	—	0.8814
HD95 Mean (mm)	—	5.62

The MONAI SwinUNETR v2 paper reports ~0.875 mean Dice on BraTS 2021 validation. This run achieves 0.9052 val / 0.8961 test using the same architecture with bfloat16 AMP and cosine annealing LR.

Quality gates (all passed at deployment)

Gate	Threshold	Achieved
Mean Dice	≥ 0.70	✅ 0.8961
Tumor Core Dice	≥ 0.65	✅ 0.8787
Whole Tumor Dice	≥ 0.75	✅ 0.9282
Enhancing Tumor Dice	≥ 0.55	✅ 0.8814

What it does

Train — SwinUNETR v2 on 4-modality BraTS 2021 MRI (FLAIR, T1, T1ce, T2); AMP bfloat16; tracked with W&B + MLflow
Serve — FastAPI inference server on AWS g5.xlarge spot (A10G 24GB): upload 4 MRI files, get back segmentation mask + slice overlays + confidence scores
Explain — Gradient saliency maps (per tumour class) + test-time augmentation uncertainty + MC Dropout epistemic uncertainty
Interpret — ExplainabilityAgent: interprets confidence scores + uncertainty into a clinical reliability report
Visualise — Browser viewer: 3D volumetric render, MPR, surgical metrics panel, clinical findings, webcam hand gesture control
Analyse — ClinicalAgent (Claude via LangGraph): clinical interpretation from voxel volumes + formal PDF report
Demo — Pre-loaded BraTS cases served from S3 — no data download needed to try the app

Project structure

├── src/
│   ├── train.py           # SwinUNETR v2 training loop (AMP bfloat16, W&B + MLflow)
│   ├── dataset.py         # BraTS 2021 MONAI DataLoader + transforms + dataset.json generator
│   ├── model.py           # SwinUNETRV2 model factory (with MC Dropout support)
│   ├── app.py             # FastAPI inference server
│   ├── explainability.py  # Gradient saliency + TTA uncertainty + MC Dropout UQ
│   ├── federated.py       # Flower FL client (BraTSClient + data partitioning)
│   ├── evaluate.py        # Standalone evaluation script with quality gates
│   ├── metrics.py         # Dice + HD95 metrics
│   └── utils.py           # Sliding-window inferer, checkpoint helpers
│
├── agents/
│   ├── base.py                  # LLM factory (Anthropic/OpenRouter) + BaseGraph wrapper
│   ├── graph_state.py           # AgentState TypedDict + Pydantic input models
│   ├── tools.py                 # @tool functions: volumes, W&B, quality gates, PDF
│   ├── clinical_agent.py        # Tumour volume interpretation → clinical report
│   ├── explainability_agent.py  # Confidence + uncertainty interpretation
│   ├── evaluation_agent.py      # Quality gate checks + PASS/FAIL verdict
│   ├── data_agent.py            # BraTS data validation + dataset.json generation
│   ├── supervisor.py            # Top-level agent: routes intent to sub-agents
│   └── report_agent.py          # Writes formal report sections + assembles PDF
│
├── viewer/
│   └── index.html         # 3D viewer (NiiVue WebGL + MediaPipe Hands gesture control)
│
├── infra/
│   ├── main.tf            # Terraform: EC2 g5.xlarge spot, S3, IAM, CloudWatch, EIP
│   ├── variables.tf       # Input variables
│   └── userdata.sh        # EC2 boot script: clone, install, download model, start server
│
├── scripts/
│   ├── orchestrate.py         # CLI for all agent modes (chat/train/eval/clinical/report)
│   ├── upload_demo_cases.sh   # Upload BraTS cases to S3 for the demo picker
│   └── fl_simulate.py         # Flower federated learning simulation
│
├── .github/workflows/
│   └── deploy.yml         # CI/CD: Terraform apply on infra changes, SSH deploy on code changes
│
├── Docker/
│   ├── Dockerfile         # Inference API image (CPU)
│   ├── Dockerfile.gpu     # Inference API image (CUDA)
│   └── Dockerfile.train   # Training image for RunPod
│
├── config/config.yaml     # Model hyperparameters, speed flags, UQ + FL settings
├── requirements.txt       # Full stack (training + API server + agents)
├── requirements-train.txt # Training only
└── .env.example           # Environment variables reference

Quickstart

Try the demo

Open the deployed viewer — no data download, no setup. Pre-loaded BraTS cases available from the in-app picker.

Run locally

git clone https://github.com/moebouassida/SwinUNETR-3D-Brain-Segmentation
cd SwinUNETR-3D-Brain-Segmentation
pip install -r requirements.txt
cp .env.example .env   # fill in API keys
export CHECKPOINT_PATH=runs/best_fold0.pt
uvicorn src.app:app --host 0.0.0.0 --port 8000
# open http://localhost:8000/viewer

Training only (minimal deps)

pip install -r requirements-train.txt   # torch, monai, wandb — nothing else

Only WANDB_API_KEY needed. Agent keys are not required for training.

Data

Download BraTS 2021 Training Data and place it under Data/:

Data/
  BraTS2021_00000/
    BraTS2021_00000_flair.nii.gz
    BraTS2021_00000_t1.nii.gz
    BraTS2021_00000_t1ce.nii.gz
    BraTS2021_00000_t2.nii.gz
    BraTS2021_00000_seg.nii.gz
  BraTS2021_00001/ ...

Generate the dataset manifest:

python -c "from src.dataset import generate_dataset_json; generate_dataset_json('./Data')"

Train on RunPod

Tested from Windows — all setup runs inside the pod.

Recommended GPUs (batch_size=1, roi=128³, gradient checkpointing on):

GPU	VRAM	~$/hr	100 epochs
RTX 4090	24 GB	$0.74	~18 hrs
A100 80GB	80 GB	$1.99	~8 hrs
H100 SXM	80 GB	$3.99	~5 hrs

Use On-Demand, not Spot, to avoid interruption mid-training. Template: RunPod PyTorch 2.x, Volume: 200 GB.

git clone https://github.com/moebouassida/SwinUNETR-3D-Brain-Segmentation.git
cd SwinUNETR-3D-Brain-Segmentation
pip install -r requirements-train.txt

# Download BraTS via Kaggle
pip install kaggle
export KAGGLE_USERNAME=your_username
export KAGGLE_KEY=your_api_key
mkdir -p /runpod-volume/data && cd /runpod-volume/data
kaggle datasets download -d dschettler8845/brats-2021-task1
unzip brats-2021-task1.zip

# Train
cd ~/SwinUNETR-3D-Brain-Segmentation
export WANDB_API_KEY=your_key
python src/train.py --config config/config.yaml \
  --data-dir /runpod-volume/data \
  --save-dir /runpod-volume/runs

Best checkpoints are automatically uploaded to W&B Artifacts and survive pod termination.

Resume after interruption:

python src/train.py --config config/config.yaml \
  --data-dir /runpod-volume/data \
  --save-dir /runpod-volume/runs \
  --resume /runpod-volume/runs/last.pt

Deploy on AWS

Full infrastructure-as-code via Terraform. Deploys to eu-central-1 (Paris) on a g5.xlarge spot instance (A10G 24GB VRAM, ~$0.40/hr). CloudWatch streams server logs automatically.

One-time setup

Create an IAM user with AdministratorAccess and download the access key CSV
Create an EC2 key pair named brain-seg in eu-central-1

Deploy from pod (or any Linux machine with AWS CLI + Terraform)

# Install tools
curl -s "https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip" -o awscliv2.zip
unzip -q awscliv2.zip && ./aws/install
wget -q https://releases.hashicorp.com/terraform/1.9.0/terraform_1.9.0_linux_amd64.zip
unzip -q terraform_1.9.0_linux_amd64.zip && mv terraform /usr/local/bin/

# Configure AWS
aws configure   # paste key ID, secret, region: eu-central-1, output: json

# Create Terraform state bucket (one-time)
aws s3api create-bucket --bucket brain-seg-tfstate-moez --region eu-central-1 \
  --create-bucket-configuration LocationConstraint=eu-central-1

# Create key pair (one-time)
aws ec2 create-key-pair --key-name brain-seg --region eu-central-1 \
  --query 'KeyMaterial' --output text > ~/.ssh/brain-seg.pem
chmod 400 ~/.ssh/brain-seg.pem

# Fill variables
cd infra
cat > terraform.tfvars <<EOF
key_name           = "brain-seg"
wandb_api_key      = "your_wandb_key"
openrouter_api_key = "your_openrouter_key"
langchain_api_key  = "your_langchain_key"
sentry_dsn         = ""
EOF

terraform init && terraform apply
terraform output viewer_url   # your live URL

Upload demo cases

bash scripts/upload_demo_cases.sh brain-seg-demo-cases-moez /path/to/Data eu-central-1 5

CI/CD

Every push to main triggers GitHub Actions:

infra changes → terraform apply recreates infrastructure
code changes → SSH into running instance, git pull, restart server

Add these secrets to your repo (Settings → Secrets → Actions):

Secret	Value
`AWS_ACCESS_KEY_ID`	from IAM CSV
`AWS_SECRET_ACCESS_KEY`	from IAM CSV
`EC2_KEY_NAME`	`brain-seg-key`
`EC2_PRIVATE_KEY`	contents of `.pem` file
`EC2_HOST`	IP from `terraform output`
`WANDB_API_KEY`	wandb.ai → Settings
`OPENROUTER_API_KEY`	openrouter.ai
`LANGCHAIN_API_KEY`	smith.langchain.com → Settings
`DEMO_CASES_BUCKET`	`brain-seg-demo-cases-moez`

API endpoints

Method	Endpoint	Description
`GET`	`/health`	Liveness check, checkpoint status, device
`POST`	`/predict`	4 NIfTI files → segmentation mask (`.nii.gz`)
`POST`	`/predict/slices`	4 NIfTI files → voxel counts + PNG overlay slices
`POST`	`/predict/explain`	4 NIfTI files → saliency maps + uncertainty + confidence
`POST`	`/explain/interpret`	Confidence scores → explainability report
`POST`	`/analyze`	Voxel counts → clinical analysis
`POST`	`/report`	Segmentation + explain data → PDF report download
`GET`	`/demo-cases`	List available pre-loaded demo cases
`GET`	`/demo-cases/{id}/{modality}`	Stream a demo case modality (flair/t1/t1ce/t2)
`GET`	`/viewer`	Serves the 3D viewer

Typical workflow:

/predict/slices  →  /analyze                              (clinical report from volumes)
                 →  /predict/explain → /explain/interpret  (confidence + uncertainty)
                 →  /report                               (PDF combining all of the above)

OR Viewer

Open at http://localhost:8000/viewer or the deployed URL.

Feature	Detail
Load case	Demo Cases picker (S3-backed, no download needed) or drag & drop your own folder
View modes	3D volumetric render, Axial, Coronal, Sagittal, MPR
MPR	3 independent NiiVue instances; each panel scrolls its own axis
Analyze	Runs inference; segmentation overlaid on MRI, surgical metrics populated
Metrics panel	TC / WT / ET volumes (cm³) + estimated diameters (top-right)
AI Analysis panel	Confidence scores + AI clinical assessment (bottom-right)
Gesture control	MediaPipe Hands via webcam — pinch-to-click, palm-to-rotate, fist-hold-to-analyze

Training details

Setting	Value
Architecture	SwinUNETR v2 (MONAI)
Input	`(1, 4, 128, 128, 128)` — 4-channel BraTS MRI
Output	`(1, 3, 128, 128, 128)` — TC + WT + ET (3 independent binary channels)
Activation	Sigmoid per channel + threshold 0.5
Loss	Dice + Cross-Entropy
Metrics	Dice (TC, WT, ET, mean) + HD95
Optimiser	AdamW, lr=1e-4, weight_decay=1e-5
LR schedule	CosineAnnealingLR over 100 epochs
Precision	AMP bfloat16
Gradient checkpointing	Enabled — saves ~6 GB VRAM
Inference	Sliding window, overlap 0.5, Gaussian weighting
Tracking	W&B + MLflow

BraTS label convention

BraTS 2021 uses 3 overlapping binary channels — not mutually exclusive classes:

Channel	Region	Nesting
Label 1	Tumor Core (TC) — necrotic core + non-enhancing tumor	ET ⊂ TC
Label 2	Whole Tumor (WT) — entire mass including peritumoral edema	TC ⊂ WT
Label 3	Enhancing Tumor (ET) — active, blood-brain barrier breakdown	ET ⊂ TC ⊂ WT

This is why sigmoid is used instead of softmax — regions overlap.

AI Agent System

Built on LangGraph 1.x with Claude (via Anthropic API or OpenRouter). LangSmith tracing is automatic when LANGCHAIN_API_KEY is set.

Agents

Agent	Role
`ClinicalAgent`	Interprets segmentation volumes → structured clinical report (WHO grade, urgency, findings)
`ExplainabilityAgent`	Interprets model confidence + TTA/MC uncertainty → prediction reliability report
`DataAgent`	Validates BraTS directory, generates dataset.json
`EvaluationAgent`	Checks checkpoint metrics against quality gates
`ReportAgent`	Writes formal radiology sections → assembles PDF
`SupervisorAgent`	Routes natural-language requests to sub-agents

Production features

Fallback: Claude Sonnet → Claude Haiku on failure
Rate limiter: 0.4 req/s, burst 5 (shared across all agents)
Retry: tenacity (3 attempts, exponential back-off) on all external API calls
Memory: MemorySaver checkpointer — multi-turn conversation history per thread
Tracing: LangSmith automatic tracing when LANGCHAIN_API_KEY is set
LLM routing: ANTHROPIC_API_KEY → direct; OPENROUTER_API_KEY → via OpenRouter

CLI

python scripts/orchestrate.py --mode chat
python scripts/orchestrate.py --mode eval --checkpoint runs/best_fold0.pt
python scripts/orchestrate.py --mode clinical --voxel-counts '{"1":5000,"2":15000,"3":2000}'

Explainability & Uncertainty Quantification

Method	What it shows	Cost
Gradient saliency	Voxels most influential per tumour class	3 backward passes
TTA uncertainty	Shannon entropy across 8 flip-augmented predictions (aleatoric)	8 forward passes
MC Dropout uncertainty	Epistemic uncertainty via N stochastic passes	N forward passes

Enable MC Dropout via ?mc_passes=20 on /predict/explain. Requires mc_dropout_rate: 0.1 in config.

Evaluate

python src/evaluate.py \
  --checkpoint runs/best_fold0.pt \
  --config config/config.yaml \
  --split test \
  --data-dir ./Data

Exits 0 if all quality gates pass, 1 if any fail — suitable for CI/CD deployment gates.

Environment variables

# LLM — choose one
ANTHROPIC_API_KEY=sk-ant-...
OPENROUTER_API_KEY=sk-or-...

# Training
WANDB_API_KEY=...

# Deployment
CHECKPOINT_PATH=runs/best_fold0.pt
DEMO_CASES_DIR=demo_cases        # local path, populated from S3 at boot

# Observability (optional)
LANGCHAIN_API_KEY=ls__...
LANGCHAIN_TRACING_V2=true
LANGCHAIN_PROJECT=swin-unetr

Tech stack

Layer	Technology
Model	SwinUNETR v2 (MONAI + PyTorch)
Precision	AMP bfloat16 + optional torch.compile
Training platform	RunPod GPU pods
Experiment tracking	W&B + MLflow
Inference API	FastAPI + MONAI sliding-window inferer
Rate limiting	slowapi (6 req/min per IP on inference endpoints)
Explainability	Gradient saliency + TTA uncertainty + MC Dropout
Agent framework	LangGraph 1.x + Claude (Anthropic / OpenRouter)
Agent tracing	LangSmith
PDF generation	fpdf2
3D viewer	NiiVue (WebGL)
Gesture control	MediaPipe Hands
Reverse proxy / TLS	nginx + self-signed certificate (HTTPS)
Process management	systemd (auto-restart on crash)
Error tracking	Sentry
Infrastructure	Terraform + AWS EC2 g5.xlarge spot + S3 + CloudWatch
CI/CD	GitHub Actions (Terraform deploy + `systemctl restart`)
Containerisation	Docker (CPU, GPU, training variants)
Federated learning	Flower (simulated FedAvg — not used for the deployed checkpoint)
Code quality	pre-commit + ruff

Federated Learning

Simulates N hospital clients each training on a private non-overlapping BraTS shard, aggregated with FedAvg via Flower.

python scripts/fl_simulate.py \
  --config config/config.yaml \
  --num-clients 4 \
  --num-rounds 10 \
  --local-epochs 1

Known limitations

Issue	Workaround
`/predict/explain` OOM on ≤16 GB VRAM	Saliency skipped gracefully; TTA still runs
MC Dropout requires retraining	Set `mc_dropout_rate: 0.1` before training
No DICOM support	Convert with `dcm2niix` before uploading

Security & Compliance

Document	Contents
docs/model-card.md	Architecture, training data, evaluation results, limitations, biases
docs/responsible-ai.md	EU AI Act Article 13 transparency — risk classification, human oversight, data governance
SECURITY.md	Vulnerability disclosure policy

EU AI Act: This system is classified as high-risk under Annex III (Category 5b — medical AI). It is deployed for research and demonstration only and has not undergone conformity assessment.

Agent security: All LLM agent inputs are validated (Pydantic), output is capped at 1 024 tokens, and all calls are traced in LangSmith. The question field is limited to 2 000 characters. Prompt injection hardening is embedded in every agent system prompt.

Data: Uploaded MRI files are processed in-memory and never written to disk or logged.

License

MIT

Navigation Menu

Repository files navigation

SwinUNETR 3D Brain Tumor Segmentation

Results

Quality gates (all passed at deployment)

What it does

Project structure

Quickstart

Try the demo

Run locally

Training only (minimal deps)

Data

Train on RunPod

Deploy on AWS

One-time setup

Deploy from pod (or any Linux machine with AWS CLI + Terraform)

Upload demo cases

CI/CD

API endpoints

OR Viewer

Training details

BraTS label convention

AI Agent System

Agents

Production features

CLI

Explainability & Uncertainty Quantification

Evaluate

Environment variables

Tech stack

Federated Learning

Known limitations

Security & Compliance

License

About

Releases

Packages

Contributors

Languages

Name		Name	Last commit message	Last commit date
Latest commit History 81 Commits
.github/workflows		.github/workflows
Docker		Docker
agents		agents
config		config
docs		docs
infra		infra
scripts		scripts
src		src
tests		tests
viewer		viewer
.Dockerignore		.Dockerignore
.env.example		.env.example
.gitattributes		.gitattributes
.gitignore		.gitignore
.pre-commit-config.yaml		.pre-commit-config.yaml
LICENSE		LICENSE
README.md		README.md
SECURITY.md		SECURITY.md
conftest.py		conftest.py
hf_spaces_app.py		hf_spaces_app.py
requirements-train.txt		requirements-train.txt
requirements.txt		requirements.txt

Folders and files

Latest commit

History

Repository files navigation

SwinUNETR 3D Brain Tumor Segmentation

Results

Quality gates (all passed at deployment)

What it does

Project structure

Quickstart

Try the demo

Run locally

Training only (minimal deps)

Data

Train on RunPod

Deploy on AWS

One-time setup

Deploy from pod (or any Linux machine with AWS CLI + Terraform)

Upload demo cases

CI/CD

API endpoints

OR Viewer

Training details

BraTS label convention

AI Agent System

Agents

Production features

CLI

Explainability & Uncertainty Quantification

Evaluate

Environment variables

Tech stack

Federated Learning

Known limitations

Security & Compliance

License

About

Resources

Security policy

Stars

Watchers

Forks

Releases

Packages

Contributors

Languages