A small, self configuring segmentation setup in the spirit of nnU-Net, built for 2D data. The idea behind nnU-Net is that you should not hand tune an architecture for every new dataset. Instead you measure a few properties of the data and let a set of heuristics derive the patch size, the network depth, the channel schedule, and the normalization. This repository reproduces that idea at a readable scale and backs it with unit tests.
The pipeline has three pieces that compose cleanly.
First it fingerprints a dataset. Given a list of images and their masks, compute_dataset_properties records the median spatial shape, the number of channels, the number of label classes, and a robust summary of the foreground intensity distribution. The median is used rather than the mean so that a handful of unusually large or small cases cannot dominate the plan.
Second it plans. plan_experiment takes the fingerprint and returns a Plan. The plan chooses a patch size whose axes are multiples of sixteen and never larger than the data or a global cap, then trims that patch to respect a rough memory budget. From the patch it derives the network depth, capped so the feature map at the bottleneck never collapses below a minimum size. It builds a doubling channel schedule with a ceiling, and it picks a normalization scheme: plain z scoring for well behaved intensities, or a clip then z score when the distribution has long tails the way CT data does.
Third it builds and trains. ConfigurableUNet reads the plan and assembles an encoder and a mirrored decoder with skip connections, where the depth and width come entirely from the plan rather than from constants in the model. A soft Dice loss and a combined Dice plus cross entropy loss are provided, matching the usual nnU-Net objective.
src/
dataset_properties.py dataset fingerprinting
planner.py the self configuring heuristics
unet.py configurable 2D U-Net
losses.py soft Dice and Dice plus cross entropy
synthetic.py synthetic disk on noise segmentation data
train.py a few step training loop and normalization
tests/
test_dataset_properties.py
test_planner.py
test_unet_and_training.py
The heuristics are deterministic, so the same fingerprint always yields the same plan. That property is asserted in the tests.
- Patch size. Each axis is the median size, snapped down to a multiple of sixteen, clamped to the global maximum, with a floor of sixteen. The patch area is then halved along its longer axis until it fits the voxel budget.
- Depth. The number of stages is set by how many times the patch can be halved before any axis would drop under the minimum feature map size, plus one for the input stage, capped at the global maximum. Larger patches therefore yield deeper networks.
- Width. Channels start at thirty two and double per stage up to a ceiling of three hundred twenty.
- Normalization. When the gap between the low and high intensity percentiles is much wider than the standard deviation, the data is treated as heavy tailed and gets the clip then z score scheme. Otherwise it gets plain z scoring.
synthetic.py draws a bright filled disk at a random location on a noisy background and returns the image together with its binary mask. The foreground sits clearly above the background, so a small planned U-Net learns the segmentation within a few dozen optimizer steps. This keeps the training test fast and free of any download.
The tests are behavior checks rather than snapshots of magic numbers. They confirm that the fingerprint counts cases and classes correctly and that its median is robust to an outlier case. They confirm that the planner produces valid settings: patch axes divisible by sixteen, depth that grows with patch size and stays capped, a bottleneck that never collapses, a doubling and capped channel schedule, and the right normalization branch for tight versus heavy tailed intensities. They confirm that the built network produces an output of the right shape for both clean and awkward input sizes, that its depth matches the plan, and that a single optimizer step actually moves parameters. The headline behavior test trains the planned U-Net on synthetic masks and asserts that the soft Dice loss falls to less than half its starting value and well below the trivial value of one.
Run them with the project virtual environment:
cd C:/Users/sharv/projects/nnunet-experiments
C:/Users/sharv/.venvs/cv/Scripts/python.exe -m pytest tests/ -q
On the machine this was developed on the full suite of twenty one tests passed in a couple of seconds on CPU.
from src.synthetic import make_batch
from src.dataset_properties import compute_dataset_properties
from src.planner import plan_experiment
from src.train import train_steps
images, masks = make_batch(num_cases=8, height=32, width=32, seed=0)
props = compute_dataset_properties(list(images), list(masks))
plan = plan_experiment(props)
model, losses = train_steps(plan, images, masks, num_steps=40)
print(plan.patch_size, plan.num_stages, plan.normalization_scheme)
print(losses[0], losses[-1])