heartseg

Binary segmentation of the heart in cardiac CT slices with a from-scratch U-Net in PyTorch.

Note on the repository name. The repository is called HeartMRISegmentationNetwork, but the dataset it was built against is CT (CT_*.jpg slices with gold_*.png masks), not MRI. The code is modality-agnostic; only the name is inaccurate.

---

What this is

A small, self-contained research codebase: a configurable U-Net, a Dice + BCE objective, a training loop with early stopping on validation F1, and reporting of per-epoch and per-image precision / recall / F1.

It is packaged as heartseg with a three-verb CLI:

Command	Purpose
heartseg train	Train, early-stop on validation F1, score the test split
heartseg evaluate	Score an existing checkpoint on any split
heartseg predict	Render qualitative prediction overlays

---

Installation

Requires Python 3.10+.

python -m venv .venv && source .venv/bin/activate
pip install -e ".[dev]"

Runs on Apple Silicon (mps), CUDA, or CPU — selected automatically in that order, or forced with --device.

---

Dataset layout

The dataset is not distributed with this repository. Point --root at a directory shaped like this:

<root>/
├── images/
│   ├── tr/    CT_11_Image_47.jpg   ...
│   ├── val/   CT_09_Image_31.jpg   ...
│   └── ts/    CT_15_Image_55.jpg   ...
└── golds/     gold_11_Image_47.png ...

An image stem CT_<x> maps to the mask stem gold_<x>; both prefixes are configurable. Masks are binary gold standards and may be encoded as 0/1 or 0/255 — both are handled.

---

Usage

Train with the default architecture (16 initial features, 4 levels):

heartseg train --root /path/to/data --output-dir runs/baseline

Sweep an architecture:

heartseg train --root /path/to/data --init-features 32 --levels 4 --dropout 0.3 --output-dir runs/feat32_do03

Score a trained checkpoint on every split:

heartseg evaluate --root /path/to/data --checkpoint runs/baseline --splits tr val ts

Render overlays for specific slices:

heartseg predict --root /path/to/data --checkpoint runs/baseline --output-dir runs/baseline --image-ids CT_11_Image_47 CT_12_Image_49

Every command accepts --device, --seed, --batch-size and --log-level. The legacy flag spellings from the original scripts (--init_feat, --bs, --lr, --out) are still accepted as aliases.

Training writes into the output directory:

File	Contents
best.pth	Best-validation-F1 checkpoint, with its config embedded
epoch_metrics.csv	epoch,split,precision,recall,f1,loss per epoch, plus a final test row
test_metrics.csv	image,precision,recall,f1 — one row per test image

---

Method

Architecture. A symmetric U-Net. The stem is a double convolution (Conv → BatchNorm → ReLU → Dropout → Conv → BatchNorm → ReLU); each of the levels encoder stages max-pools by 2 and doubles the channel width; each decoder stage transposed-convolves back up, concatenates the matching encoder skip, and halves the width. A 1×1 convolution produces a single logit map. With the defaults (init_features=16, levels=4) the widths are 16-32-64-128-256.

Objective. ComboLoss = Dice + BCEWithLogits. Dice supplies a region-overlap signal that tolerates the heavy foreground/background imbalance of these slices; BCE keeps per-pixel gradients well-behaved early in training when the Dice term is nearly flat.

Preprocessing. Slices are read as RGB, resized to 256×256 and normalised with ImageNet statistics. Training additionally applies horizontal/vertical flips, ±20° rotation and mild brightness/contrast jitter. Geometric transforms run through torchvision.transforms.v2 on an Image/Mask pair, so the image is interpolated smoothly while the mask is resampled nearest-neighbour and stays strictly binary.

Optimisation. AdamW (lr=1e-3, weight_decay=1e-4) with ReduceLROnPlateau on validation F1, and early stopping after --patience epochs without improvement.

Metrics. Precision, recall and F1 are computed per image at a 0.5 probability threshold and then macro-averaged, so every slice counts equally regardless of how much heart it contains. An image with no predicted and no ground-truth foreground scores 1.0.

---

Results

Test-split scores of the committed runs, macro-averaged over 21 slices:

Run	Init features	Levels	Dropout	Epochs	Precision	Recall	F1
runs_mps16__do00	16	4	0.0	60	0.850	0.953	0.888
runs_feat32	32	4	—	69	0.851	0.948	0.884
runs_lvl3	16	3	—	43	0.838	0.960	0.879
runs_mps16	16	4	—	53	0.834	0.956	0.876
runs_mps16__do05	16	4	0.5	84	0.828	0.960	0.875
runs_mps16__do03	16	4	0.3	42	0.790	0.956	0.850

Recall is consistently ~12 points above precision: the model over-segments, spilling past the myocardial border rather than missing it. Dropout did not help at this dataset size — the plain 16-feature model is the best of the six, and 0.3 dropout is clearly the worst. Widening to 32 features and shrinking to 3 levels both land within noise of the baseline, which suggests capacity was not the binding constraint.

⚠️ These numbers are not reproducible from the current code

They were produced by the pre-1.0 pipeline, which fed the loss unnormalised 0–255 masks and bilinearly interpolated the mask during resize and rotation. Both bugs are fixed (see CHANGELOG), and both affected the reported score — soft targets in [0, 1] make the Dice term and the thresholded overlap easier to satisfy than crisp binary labels. The table is kept as a historical record.

The checkpoints still load, so heartseg evaluate will score them, but against the corrected input pipeline — expect different numbers.

The Dropout column reads — for three runs whose recorded architecture contains a dropout slot but whose probability was not persisted; Dropout2d has no parameters, so the value cannot be recovered from the checkpoint. This is exactly what the 1.0 checkpoint format fixes.

Trained weights are no longer tracked by git (they remain in history). The metric CSVs and overlays under runs_*/ are.

---

Development

pytest                 # test suite; synthesises its own data, no dataset needed
ruff check . && ruff format --check .
mypy

Tests that need the committed checkpoints are marked checkpoints and skip cleanly when the .pth files are absent; the end-to-end training test is marked slow.

src/heartseg/
├── cli.py              argparse front end for all three subcommands
├── config.py           frozen, validating configuration dataclasses
├── training.py         the training loop
├── evaluation.py       checkpoint scoring
├── engine.py           train_one_epoch / evaluate / predict_per_image
├── losses.py           DiceLoss, ComboLoss
├── metrics.py          confusion counts, precision/recall/F1
├── checkpoint.py       save/load + legacy state_dict translation
├── reporting.py        CSV writers
├── visualization.py    overlay rendering
├── device.py           device selection
├── reproducibility.py  seeding
├── data/               dataset and v2 transform pipeline
└── models/             the U-Net

License

MIT — see LICENSE.