Singapore team selection 2025 · Correction by IOAI Training

Semantic Segmentation with U-Net

This page gives the full approach. Read it once you have searched.

Metric: mIoU (5 object classes)
Evaluated on: official test
Our score: 0.609 mIoU on VOC val (0.490 batch-avg)
Expected level: about 35 % mIoU for ResNet-34 U-Net, 5 epochs

Expected level reached

We solved the problem ourselves: downloaded the data, ran the experiments and measured the score below. The statement and data belong to their organisers.

The problem in brief

You must segment PASCAL VOC 2012 images at the pixel level into six classes: background plus person, car, dog, bicycle and chair (all other VOC objects count as background). Only images containing at least one of the five classes are kept: 726 training images (from VOC train) and 719 evaluation images (from VOC val), resized to $256\times256$ . The notebook walks through a progression: data augmentation (Q1.1), a U-Net (Q1.2), a class-weighted loss (Q1.3), a U-Net whose encoder is a pretrained ResNet-34 (Q1.4, plus a bonus for higher mIoU), a residual U-Net and the focal loss (Q2), then an analysis of the class distribution, hyperparameter tuning, other pretrained backbones and other architectures (Q3).

How you are scored

100 points, mostly for correct code and answers, with the bonus of Q1.4 "based on the final mIoU". The metric is the mean Intersection over Union over the five object classes (background excluded, pixels labelled 255 ignored):

\text{IoU}_c = \frac{TP_c}{TP_c + FP_c + FN_c}, \qquad \text{mIoU} = \frac{1}{5}\sum_{c=1}^{5}\text{IoU}_c .

The notebook states that the ResNet-34 U-Net trained for 5 epochs should reach "around 35 % mIoU": that is our reference.

A subtlety: the notebook's validate_model computes the IoU per batch of 8 images and averages over batches, skipping classes absent from a batch. We report the standard dataset-level mIoU (one confusion matrix accumulated over all validation pixels), and give the notebook's number for the final model too. They differ: a few false-positive pixels of a class absent from a batch give that batch an IoU of 0 for the class, which drags the batch average down.

First look at the data

Q3 Task 1 asks to look for a long tail, and it is there.

Left: share of training pixels per class on a log scale: background 80.8 %, person 9.9 %, dog 3.5 %, car 2.9 %, chair 2.4 %, bicycle 0.6 %. Right: number of training images containing each class: person 442, chair 148, car 128, dog 121, bicycle 65.

Background is 81 % of the pixels; a model predicting "background" everywhere has 81 % pixel accuracy and 0 mIoU. That is why accuracy is useless here and why the loss needs care.
Person dominates the objects (442 of 726 images); bicycle is the rarest (0.6 % of pixels, 65 images) and also the hardest shape: thin frames and spokes, mostly made of background pixels.
Chairs are often small, partly hidden and seen from any angle.

A first baseline

The notebook's warm-up model (two convolutions, one pooling, upsampling) reaches 0.068 mIoU after 5 epochs, and the U-Net from scratch with weighted cross-entropy 0.062. Neither has learnt to recognise objects: with 726 images, learning "what a dog is" from random weights in 5 epochs is hopeless. The U-Net (31 M parameters) is not worse because of its architecture; it simply has far more to learn from too little data.

The key idea

Segmentation needs two things: knowing what is in the image and knowing where exactly. A U-Net gets the "what" from its deep, low-resolution encoder and the "where" from skip connections; with little data, the "what" must be borrowed from a pretrained encoder.

The code questions.

Q1.1, augmentation. A horizontal flip with probability 0.5 and a random rotation in $[-10°, 10°]$ , with the same random draw applied to the image and to the mask (otherwise labels no longer match pixels). The mask must be rotated with nearest-neighbour interpolation (no in-between class values), and the corners that the rotation brings in are filled with the "ignore" label 255 rather than with background. The notebook re-seeds torch before each of its two Compose pipelines to synchronise them; drawing the parameters once and calling TF.hflip / TF.rotate on both is cleaner.
Q1.2, U-Net. self.inc = DoubleConv(3, 64); each downK is nn.Sequential(nn.MaxPool2d(2), DoubleConv(c, 2c)) up to 1024 channels; the skip connection is torch.cat([x, x4], dim=1) (channel-wise concatenation, which is why conv1 takes 1024 channels); the head is a $1\times1$ convolution nn.Conv2d(64, n_classes, 1) that turns each pixel's 64 features into 6 class scores.
Q1.3, weighted loss. nn.CrossEntropyLoss(weight=class_weights, ignore_index=255), with weights proportional to the inverse pixel frequency (normalised to sum to 6): background gets 0.026, bicycle 3.5. A bicycle pixel then counts 135 times more than a background pixel.
Q1.4, pretrained encoder. self.encoder1 = nn.Sequential(resnet.conv1, resnet.bn1, resnet.relu): the first ResNet stage (a $7\times7$ convolution with stride 2, then BatchNorm and ReLU), at half resolution, whose output is the shallowest skip connection. The pretrained encoder gets learning rate $10^{-4}$ and the new decoder $10^{-3}$ : we want to adjust ImageNet features gently, not destroy them, while the randomly initialised decoder must learn fast.
Q2.1, residual block. out = out + self.skip(residual), where skip is the identity or a $1\times1$ convolution when the channel count changes. The block then learns a correction $F(x)$ to its input, and the gradient flows through the "+" unchanged, which is what makes deep networks trainable.
Q2.2, focal loss. With $p_t$ the predicted probability of the true class,

\text{FL}(p_t) = -\alpha\,(1 - p_t)^{\gamma}\,\log p_t .

$\gamma$ (focusing) multiplies the cross-entropy by $(1-p_t)^\gamma$ : a pixel already classified with $p_t = 0.9$ has its loss divided by 100 for $\gamma = 2$ , so training concentrates on hard, misclassified pixels (object boundaries, rare classes) instead of the millions of easy background pixels. $\alpha$ (balancing) rescales the loss, globally or per class, to compensate for class frequency, the same role as the class weights of Q1.3. Our implementation computes $p_t$ from the unweighted cross-entropy and applies the class weights and $\alpha$ on top.

What actually moves the mIoU (Q1.4 bonus and Q3). Three changes, each with a reason:

Train longer, with a cosine learning-rate schedule. Five epochs on 726 images are only 455 gradient steps; the curve below is still climbing steeply at epoch 5. Twenty epochs with a cosine decay take the ResNet-34 U-Net from 0.243 to 0.388.
Replace the inverse-frequency weights by cross-entropy + soft Dice. The weights are extreme (background $\times 0.026$ ), so the model is pushed to paint uncertain regions with rare classes, which creates false positives and ruins IoU, a metric that punishes false positives as much as misses. The soft Dice loss optimises an overlap measure directly, per class, and is insensitive to how many pixels a class has:

\mathcal L_{\text{Dice}} = 1 - \frac{1}{5}\sum_{c=1}^{5} \frac{2\sum_i p_{i,c}\, y_{i,c} + 1}{\sum_i p_{i,c} + \sum_i y_{i,c} + 1},

where $p_{i,c}$ is the softmax probability of class $c$ at pixel $i$ . Plain (unweighted) cross-entropy keeps the per-pixel calibration. Result: 0.484, and background IoU jumps from 0.74 to 0.92. 3. A stronger pretrained encoder (Q3 Task 3). ResNet-50 (feature maps of 256 to 2048 channels instead of 64 to 512) with the same decoder recipe: 0.609. Better features are the single largest gain: dog 0.71 to 0.79, bicycle 0.0 to 0.35.

The full solution

solution.py contains the data pipeline, all the models (simple CNN, U-Net, ResNet-34 and ResNet-50 U-Nets), the weighted CE, focal and Dice losses, and the evaluation; python solution.py <model> <epochs> trains one configuration. The data come from the Hugging Face mirror nateraw/pascal-voc-2012 (the official 1 464 / 1 449 segmentation split, masks stored as colour PNGs, which we map back to class indices with the standard VOC colour map). Each image is resized to $256\times256$ once at load time; all models use batches of 8 and Adam; one epoch takes about 20 seconds on a laptop GPU (Apple MPS).

loss = F.cross_entropy(logits, target, ignore_index=255) + dice_loss(logits, target)
opt = Adam([{"params": encoder, "lr": 1e-4}, {"params": decoder, "lr": 1e-3}])
sched = CosineAnnealingLR(opt, T_max=20)

Results

Dataset-level mIoU over the 5 object classes on the 719 validation images (VOC val, whose labels are public; VOC has no public test labels). All numbers are those of the last epoch, not the best one, so that the validation set is not used to pick a checkpoint.

model	loss	epochs	mIoU	person	car	dog	bicycle	chair
simple CNN (warm-up)	weighted CE	5	0.068	0.16	0.06	0.08	0.02	0.01
U-Net from scratch	weighted CE	5	0.062	0.14	0.08	0.06	0.02	0.01
ResNet-34 U-Net (notebook setting)	weighted CE	5	0.243	0.38	0.31	0.42	0.03	0.07
ResNet-34 U-Net	weighted CE, cosine	20	0.388	0.58	0.53	0.47	0.21	0.15
ResNet-34 U-Net	CE + Dice, cosine	20	0.484	0.76	0.75	0.71	0.00	0.20
ResNet-50 U-Net	CE + Dice, cosine	20	0.609	0.81	0.81	0.79	0.35	0.28

For the final model, the notebook's batch-averaged metric gives 0.490 (against 0.609 dataset-level): whichever way it is measured, it is well above the expected 35 %. Our 5-epoch ResNet-34 run (0.243) is below the notebook's "around 35 %"; differences in metric, augmentation and data loading probably explain part of the gap, and the curve shows that 5 epochs is simply too early to judge.

Validation mIoU per epoch for the six configurations: the two models without pretraining stay below 0.08; the ResNet-34 U-Net climbs to 0.39 with weighted CE and to 0.48 with CE + Dice; the ResNet-50 U-Net reaches 0.61.

Three validation images with ground truth and the ResNet-50 U-Net prediction: the riders and bicycles are well segmented (the prediction even follows the bicycle frames, while the VOC annotation fills the whole wheel area), the small chair in the living room is missed, and the bicycle by the wall is found but as a blob.

What remains hard: chairs (0.28) and bicycles (0.35). Both are thin or occluded, with few training examples, and VOC annotates bicycles as filled regions while the model learns to follow the frame. The ResUNet of Q2.1 and the focal loss of Q2.2 are implemented in the code, and Q3 Task 4 suggests comparing with Transformer-based segmenters (SegFormer, Mask2Former); we did not train those three variants within our compute budget. Given the table, a pretrained Transformer backbone is the most promising next step; a residual U-Net trained from scratch would face the same lack of data as the plain U-Net.

What to remember

Look at the class distribution first: 81 % background means pixel accuracy is meaningless, and rare thin classes (bicycle) will dominate the mIoU's weak spots.
Augment image and mask with the same random transform, nearest-neighbour for masks, ignore label for padded corners.
Pretrained encoders are what make segmentation possible on a few hundred images: from scratch, 0.06; with ImageNet ResNet-34, 0.39; with ResNet-50, 0.61. Use a smaller learning rate for the pretrained part.
Pick a loss that matches the metric: extreme inverse-frequency weights create false positives that IoU punishes; CE + soft Dice optimises overlap directly.
Check how the metric is aggregated (per batch or over the whole dataset): the same model scores 0.49 or 0.61.

Full code (solution.py)

"""Singapore IOAI 2025 selection test, CV: semantic segmentation of 5 PASCAL VOC classes with U-Net.

Classes: background + person, car, dog, bicycle, chair. Images resized to 256 x 256.
Train = VOC2012 'train' images containing at least one of the 5 classes, evaluation = same filter on 'val'.
Metric: mIoU over the 5 object classes (background excluded), as in the notebook.
Data: HF nateraw/pascal-voc-2012 (two parquet files, the official 1464/1449 segmentation split,
masks stored as VOC colour PNGs) saved in work/.

  python solution.py simple | unet | resunet34 | resunet34_long | resunet34_dice | resunet50 [epochs]
"""
import io, json, sys, time
import numpy as np, pandas as pd, torch, torch.nn as nn, torch.nn.functional as F
from PIL import Image
from torch.utils.data import Dataset, DataLoader
import torchvision.transforms as T
import torchvision.transforms.functional as TF
from torchvision.models import resnet34, ResNet34_Weights, resnet50, ResNet50_Weights

torch.set_num_threads(2); torch.manual_seed(0); np.random.seed(0)
dev = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
SELECTED = {0: 0, 15: 1, 7: 2, 12: 3, 2: 4, 9: 5}           # VOC id -> our id (person car dog bicycle chair)
NAMES = ["background", "person", "car", "dog", "bicycle", "chair"]

def voc_palette():
    """The standard VOC colour map: class k -> RGB built from the bits of k."""
    pal = np.zeros((256, 3), dtype=np.uint8)
    for k in range(256):
        c, r, g, b = k, 0, 0, 0
        for j in range(8):
            r |= ((c >> 0) & 1) << (7 - j); g |= ((c >> 1) & 1) << (7 - j); b |= ((c >> 2) & 1) << (7 - j); c >>= 3
        pal[k] = (r, g, b)
    return pal
PAL = voc_palette()
COLOR2ID = {tuple(int(v) for v in PAL[k]): k for k in list(range(21)) + [255]}

def rgb_to_ids(m):
    a = np.asarray(m.convert("RGB")); key = (a[..., 0].astype(np.int32) << 16) | (a[..., 1].astype(np.int32) << 8) | a[..., 2]
    out = np.full(key.shape, 255, np.uint8)                   # unknown colours -> ignore
    for (r, g, b), k in COLOR2ID.items(): out[key == ((r << 16) | (g << 8) | b)] = k
    return out

def load(split):
    df = pd.read_parquet(f"work/{split}.parquet"); items = []
    for _, r in df.iterrows():
        ids = rgb_to_ids(Image.open(io.BytesIO(r["mask"]["bytes"])))
        if not np.isin(ids, [15, 7, 12, 2, 9]).any(): continue          # the notebook's filter
        lut = np.zeros(256, np.uint8); lut[255] = 255                   # other VOC classes -> background
        for v, s in SELECTED.items(): lut[v] = s
        img = Image.open(io.BytesIO(r["image"]["bytes"])).convert("RGB").resize((256, 256), Image.BILINEAR)
        items.append((img, Image.fromarray(lut[ids]).resize((256, 256), Image.NEAREST)))   # resize once: low RAM
    del df; return items

class SegData(Dataset):
    def __init__(s, items, augment): s.items, s.augment = items, augment
    def __len__(s): return len(s.items)
    def __getitem__(s, i):
        img, msk = s.items[i]
        if s.augment:                                    # Q1.1: SAME random flip and rotation for image and mask
            if np.random.rand() < 0.5: img, msk = TF.hflip(img), TF.hflip(msk)
            ang = float(np.random.uniform(-10, 10))
            img = TF.rotate(img, ang); msk = TF.rotate(msk, ang, fill=255)   # rotated-in corners are ignored
        x = TF.normalize(TF.to_tensor(img), [0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        return x, torch.from_numpy(np.array(msk)).long()

# ------------------------------------------------------------------ models
class Simple(nn.Module):                         # the notebook's warm-up model
    def __init__(s, n=6):
        super().__init__()
        s.net = nn.Sequential(nn.Conv2d(3, 64, 3, padding=1), nn.ReLU(True), nn.Conv2d(64, 64, 3, padding=1), nn.ReLU(True),
                              nn.MaxPool2d(2), nn.Conv2d(64, 64, 3, padding=1), nn.ReLU(True), nn.Conv2d(64, n, 3, padding=1),
                              nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True))
    def forward(s, x): return s.net(x)

def DoubleConv(i, o):
    return nn.Sequential(nn.Conv2d(i, o, 3, padding=1), nn.BatchNorm2d(o), nn.ReLU(True),
                         nn.Conv2d(o, o, 3, padding=1), nn.BatchNorm2d(o), nn.ReLU(True))

class UNet(nn.Module):                           # Q1.2
    def __init__(s, n=6):
        super().__init__()
        s.inc = DoubleConv(3, 64)                                             # line 1
        s.downs = nn.ModuleList([nn.Sequential(nn.MaxPool2d(2), DoubleConv(a, b))
                                 for a, b in [(64, 128), (128, 256), (256, 512), (512, 1024)]])
        s.ups = nn.ModuleList([nn.ConvTranspose2d(a, a // 2, 2, stride=2) for a in (1024, 512, 256, 128)])
        s.convs = nn.ModuleList([DoubleConv(a, a // 2) for a in (1024, 512, 256, 128)])
        s.outc = nn.Conv2d(64, n, 1)                                          # line 3: 1x1 conv to classes
    def forward(s, x):
        skips = [s.inc(x)]
        for d in s.downs: skips.append(d(skips[-1]))
        x = skips.pop()
        for up, conv in zip(s.ups, s.convs):
            x = conv(torch.cat([up(x), skips.pop()], dim=1))                  # line 2: skip = concatenation
        return s.outc(x)

class ResNetUNet(nn.Module):                     # Q1.4 / Q3.3: pretrained ResNet encoder
    def __init__(s, n=6, depth=34):
        super().__init__()
        r = resnet34(weights=ResNet34_Weights.DEFAULT) if depth == 34 else resnet50(weights=ResNet50_Weights.DEFAULT)
        c = [64, 64, 128, 256, 512] if depth == 34 else [64, 256, 512, 1024, 2048]
        s.enc1 = nn.Sequential(r.conv1, r.bn1, r.relu)                         # the 1 line of Q1.4
        s.pool, s.enc2, s.enc3, s.enc4, s.enc5 = r.maxpool, r.layer1, r.layer2, r.layer3, r.layer4
        s.up5, s.dec5 = nn.ConvTranspose2d(c[4], 256, 2, 2), DoubleConv(256 + c[3], 256)
        s.up4, s.dec4 = nn.ConvTranspose2d(256, 128, 2, 2), DoubleConv(128 + c[2], 128)
        s.up3, s.dec3 = nn.ConvTranspose2d(128, 64, 2, 2), DoubleConv(64 + c[1], 64)
        s.up2, s.dec2 = nn.ConvTranspose2d(64, 32, 2, 2), DoubleConv(32 + c[0], 32)
        s.up1, s.dec1 = nn.ConvTranspose2d(32, 16, 2, 2), DoubleConv(16, 16)
        s.last = nn.Conv2d(16, n, 1)
        s.encoder = nn.ModuleList([s.enc1, s.enc2, s.enc3, s.enc4, s.enc5])
    def forward(s, x):
        e1 = s.enc1(x); e2 = s.enc2(s.pool(e1)); e3 = s.enc3(e2); e4 = s.enc4(e3); e5 = s.enc5(e4)
        d = s.dec5(torch.cat([s.up5(e5), e4], 1)); d = s.dec4(torch.cat([s.up4(d), e3], 1))
        d = s.dec3(torch.cat([s.up3(d), e2], 1)); d = s.dec2(torch.cat([s.up2(d), e1], 1))
        return s.last(s.dec1(s.up1(d)))

# ------------------------------------------------------------------ losses and metric
class FocalLoss(nn.Module):                      # Q2.2: FL = alpha * (1 - p_t)^gamma * CE
    def __init__(s, alpha=1.0, gamma=2.0, weight=None, ignore_index=255):
        super().__init__(); s.alpha, s.gamma = alpha, gamma
        s.ce = nn.CrossEntropyLoss(weight=weight, ignore_index=ignore_index, reduction="none")
    def forward(s, logits, t):
        ce = s.ce(logits, t)                                         # (class-weighted) CE per pixel
        pt = torch.exp(-F.cross_entropy(logits, t, ignore_index=255, reduction="none"))   # prob. of the true class
        valid = t != 255
        return (s.alpha * (1 - pt) ** s.gamma * ce)[valid].mean()

def dice_loss(logits, t, n=6):
    valid = (t != 255).unsqueeze(1); p = logits.softmax(1) * valid
    oh = F.one_hot(t.clamp(max=n - 1), n).permute(0, 3, 1, 2) * valid
    inter = (p * oh).sum((0, 2, 3)); den = p.sum((0, 2, 3)) + oh.sum((0, 2, 3))
    return 1 - ((2 * inter + 1) / (den + 1))[1:].mean()          # soft Dice on the 5 object classes

def class_weights(items):                       # Q1.3: inverse frequency, normalised to sum to C
    cnt = np.zeros(6)
    for _, m in items:
        a = np.array(m); cnt += np.bincount(a[a != 255], minlength=6)[:6]
    w = 1 / (cnt + 1e-10); return torch.tensor(w / w.sum() * 6, dtype=torch.float32), cnt

@torch.no_grad()
def evaluate(model, loader):
    """Dataset-level IoU per class (confusion matrix accumulated over all val pixels)."""
    model.eval(); cm = torch.zeros(6, 6, dtype=torch.long)
    for x, t in loader:
        p = model(x.to(dev)).argmax(1).cpu(); v = t != 255
        cm += torch.bincount(t[v] * 6 + p[v], minlength=36).view(6, 6)
    model.train()
    tp = cm.diag().float(); iou = tp / (cm.sum(0) + cm.sum(1) - tp).float()
    return float(iou[1:].mean()), iou.tolist()

def run(name, epochs):
    tr_items, va_items = load("train"), load("val")
    print(f"train {len(tr_items)} images, val {len(va_items)} images", flush=True)
    w, cnt = class_weights(tr_items); print("pixel share", (cnt / cnt.sum()).round(4).tolist(), "weights", w.numpy().round(3).tolist())
    tl = DataLoader(SegData(tr_items, augment=name != "simple"), batch_size=8, shuffle=True, num_workers=0)
    vl = DataLoader(SegData(va_items, augment=False), batch_size=8)
    if name == "simple": model = Simple()
    elif name == "unet": model = UNet()
    else: model = ResNetUNet(depth=50 if name == "resunet50" else 34)
    model.to(dev)
    wce = nn.CrossEntropyLoss(weight=w.to(dev), ignore_index=255)
    if name in ("simple", "unet"):
        opt = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-4)
    else:                                        # small lr for the pretrained encoder, larger for the new decoder
        enc = list(model.encoder.parameters()); ids = {id(p) for p in enc}
        opt = torch.optim.Adam([{"params": enc, "lr": 1e-4}, {"params": [p for p in model.parameters() if id(p) not in ids], "lr": 1e-3}])
    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, epochs) if name.endswith(("long", "dice", "50")) else None
    if name.endswith(("long", "dice", "50")):
        lossf = (lambda o, t: nn.functional.cross_entropy(o, t, ignore_index=255) + dice_loss(o, t)) if name != "resunet34_long" else wce
    else:
        lossf = wce
    hist = []
    for ep in range(1, epochs + 1):
        t0 = time.time(); tot = 0
        for x, t in tl:
            x, t = x.to(dev), t.to(dev)
            opt.zero_grad(); loss = lossf(model(x), t); loss.backward(); opt.step(); tot += loss.item()
        if sched: sched.step()
        miou, ious = evaluate(model, vl)
        hist.append({"epoch": ep, "loss": tot / len(tl), "miou": miou, "ious": ious, "sec": round(time.time() - t0)})
        print(name, hist[-1], flush=True)
    json.dump(hist, open(f"work/res_{name}.json", "w"))
    torch.save(model.state_dict(), f"work/{name}.pth")

if __name__ == "__main__":
    run(sys.argv[1], int(sys.argv[2]) if len(sys.argv) > 2 else 5)