The baseline before the fancy model: medical image segmentation

In short: before training a U-Net, write a threshold with a little cleanup and score it with the same metric. It's quick to write, it tells you how hard the problem really is, and on high-contrast images it is sometimes all you need.

This month I put a small segmentation project on GitHub, image_segmentation. It trains a U-Net with a ResNet-34 encoder, runs inference, and scores the result with IoU and Dice. It works, and the predicted masks look right next to the true ones.

Looking back at it, one thing is missing. There is no baseline. I went straight to the deep model, so I can say the U-Net is good but not how much better it is than something simple. This post is the baseline I'd write first next time, and how to tell whether you need anything more.

What the repo does

The setup is standard. A U-Net, with a ResNet-34 encoder, trained on image and mask pairs stored in one pickle file. A train.py that takes batch size, epochs and learning rate as arguments, and an inference.py that loads a checkpoint and draws the overlay. The sample overlay in the README came from this configuration:

Two trained checkpoints ship with the repo, a short run and a longer one, so anyone can run inference without training. All of this is fine. None of it answers the question a reviewer, a clinician or a client would ask first: compared to what?

The short baseline

Many medical segmentation targets are brighter or darker than what surrounds them. A global threshold catches most of that. Two cheap cleanup steps fix its most common mistakes: keep only the largest connected region, and fill the holes inside it.

import numpy as np
from scipy.ndimage import binary_fill_holes
from skimage.filters import threshold_otsu
from skimage.measure import label

def threshold_baseline(img):
    mask = img > threshold_otsu(img)
    regions = label(mask)
    if regions.max() == 0:
        return mask
    largest = np.argmax(np.bincount(regions.ravel())[1:]) + 1
    return binary_fill_holes(regions == largest)

def dice(a, b, eps=1e-7):
    a, b = a.astype(bool), b.astype(bool)
    return (2 * (a & b).sum() + eps) / (a.sum() + b.sum() + eps)

No training, no GPU, no hyperparameters except the choice of cleanup. Score it with the same Dice function you'll use for the U-Net, on the same test images, and you have a floor.

How each one goes wrong

The two methods fail in different ways, and seeing that is more useful than the score alone. Here is the shape of it, drawn on a toy example: one oval target, and a second bright structure nearby that isn't part of it.

I rasterised those three toy masks and scored them against the truth. These numbers describe the drawings above, not a real dataset, but the pattern is the one you see in practice.

Look at the middle row. One line of cleanup closed more than half the gap between the raw threshold and the U-Net. That is the kind of thing a baseline shows you. If your real numbers look like this, the honest question is whether the last few points of Dice are worth a training pipeline, a GPU and a model to maintain.

When the baseline is enough

Back in 2020 I wrote a simple MATLAB segmentation for brain MRI, separating the brain from the skull and looking for multiple sclerosis lesions. It is still on my GitHub as segment_brain_from_skull. No deep learning, and for that job it did what was needed.

A threshold baseline tends to be enough when:

• The target has clear contrast with its surroundings.
• The images come from one scanner or one protocol, so intensities are consistent.
• What you need downstream is an area or a rough location, not a precise border.
• You have very few labelled images, too few to train anything without overfitting.

It stops being enough when boundaries are faint, when structures touch each other, or when intensities shift between machines. That is where a U-Net earns its cost, and where you'll be glad you have the baseline number to show it.

Before your next segmentation model

Copy the two functions above into your project. Run the baseline on your test set and write the Dice next to your model's in the README. If the gap is small, ship the simple thing. If it's big, you now have a sentence that justifies the model.