Training Your
Own Models
Browse the book
Chapter 10 / 4010 min read

10 Project predict a mask then extract an outline

Now the input is still an image, but “what kind of image is this?” is no longer enough. You want the foreground's shape. The output must retain spatial information.

The synthetic exercise draws irregular foreground regions, occasional holes, varying illumination, and noise. It deliberately supplies a learnable color cue. Open a fresh terminal at the companion root, run cd examples/tiny-ml, and activate the neural environment you checked in the image-classification project. If you skipped that project, complete its explicit installation block first. Then run:

python make_shapes.py --out data/shapes --size 128 \
  --train 256 --val 64 --test 64 --seed 42
python mask_metrics.py
python train_segmenter.py --data data/shapes --out runs/segmenter-smoke \
  --epochs 1 --batch-size 2 --width 16 --device cpu
# After checking save and reload, use a new longer run:
python train_segmenter.py --data data/shapes --out runs/segmenter \
  --epochs 15 --batch-size 8 --width 16 --device cpu

For an installed compatible GPU environment, replace the last command's device with --device cuda --amp. Begin with a one-epoch smoke run and inspect the saved images before spending time on a longer run. The generator and mask-metric unit checks were executed on CPU; the neural training command was not executed in the authoring environment.

Classification detection semantic segmentation and instances

These outputs are not interchangeable:

  • Classification: one label or set of labels for the whole image
  • Object detection: an object class and a bounding box for each detected object; it localizes objects approximately
  • Semantic segmentation: a class at each pixel; two touching objects of the same class may form one region
  • Instance segmentation: a separate mask and identity for each object, even when neighboring objects share a class
  • Contour extraction: convert an existing region boundary into a curve or polygon; this is often postprocessing of a mask

A box says where an object lies, not which pixels belong to it. It cannot recover an exact silhouette by itself. Semantic segmentation also cannot guarantee separate identities for touching foreground objects. If you need to count overlapping identical parts, label instances and choose an instance-capable model or a justified separation procedure.

The exercise is binary semantic segmentation: target values are background 0 and foreground 1. Files encode them as 0 and 255, but the dataset loader converts them back to floating-point 0 and 1. Each batch's target shape is [batch, 1, height, width]. The one-channel model output has the same shape, containing logits rather than binary masks.

A tiny encoder decoder with skip connections

TinyUNet downsamples twice to learn broader context, then upsamples twice to return to the original resolution. At each upsampling stage, a skip connection concatenates higher-resolution features from the encoder. These features help retain location information lost during downsampling. This is a deliberately smaller U-Net-like teaching network, not a reimplementation of the original paper's exact architecture. M07

At width 16, the channel progression is 16, 32, and 64. --width scales that capacity. Doubling width roughly quadruples many convolution weight and compute costs because both input and output channel counts grow. Increasing image width and height together increases pixel count quadratically. Training memory also stores intermediate activations and gradients; a small parameter count does not imply small high-resolution activation memory.

Bilinear upsampling uses align_corners=False, and its target size comes from the corresponding skip tensor. This makes tensor sizes match rather than relying on a fragile assumption about division rounding. The final 1 × 1 convolution maps features to one foreground logit per pixel.

What a segmentation architecture needs from its labels

The scratch network needs aligned pixel labels, not only image-level class names or bounding boxes. A box-supervised task requires a different method; pretending every pixel inside a box is foreground teaches a different target. A pretrained encoder can reduce the amount of visual feature learning, but it does not supply your target's boundary convention. Fine structures, touching objects, holes, and difficult backgrounds must appear in labeled examples if they matter in deployment.

More layers may enlarge the receptive field but cannot recover omitted labels or distinguish instances when all instance identities were collapsed into one semantic mask. Higher resolution is often a more appropriate experiment than a much deeper network when the failure is a narrow edge. Estimate data needs through learning curves and labeled error slices, with independent scenes as the counting unit.

Choose a pixel loss that does not reward empty predictions

The example adds two terms:

loss = binary_cross_entropy_with_logits(logits, target)
       + mean_per_image_soft_dice_loss

Binary cross entropy penalizes incorrect pixel-level probability assignments. Its logits variant combines the sigmoid and logarithmic calculation stably. Soft Dice uses probabilities to measure foreground overlap before thresholding. In the script:

Dice = (2 × sum(probability × target) + 1)
       / (sum(probability) + sum(target) + 1)
Dice loss = 1 − Dice

The sums are per image, then averaged. The smoothing constant 1 keeps empty and tiny regions numerically manageable. It influences the loss more strongly for tiny objects; it is not a physical constant. Equal weighting of BCE and Dice is a starting choice. If foreground is extremely rare, compare suitable class weighting or focal-style objectives, but judge the result on actual foreground and boundary errors rather than the loss value alone. M05

An all-background prediction can have excellent pixel accuracy when only 1% of pixels are foreground. Always include foreground overlap, empty-image behavior, and small-object errors. Include genuinely empty images in training and evaluation if the deployed camera will encounter them.

For K mutually exclusive semantic classes, use K logits per pixel and integer target maps with values from zero to K−1. For overlapping labels, use independent binary channels instead. If some pixels are unlabeled, mask them out of the loss; do not silently relabel them as background.

Keep image transformations and targets aligned

The dataset applies the same random horizontal flip and 90-degree rotation to image and mask. A transform applied to only one of them creates supervision for the wrong pixels. If you resize a categorical mask, nearest-neighbor interpolation preserves class IDs; bilinear interpolation invents fractional labels unless that soft-label behavior is explicitly intended. For an image, continuous interpolation is usually appropriate.

Real segmentation datasets require a labeling convention: does an outline include anti-aliased edge pixels, shadows, transparent regions, holes, and occluded sections? Are you labeling visible extent or inferred complete shape? Two annotators may disagree even when both are careful. Measure disagreement on a small independently labeled subset before demanding one-pixel model accuracy.

Split by original scene, video, subject, or acquisition session before making patches. Crops from the same original image should not leak across train and test. A giant image cut into thousands of adjacent tiles is not thousands of independent scenes.

High IoU is not exact outline accuracy

Intersection over union, or IoU, is overlapping foreground area divided by foreground area covered by either prediction or target. It is useful, but it is area-based. A one-pixel displacement of a 96 × 96 square gives IoU about 0.979. In an executed metric check, its exact-pixel boundary F1 was only 0.50. An application that needs a precise cutting outline should not call that “98% correct boundary.”

The supplied mask_metrics.py reports both IoU and a symmetric boundary F1. It extracts boundary pixels, measures how many predicted boundary pixels fall within a tolerance of a true boundary, measures the reverse fraction, and combines the two. Tolerance 2 means a two-pixel neighborhood under a square, Chebyshev-distance structuring element. This is not Euclidean distance and is not the paper's separately defined Boundary IoU metric. Both-empty masks score one; exactly one empty mask scores zero. Report that convention because datasets with many empty images can inflate mean scores. The Boundary IoU paper explains why boundary-sensitive evaluation complements regional overlap. M08

For a real outline product, report error in the coordinate system that matters: original-image pixels, millimeters after calibration, or another physical unit. Add boundary distance summaries, small-object stratification, and checks for disconnected fragments and holes when relevant. A large region can dominate an aggregate metric while a narrow critical feature is missing.

From probability map to polygon

After training:

python predict_segmenter.py runs/segmenter/best.pt my_image.png \
  --out prediction --threshold 0.5
python contours.py prediction/mask.png --epsilon-px 1 \
  --out prediction/contours.json

Add --device cuda to predict_segmenter.py for GPU inference; the default is CPU. Both load the same state dictionary and saved width, convert the image to RGB float32 divided by 255, run evaluation/inference mode, sigmoid the output, and apply the requested threshold. The CUDA path moves inputs and weights to CUDA and returns the mask to CPU for PNG writing. Both inference paths remain FP32; the training --amp setting does not force inference precision.

The first script writes a probability visualization and binary mask at original input resolution. It does not resize the image. Very large images can therefore need substantial RAM; use a validated tiled strategy if required. The contour script requires OpenCV and was syntax-checked only; it is an optional downstream recipe, not a dependency of training.

Contour extraction uses RETR_TREE so the hierarchy can preserve holes, and CHAIN_APPROX_NONE before optional polygon simplification. Each output contour stores x/y points plus parent/child hierarchy. epsilon-px=1 is the polygon approximation tolerance in mask pixels. It controls geometric simplification, not the neural model's uncertainty. Raising it removes vertices and can erase important details. RETR_EXTERNAL would discard interior holes. M09

A practical pipeline is probability map → threshold → optional validated cleanup → contour extraction → optional simplification → coordinate mapping. Threshold, minimum component size, hole filling, and smoothing are model-selection choices: tune them on validation, then freeze them before final testing. Blind hole filling is wrong for objects that genuinely have holes. Keeping only the largest component is wrong when multiple objects are legitimate.

If you resized, padded, or cropped an image, undo that geometry when returning the contour to the caller. A perfect mask in the wrong coordinate system is still a wrong result. A segmentation network also cannot restore subpixel detail that was destroyed by a low-resolution input.

Understand the segmentation controls and artifacts

Control Default Purpose
generator --size 128 Image side length; at least 32, divisible by four
generator --train/--val/--test 256/64/64 Independently generated examples in each split
--seed 42 Reproducible generator or training random state
training --epochs 15 Maximum passes through training images
--batch-size 8 Images contributing to each update; lower first if out of memory
--width 16 Base number of feature channels
--lr 0.001 AdamW learning rate
--weight-decay 0.0001 Decoupled weight regularization
--device cpu CPU or CUDA execution
--amp off CUDA FP16 autocast plus gradient scaling; supported only with CUDA here
inference --threshold 0.5 Probability cutoff for foreground
contour --epsilon-px 1 Polygon simplification tolerance

The loop zeroes gradients, runs the forward pass, computes the loss, backpropagates, clips the gradient norm to 1, and updates with AdamW. In AMP mode, it unscales gradients before clipping. PyTorch's current namespace for the pinned API is torch.amp.GradScaler; the example avoids older deprecated torch.cuda.amp calls. M12

best.pt holds the lowest-validation-loss state dictionary and network width. metrics.json contains configuration, parameter count, training history, test metrics, elapsed time, and peak allocated CUDA memory if applicable. Eight input/probability/prediction/target groups support visual inspection. Allocated CUDA memory is not the complete device footprint; compare it with system-level observations before claiming a fit.

Training examples are procedural and much simpler than real camera data. Their purpose is to expose tensor shapes, paired transforms, output decoding, and evaluation. A strong synthetic score does not establish robustness to new lenses, reflections, clutter, textures, occlusion, or unknown classes.

When a detector or pretrained segmenter is the right next step

If your requested output is a class and rectangle for each object, use a detector. If it is a separate mask for each object, use an instance segmenter. The official torchvision detection tutorial demonstrates the dataset contract and fine-tuning a Mask R-CNN model; its example dataset is educational, not an automatic commercial data license. M10

A custom detection sample returns an image and a target dictionary containing boxes in [x_min, y_min, x_max, y_max] format, labels, and, for instance segmentation, per-instance masks. Preserve the model's conventions: torchvision detectors reserve label zero for background. Empty scenes need valid empty tensors, not invented background boxes. Detection training usually returns a dictionary of losses; sum the intended components and inspect each rather than assuming it behaves like image classification.

fasterrcnn_mobilenet_v3_large_320_fpn is an official compact-backbone option for a box-detection experiment. The API is versioned and the detection module's stability warnings should be read before production use. Start with a very small batch and representative resolutions on a 24 GB GPU; proposal counts and image content affect memory. Review pretrained-weight terms independently of the permissive torchvision source-code license. M25 M11

Detection average precision depends on class, confidence ranking, and matching at specified IoU thresholds. AP at IoU 0.5 is not equivalent to COCO-style AP averaged over multiple thresholds, and neither proves exact edges. For segmentation transfer learning, a small pretrained encoder can be useful; retaining its normalization and adapting its output head is essential. Compare a frozen encoder, partial fine-tuning, and the scratch baseline before increasing model size.