Cardiac-Scribble-Segmentation

Weakly-supervised multi-class cardiac MRI segmentation using sparse scribble annotations, compared against a fully-supervised baseline, with a confidence-masked pseudo-labeling improvement to close the gap. Built on the ACDC dataset.

Problem Statement

Develop a weakly supervised deep learning framework for multi-class cardiac MRI segmentation using sparse scribble annotations, aiming to produce accurate segmentation masks while reducing the need for dense pixel-level annotations.

Status — Complete

• Phase 1 — Dataset analysis (notebooks/eda.ipynb)
• Phase 2 — Data pipeline (dataset/cardiac_dataset.py)
• Phase 3 — Baseline U-Net, full supervision
• Phase 4 — Scribble-only training
• Phase 5 — Improvement: confidence-masked pseudo-labels
• Phase 6 — Evaluation (Dice / IoU / HD95 comparison)
• Phase 7 — Inference + demo

Results

Evaluated on a held-out patient-level test set (15 patients, 278 slices, never touched during training or hyperparameter selection).

Model	Dice	IoU	HD95 (px)
U-Net (full supervision)	0.9255	0.8669	2.97
Scribble U-Net	0.5977	0.4739	113.10
Scribble + pseudo-labels	0.6829	0.5580	100.52

Confidence-masked pseudo-labeling improved scribble-only Dice by +8.5 points (+13% relative) and reduced HD95 by ~11%, using only the same sparse scribble annotations plus the model's own high-confidence predictions — no additional annotation cost.

Per-class breakdown

Class	Full-sup Dice	Scribble Dice	Scribble+Pseudo Dice
Background	0.998	0.961	0.975
RV	0.896	0.477	0.701
Myocardium	0.866	0.579	0.570
LV	0.942	0.373	0.485

Notable asymmetry: pseudo-labeling improved RV substantially (+0.22 Dice) but left myocardium essentially flat (-0.01). Myocardium is a thin, ring-shaped structure surrounding the LV cavity — a wrong-but-confident pseudo-label near its narrow boundary does proportionally more damage than a similar error on RV's more compact blob shape, likely explaining the gap. This is a natural target for the boundary-aware improvement suggested below.

Repo Structure

dataset/        CardiacDataset — reads .h5 slices, resizes, augments
models/         U-Net architecture
losses/         DiceCELoss (full supervision), ScribbleLoss (partial CE/Dice),
                PseudoLabelScribbleLoss (Phase 5)
utils/          Dice/IoU/HD95 metrics
configs/        YAML configs per experiment
train.py        Shared training entry point
notebooks/      EDA + full Kaggle training/evaluation notebooks
checkpoints/    Saved model weights (best + last, per run)
results/        final_comparison.csv, comparison_plot.png, demo images
figures/        Additional plots for README/presentation

Data

ACDC (Automated Cardiac Diagnosis Challenge), 100 patients, 1902 2D slices after deduplication. Each .h5 file contains image (float32), label (dense mask, uint8, classes 0-3: background/RV/myocardium/LV), and scribble (sparse mask, uint16, classes 0-3 same as label + class 4 = unannotated/ignore).

Patient-level 70/15/15 train/val/test split — no slice from the same patient appears in more than one split. No pixel-spacing metadata present in the source .h5 files, so all images/masks are resized to 256x256 (bilinear for images, nearest-neighbor for masks).

Method

Phase 3 — Baseline: standard 4-level U-Net (~31M params), trained on dense masks with a combined Dice + Cross-Entropy loss.

Phase 4 — Scribble supervision: same architecture, trained only on sparse scribble annotations using Partial Cross-Entropy + Partial Dice (loss computed only on annotated pixels; unannotated pixels are masked out via ignore_index).

Phase 5 — Confidence-masked pseudo-labels: on top of the real scribble loss, the model's own predictions on unannotated pixels are used as additional supervision — but only where prediction confidence exceeds 0.9, and only where no real scribble annotation exists (never overrides real labels). Pseudo-label loss weight ramps linearly from 0 to 0.3 over the first 15 epochs, preventing the model from reinforcing early, unreliable predictions.

All models are validated/tested against the same dense ground-truth label masks regardless of what supervision they were trained with, ensuring a fair comparison.

Setup

pip install -r requirements.txt

Training

See notebooks/ for the full Kaggle-ready training and evaluation notebooks (recommended — includes GPU-friendly checkpointing). Or use the standalone scripts:

python train.py --config configs/baseline.yaml --epochs 2   # smoke test first
python train.py --config configs/baseline.yaml               # Phase 3
python train.py --config configs/scribble.yaml                # Phase 4

Future Improvements

Not implemented here, but the natural next steps to push results further:

• EMA teacher + consistency loss: a slower-moving "teacher" copy of the model (updated via exponential moving average of the student's weights) generates more stable pseudo-labels than the student's own live predictions, typically reducing pseudo-label noise significantly.
• Class-specific confidence thresholds: given myocardium didn't benefit from pseudo-labeling while RV did, a higher confidence threshold specifically for myocardium's ring-boundary pixels (or excluding pseudo-labels within N pixels of the predicted myocardium boundary) could avoid the thin-structure error amplification seen here.
• Boundary-aware loss term: an additional loss penalizing boundary disagreement directly (e.g. a signed distance transform loss) would likely help HD95 more directly than Dice/CE alone, since HD95 is a boundary metric and neither current loss optimizes for it explicitly.
• Test-time augmentation: averaging predictions over flipped/rotated versions of each test slice is a cheap way to typically gain 1-2 Dice points with no retraining.
• 5-fold cross-validation: current results are from a single 70/15/15 split; averaging over multiple folds would give tighter, more publication-grade confidence in the reported numbers.

Demo

See results/demo_*.png for example inference outputs (MRI -> predicted mask -> overlay -> Dice).