Four model families, benchmarked honestly — with a test suite, CI and a containerised inference service.

Results · Where it fails · Install · Serving · Testing · Limitations

---

---

📋 Overview

A medical imaging pipeline covering tumor segmentation, detection, and through-plane slice interpolation, built end-to-end on one 4 GB laptop GPU.

The starting point was a published reference implementation (sdsubhajitdas/Brain-Tumor-Segmentation, 2-D U-Net, 0.745 Dice on the Cheng dataset). Phase 0 reproduces that result in a clean codebase as an honest baseline; the rest builds on it.

Warning

Not a clinical tool. Trained on public research data, never validated prospectively, never reviewed by a radiologist. See Limitations.

What's here that the reference doesn't have

	Reference	NeuroSeg
Data	2-D, single modality, 3,064 slices	3-D multimodal BraTS + the 2-D Cheng set
Architectures	One 2-D U-Net	U-Net → Attention U-Net → 3-D U-Net → UNETR
Detection	—	YOLOv8 on mask-derived boxes
Interpolation	—	Residual CNN for anisotropic volumes
Metrics	Dice	Dice, IoU, Hausdorff, sens/spec, mAP, PSNR/SSIM
Splits	slice-level	patient-disjoint (harder, and correct)
Engineering	notebook + one script	354 tests · CI · Docker · FastAPI · MLflow · DVC

🗺️ Pipeline architecture

graph LR
    subgraph Data
        D1["Cheng Dataset<br/>2D T1ce slices"]
        D2["BraTS / MSD<br/>3D multimodal"]
    end
    D1 --> Prep["Preprocessing<br/>patient-disjoint splits"]
    D2 --> Prep
    Prep --> Seg2D["2D Segmentation<br/>U-Net / Attention U-Net"]
    Prep --> Seg3D["3D Segmentation<br/>3D U-Net / UNETR"]
    Seg2D --> Det["Detection<br/>YOLOv8 on mask-derived boxes"]
    Prep --> Interp["Slice Interpolation<br/>Residual CNN"]
    Seg2D --> Eval["Evaluation<br/>Dice · IoU · Hausdorff · mAP · PSNR/SSIM"]
    Seg3D --> Eval
    Det --> Eval
    Interp --> Eval
    Eval --> Gate["Regression Gate (CI)"]
    Seg2D --> API["FastAPI Serving<br/>Dockerized"]
    Det --> API
    Interp --> API

Loading

---

📊 Results

Every number was produced on an RTX 3050 Laptop (4 GB) and is reproducible from this repo. Regenerate the whole table with python evaluate.py --task all --split test.

Task	Model	Test set	Headline	Claim
2-D segmentation	U-Net	462 slices	0.8244 Dice	✅ verified
2-D segmentation	Attention U-Net	462 slices	0.8202 Dice	⚪ not significant
3-D segmentation	3-D U-Net	9 volumes	0.8321 Dice	🟡 subset
3-D segmentation	UNETR	9 volumes	0.7348 Dice	🟡 under-trained
Detection	YOLOv8n	462 images	0.913 mAP@0.5	✅ verified
Interpolation	Residual CNN	360 slices	32.47 dB	✅ beats baseline

The four tasks are not comparable to one another — Dice, mAP and PSNR measure different things — which is why the harness emits one table per family rather than one wide table.

2-D segmentation · Cheng dataset

Model	Params	Dice	IoU	HD95 px	Sens.	Spec.	ms/case
Reference (published)	~7.8 M	0.745	—	—	—	—	—
U-Net	7.76 M	0.8244 ±.178	0.7315	8.27	0.8556	0.9974	9.85
Attention U-Net	7.89 M	0.8202 ±.184	0.7274	8.63	0.8617	0.9972	11.02

Important

The comparison against the reference is harder than it looks — in this repo's disfavour. The reference splits at the slice level. Cheng has 3,064 slices from only 233 patients — ~13 near-duplicate slices each — so a slice-level split puts near-identical images in both train and test and inflates the score. Every split here is patient-disjoint. The gain over 0.745 comes from the training setup, not an easier evaluation.

Attention gates: a negative result, reported as one

Both runs are identical in depth, width, loss, optimiser, schedule and split — only the gates differ. Attention led on validation (0.7727 vs 0.7682) and trailed on test. A ranking that flips between splits is the signature of noise, so it was tested rather than asserted:

mean difference (U-Net − Attention)   +0.0041
95% CI of difference                  [-0.0047, +0.0131]   ← includes zero
Wilcoxon signed-rank p                0.3125
paired t-test p                       0.3638
Attention better on                   217 / 462 cases  (47%)

The interval spans zero and the win rate is a coin flip: the extra 0.13 M parameters buy nothing measurable here. Reporting only the validation number — where attention "wins" — would have been misleading.

python scripts/compare_models.py \
    --a artifacts/phase0_unet2d/best.pt \
    --b artifacts/phase2_attention_unet/best.pt

3-D segmentation · BraTS volumes

60 Decathlon cases, split by case (42 / 9 / 9). Patch-based training at 96³ with foreground oversampling; evaluated by sliding-window inference over whole 240×240×155 volumes.

Model	Params	Whole	Core	Enhancing	Mean
3-D U-Net	1.40 M	0.8344	0.8217	0.8402	0.8321
UNETR (transformer)	102.24 M	0.7655	0.6566	0.7823	0.7348

Note

A 1.4 M CNN beats a 102 M transformer — because of data scale, not architecture. A U-Net has locality and translation equivariance built in; a ViT must learn them from data. This experiment has 42 training cases. UNETR's published results use the full BraTS set with far longer schedules. Both models' validation Dice was still rising at epoch 40 — both are under-trained, the transformer more severely. Ranking architectures from this would be a mistake.

Detection · YOLOv8

Metric	Ultralytics	This repo's implementation
mAP@0.5	0.9131	0.8967
mAP@0.5:0.95	0.5473	0.5100
Precision @0.25	0.8949	0.7822
Recall @0.25	0.8528	0.8939
TP / FP / FN	—	413 / 115 / 49

Both columns shown deliberately: the second is an independent implementation validated against hand-computed cases in the test suite. Agreeing to within ~0.02 mAP is a sanity check on both.

The gap between mAP@0.5 (0.91) and mAP@0.5:0.95 (0.55) is the real finding — the detector reliably finds the tumor but draws loose boxes, and at stricter IoU the score collapses.

Eight worst detections by IoU. 🟩 ground truth · 🟨 prediction with confidence. Most errors are loose boxes on the right lesion, not misses.

Slice interpolation

360 held-out slices, stride 2 — every second axial slice withheld and reconstructed from neighbours two apart.

Method	PSNR (dB)	SSIM
Linear average of neighbours	30.836 ±2.710	0.89911
Residual CNN	32.469 ±2.484	0.92298
Gain	+1.633	+0.02387

The baseline is the point of this table. Averaging two neighbouring slices is already strong, so a standalone 32 dB means nothing without knowing averaging alone reaches 30.8. The model predicts a residual on the linear average through a zero-initialised head — it starts training at exactly the baseline, so every parameter is spent on what averaging gets wrong.

---

🔍 Where the model fails

71.9% of cases score Dice ≥ 0.8, and the median (0.8837) sits well above the mean (0.8244) — the average is dragged down by a small tail of hard cases, not by broad mediocrity.

Tumor type	n	Dice		Lesion size	n	Median px	Dice
Meningioma	90	0.8928		Q1 smallest	116	260	0.7951
Pituitary	162	0.8743		Q2	115	664	0.8423
Glioma	210	0.7566		Q3	116	1133	0.8211
				Q4 largest	115	1906	0.8393

Gliomas are hardest by a wide margin — which is what a radiologist would predict. They are infiltrative with genuinely ambiguous margins, whereas meningiomas and pituitary adenomas are well-circumscribed. Lesion size matters much less than expected.

Caution

All five total failures predicted a tumor in the wrong place — none declined to predict. Five cases (1.1%) scored exactly 0.000 Dice. In every one, the model produced a confident, plausible-looking mask at the wrong location — in one case the contralateral hemisphere. That is the more dangerous failure: an empty prediction is obviously useless, but a well-formed mask in the wrong place is the kind of output a reader could act on.

Six worst test cases. 🟩 true positive · 🟥 false negative · 🟦 false positive

---

🚀 Quick start

# 1 · Install — GPU users: torch from the CUDA index FIRST (see warning below)
pip install torch==2.7.1+cu118 torchvision==0.22.1+cu118 \
    --index-url https://download.pytorch.org/whl/cu118
pip install -r requirements.txt

# 2 · Data — both public, no registration
python scripts/prepare_cheng.py        # ~880 MB
python scripts/fetch_msd_labels.py     # ~24 MB, not 7.6 GB (see below)

# 3 · Train
python -m src.train --config configs/phase0_unet2d.yaml
python scripts/train_brats3d.py --model unet3d --n-cases 60 --epochs 40
python scripts/train_yolo.py --epochs 40
python scripts/train_interpolation.py --n-cases 60 --epochs 25 --stride 2

# 4 · Benchmark, figures, error analysis
python evaluate.py --task all --split test
python scripts/visualize.py --checkpoint artifacts/phase0_unet2d/best.pt
python scripts/error_analysis.py

Warning

Install order matters. Installing MONAI without pinning torch lets pip resolve the CPU-only wheel over a CUDA build of the same version. Training then silently falls back to CPU — no error, ~30× slower. Verify with python -c "import torch; print(torch.version.cuda)" → must not be None.

Getting the BraTS labels: 24 MB instead of 7.6 GB

BraTS 2021 needs registration, so this uses the Medical Segmentation Decathlon archive — same data, redistributable, but one 7.6 GB tar. The obvious approach wastes hours:

curl -sL "$URL" | tar -x Task01_BrainTumour/labelsTr   # ran 1 h, extracted nothing

The assumption underneath it — that the archive is ordered imagesTr, labelsTr, imagesTs — is wrong. Probing with HTTP range requests shows the real layout:

dataset.json → imagesTr (484 vols, ~4.5 GB)
             → imagesTs (266 vols, ~2.9 GB)
             → labelsTr (~15 MB)   ← LAST

The labels are the last 15 MB of the file. Since S3 honours range requests, the whole problem collapses to a tail fetch: pull the last 24 MB, find the first 512-byte-aligned labelsTr tar header, and stream the remainder through Python's tarfile. All 484 labels, 24 MB instead of 7.6 GB — seconds instead of hours. That is scripts/fetch_msd_labels.py.

---

🐳 Serving

docker build -t neuroseg:latest .
docker run --rm -p 8000:8000 -v "$(pwd)/artifacts:/app/artifacts:ro" neuroseg:latest
curl localhost:8000/health

Endpoint	Purpose
GET /health	liveness + which capabilities actually loaded
GET /models	metadata for loaded checkpoints
POST /segment	slice → mask (JSON with RLE, or a PNG)
POST /detect	slice → bounding boxes derived from the mask
POST /interpolate	.npy stack → reconstructed intermediate slices

Verified end-to-end: image builds (2.2 GB), container healthy ~5 s after start, and posting a real test slice to /segment returned 1,032 predicted tumor pixels against 928 in the ground truth, with /detect returning one coherent box.

Missing checkpoints yield 503 on the affected endpoint, not a crash on boot — the container stays deployable while models are still training. Uploads accept .npy as well as PNG, because 8-bit images cannot represent the dynamic range of MRI.

---

🧪 Testing

pytest tests -m "not data"     # 354 tests, 82% coverage, no dataset required
pytest tests                   # adds assertions against the real cache

Module	Tests		Module	Tests
test_detection.py	49		test_benchmark_gate.py	31
test_models.py	47		test_interpolation.py	29
test_msd_and_config.py	42		test_api.py	23
test_patching.py	38		test_inference.py	16
test_data.py	37		test_patch_sampler.py	9
test_metrics.py	35

Everything runs on synthetic fixtures, so CI needs no dataset. Metrics are checked against hand-computed values, not the implementation's own output — Dice on a known half-overlap, IoU as 4/12, Hausdorff on a square shifted exactly 3 px, AP on a curve worked out by hand.

Six bugs found — four produced no error at all

Bug	Why it was invisible
Patch sampler replayed identical patches every epoch	Seeded seed + index; the loader walks the same range each epoch, so training saw ~40× less data than intended. Loss fell normally.
MONAI silently replaced CUDA torch with a CPU build	pip resolved the CPU wheel over the CUDA build of the same version. ~30× slowdown, no error.
mAP computed at the wrong confidence floor	Filtering at 0.25 truncates the PR curve — understated mAP by ~0.04.
Ultralytics wrote weights elsewhere	A relative project path resolves against Ultralytics' own runs_dir.
UNETR diverged at the shared learning rate	1e-3 works for the CNNs; the 102 M ViT went non-finite in one epoch. Now per-architecture LR + warmup.
2-D benchmark evaluated 3-D checkpoints	Once artifacts/ held both, the 2-D evaluator fed a 3-D model 2-D slices.

Two tests exist because the obvious assumption is wrong: Tversky(0.5, 0.5) equals Dice only at twice the smoothing constant, and a duplicate detection ranked below a true positive does not reduce AP — it only costs precision at a fixed operating point.

The regression gate handles metric direction explicitly: a fall in Dice is a regression, a rise in Hausdorff distance is. Both directions are tested — a gate never observed to fail is indistinguishable from one wired to always pass.

---

🗂️ Repository layout

graph TD
    Root["NeuroSeg/"]
    Root --> Src["src/"]
    Src --> Data["data/ — cheng · brats · msd loaders, transforms, patching, splits"]
    Src --> Models["models/ — unet2d · attention_unet · unet3d · unetr · losses · registry"]
    Src --> Det["detection/ — mask→bbox conversion, YOLO export"]
    Src --> Interp["interpolation/ — residual slice-interpolation net + triplet dataset"]
    Src --> Eval["eval/ — metrics, detection metrics, inference, report generation"]
    Src --> Serve["serving/ — FastAPI app"]
    Src --> Train["train.py — config-driven training entrypoint"]
    Root --> Scripts["scripts/ — prepare_cheng · fetch_msd_labels · train_yolo/brats3d/interpolation · visualize · error_analysis · compare_models"]
    Root --> Configs["configs/ — YAML experiment configs"]
    Root --> Tests["tests/ — 354 tests"]
    Root --> EvalPy["evaluate.py — benchmarking harness + regression gate"]

Loading

Architectures are built through src/models/registry.py, so a YAML config names a model and nothing downstream needs to know the class — which is what makes the benchmark a loop over configs rather than a chain of if-statements.

---

🧬 Datasets

Cheng (Figshare) — 3,064 T1ce slices, 233 patients, 708 meningioma / 1,426 glioma / 930 pituitary, with pixel-level masks. Public, no registration. Tumor occupies 1.69% of all pixels, which is why every loss here is imbalance-aware: plain cross-entropy is minimised well by predicting background everywhere, scoring >98% pixel accuracy and 0.0 Dice.

Medical Segmentation Decathlon Task01 — BraTS 2016/2017, 484 cases, 4 co-registered sequences at 1 mm isotropic. Two format differences are reconciled in src/data/msd.py, both silent failure modes if missed:

• Modality order — Decathlon packs FLAIR, T1w, T1gd, T2w; this codebase uses t1, t1ce, t2, flair. Getting it wrong trains on FLAIR while labelling the channel T1.
• Label codes — Decathlon uses 1=edema, 2=non-enhancing, 3=enhancing; BraTS uses 1=necrotic, 2=edema, 4=enhancing.

Results are reported on the three nested BraTS sub-regions (whole / core / enhancing), treated as independent binary channels because they overlap.

---

⚠ Limitations

• Not clinically validated. Public research data, retrospective, no prospective evaluation, no radiologist review. Nothing here is a medical device.
• The confident-wrong-location failure mode (1.1% of cases) is the most serious practical issue and is not mitigated in this codebase.
• Glioma performance (0.757) is meaningfully worse than the headline 0.824 — and gliomas are the hardest, most common malignant case.
• Cheng is single-modality, single-slice. A radiologist reads multiple sequences and adjacent slices together.
• The volumetric results use 60 of 484 cases with a 9-volume test set. Indicative, not comparable to published BraTS results.
• Both 3-D models are under-trained — validation Dice still climbing at epoch 40.
• Interpolation is evaluated on simulated anisotropy. BraTS volumes were already resampled to 1 mm isotropic, so neighbours are smoother than genuinely thick-slice clinical data; PSNR here is optimistic.
• Detection boxes are derived from masks, so the detector inherits every segmentation error rather than providing an independent check.
• No histology / whole-slide imaging — a different modality with different tooling. An adjacent interest, not covered work.

---

📚 References

• Ronneberger et al. (2015), U-Net: Convolutional Networks for Biomedical Image Segmentation
• Çiçek et al. (2016), 3D U-Net: Learning Dense Volumetric Segmentation from Sparse Annotation
• Oktay et al. (2018), Attention U-Net: Learning Where to Look for the Pancreas
• Hatamizadeh et al. (2021), UNETR: Transformers for 3D Medical Image Segmentation
• Cheng et al. (2015), brain tumor dataset — Figshare
• Antonelli et al. (2022), The Medical Segmentation Decathlon
• Baid et al. (2021), The RSNA-ASNR-MICCAI BraTS 2021 Benchmark

---

_{Built on one RTX 3050 Laptop GPU (4 GB). Every figure and number in this README is reproducible from the commands above.}