View source Download .ipynb

Train and Diagnose SegFormer on Sugar Beets 2016

In this tutorial, we fine-tune a SegFormer semantic-segmentation model on the University of Bonn Sugar Beets 2016 agricultural dataset. This is a 3-class problem — soil, sugarbeet (crop), and weed.

The dataset is prepared as single-channel class-indexed masks with an 80/20 train/val split.

This tutorial covers the following:

1. Semantic segmentation 3LC project setup

2. Semantic segmentation 3LC table creation

3. Training a SegFormer model and collecting per-sample 3LC metrics

4. Collecting diagnostic metrics with a pass over the best model

5. Data curation techniques in semantic segmentation

Part 1: Training and Diagnostics

The first part contains the executable code that creates the Sugar Beets 2016 3LC project, trains SegFormer, and collects per-sample metrics.

Project setup

# Parameters
PROJECT_NAME = "3LC Tutorials - Sugar Beets 2016"
DATASET_NAME = "sugarbeets-2016"
TRAIN_TABLE_NAME = "train"
VAL_TABLE_NAME = "val"
RUN_NAME = "Train SegFormer"
DIAGNOSTICS_RUN_NAME = "Diagnose SegFormer"

# Data source
# DOWNLOAD_PATH must contain the prepared dataset folder "sugarbeets-2016" (images/, masks/, splits/).
DOWNLOAD_PATH = "../../transient_data"
# >>> LOCAL SWITCH: point this at your local copy to run the notebook here, e.g.
# DOWNLOAD_PATH = "your_local_path"

# Model and training
BACKBONE = "nvidia/mit-b0"  # SegFormer encoder (ImageNet-pretrained)
EPOCHS = 30
BATCH_SIZE = 8
LR = 6e-5  # encoder LR (AdamW) — SegFormer's standard base LR
HEAD_LR_MULT = 10.0  # the fresh decode head trains at LR * HEAD_LR_MULT
WEIGHT_DECAY = 0.01
CROP_SIZE = 512  # training random-crop size (multiple of 32)
EVAL_SIZE = 768  # eval/full-image inference size (multiple of 32; native is 1296x966)
AUGMENT = True
COLLECT_FREQUENCY = 10  # write the heavy per-sample metrics every N epochs (and the final epoch)
EVAL_BATCH_SIZE = 8  # eval at 768 (~0.56x the memory of 1024) leaves room for a larger batch
NUM_WORKERS = 4  # set to 0 if running locally in a notebook kernel on macOS
N_TRAIN = 500  # cap train rows for a quick run; None keeps the whole split
N_VAL = 100  # cap val rows for a quick run; None keeps the whole split
REDUCTION_METHOD = "umap"  # embedding reducer: "pacmap" or "umap"
INV_LOG_C = 1.10  # ENet/ERFNet inverse-log constant for the per-class weights
SEED = 42
INSTALL_DEPENDENCIES = True

# Writable scratch dir for SegFormer checkpoints.
# The best checkpoint is written in the training run and re-read by the diagnostics run.
CHECKPOINT_DIR = "../../tmp/segformer_sugarbeets"

Install dependencies

if INSTALL_DEPENDENCIES:
    %pip install -q 3lc torch torchvision transformers tqdm pacmap umap-learn matplotlib

Runtime setup

Set a CUDA allocator option that reduces memory fragmentation during the memory-heavy evaluation (a no-op when not running on CUDA).

import os

# Reduce CUDA memory fragmentation ahead of the memory-heavy evaluation (set before CUDA initializes; no-op off-CUDA).
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

Imports

import contextlib
import random
import time
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import tlc
import torch
import torch.nn as nn
import torch.nn.functional as F  # noqa: N812
import torchvision.transforms.functional as TF  # noqa: N812
from matplotlib import colormaps
from PIL import Image
from tlc.constants import EXAMPLE_ID, SAMPLE_WEIGHT
from tlc.integration.torch.samplers import create_sampler
from tlc.metrics.semantic_segmentation import semantic_segmentation_metrics
from tlc.schemas import MapElement, SampleWeightSchema, SemanticSegmentationRleSchema
from tlc.schemas.values import DimensionNumericValue, Float32Value
from torch.utils.data import DataLoader, Dataset
from tqdm.auto import tqdm
from transformers import SegformerForSemanticSegmentation

Class scheme

There are three classes:

• soil

• sugarbeet

• weed

All classes are treated as real classes, with no implicit background in this dataset. (If you instead pass a class id as background= to the mask schema, it is dropped from the value map and rendered transparent in the Dashboard.) Here we keep soil as a normal class with a blue display_color so it stays visible and is scored like the others. This dataset has no void/ignore class.

CLASS_NAMES = ["soil", "sugarbeet", "weed"]  # class ids 0 / 1 / 2 — all real classes (no implicit background)
SOIL_ID = 0
NUM_CLASSES = len(CLASS_NAMES)

CLASS_COLORS = {
    0: MapElement("soil", display_color="#2563ebff"),  # blue
    1: MapElement("sugarbeet", display_color="#22c55eff"),  # green (crop)
    2: MapElement("weed", display_color="#ef4444ff"),  # red
}
# No implicit background in this dataset: every id (incl. soil) stays a real class in the value map,
# so soil renders in its blue color (not transparent) and is scored like any other class.
GT_CLASSES = CLASS_COLORS  # frozen API: pass the classes dict directly (no semseg_classes helper)
PREDICTED_CLASSES = CLASS_COLORS
CLASS_IDS = sorted(GT_CLASSES)

Locate the dataset

DOWNLOAD_PATH must contain the prepared sugarbeets-2016/ folder (images/, masks/, splits/).

Running this notebook yourself?

To run this notebook you need the Sugar Beets 2016 dataset in the images/ · masks/ · splits/ layout the cells below expect, with DOWNLOAD_PATH pointing at the folder that contains it. The steps below build that layout from the raw University-of-Bonn dataset.

1. Download & extract the raw crop/weed annotations — Sugar Beets 2016 (Chebrolu et al., IJRR 2017), ~21.4 GB.

   The annotations ship as 11 multipart RAR archives. Split RAR needs a real extractor (macOS’s bsdtar is unreliable for it), so install unar first. The download is resumable — just re-run the loop if it drops.

   # One-time: a multipart-RAR extractor.
   brew install unar                       # or: brew install sevenzip  (gives 7zz)
   
   BASE="http://www.ipb.uni-bonn.de/datasets_IJRR2017/annotations/cropweed"
   DEST="$HOME/Downloads/sugarbeets_2016"
   mkdir -p "$DEST/archives" && cd "$DEST/archives"
   
   # Resumable download of part01 … part11 (curl -C - resumes a partial file).
   for n in $(seq -w 1 11); do
   curl -L -f -C - --retry 5 --retry-delay 5 -O \
       "$BASE/ijrr_sugarbeets_2016_annotations.part${n}.rar"
   done
   
   # Extract by pointing the tool at part01 — it pulls in the rest of the set.
   unar -force-overwrite -output-directory "$DEST" \
   "ijrr_sugarbeets_2016_annotations.part01.rar"
   
   # Extracted layout (21 CKA_* capture folders):
   #   $DEST/ijrr_sugarbeets_2016_annotations/CKA_*/images/rgb/*.png
   #   $DEST/ijrr_sugarbeets_2016_annotations/CKA_*/annotations/dlp/iMapCleaned/*.png
   

2. Convert to the 3-class layout — pair, filter, remap, split, write.

   The raw iMapCleaned masks use the authors’ own integer label scheme (soil / crop / several weed types), and the dataset ships no train/val split. This step (a) pairs each RGB with its mask by shared stem and drops any frame with no matching annotation; (b) remaps the indices to 0 = soil, 1 = sugarbeet, 2 = weed, raising on any value it doesn’t recognise so nothing is silently mislabeled; and (c) writes a seeded, deterministic 80/20 train/val split into the flat images/ · masks/ · splits/ layout (~12,330 frames → 9,864 train / 2,466 val).

   """Reshape the raw Bonn dataset into sugarbeets-2016/{images,masks,splits}.
   Run once, outside the notebook. Requires: numpy, Pillow."""
   import random
   import shutil
   from pathlib import Path
   
   import numpy as np
   from PIL import Image
   
   SRC = Path.home() / "Downloads/sugarbeets_2016/ijrr_sugarbeets_2016_annotations"
   DOWNLOAD_PATH = Path("~/Datasets/sugarbeets-2016-3lc").expanduser()   # match this notebook's DOWNLOAD_PATH
   OUT = DOWNLOAD_PATH / "sugarbeets-2016"
   SEED, TRAIN_FRAC = 42, 0.80
   
   # iMapCleaned integer -> 3-class id, from the dataset's own labelMap:
   #   soil {0,1,97}; crop {10000..10002}; weed = generic {2} + dycot {20000..20011} + grass {20100..20105}.
   INDEX_TO_CLASS = {
       0: 0, 1: 0, 97: 0,
       10000: 1, 10001: 1, 10002: 1,
       2: 2,
       **{i: 2 for i in range(20000, 20012)},
       **{i: 2 for i in range(20100, 20106)},
   }
   RGB_REL, INDEX_REL = Path("images/rgb"), Path("annotations/dlp/iMapCleaned")
   
   # (a) Pair RGB <-> mask within each CKA_* capture folder; drop RGBs with no matching mask.
   #     key = <cra>__<stem> so identical stems in different folders never collide.
   pairs, unpaired = {}, 0
   for cra in sorted(p for p in SRC.iterdir() if (p / RGB_REL).is_dir()):
       masks = {p.stem: p for p in (cra / INDEX_REL).glob("*.png")}
       for rgb in sorted((cra / RGB_REL).glob("*.png")):
           mask = masks.get(rgb.stem)
           if mask is None:            # no proper annotation -> drop this frame
               unpaired += 1
               continue
           pairs[f"{cra.name}__{rgb.stem}"] = (rgb, mask)
   print(f"paired {len(pairs)} samples; dropped {unpaired} RGB with no matching mask")
   
   def remap(mask_path: Path) -> np.ndarray:
       """iMapCleaned -> uint8 (H, W) in {0,1,2}; raises on any index not in INDEX_TO_CLASS."""
       arr = np.asarray(Image.open(mask_path)).astype(np.int64)
       out, seen = np.zeros(arr.shape, np.uint8), np.zeros(arr.shape, bool)
       for idx, cls in INDEX_TO_CLASS.items():
           hit = arr == idx
           out[hit], seen = cls, seen | hit
       if not seen.all():
           bad = sorted(int(v) for v in np.unique(arr[~seen]))[:10]
           raise SystemExit(f"{mask_path.name}: unmapped indices {bad} — update INDEX_TO_CLASS")
       return out
   
   # (b) Seeded, deterministic 80/20 train/val split over all kept keys (incl. soil-only frames).
   #     sorted-then-shuffle makes it independent of filesystem walk order.
   keys = sorted(pairs)
   random.Random(SEED).shuffle(keys)
   n_train = round(len(keys) * TRAIN_FRAC)
   splits = {"train": keys[:n_train], "val": keys[n_train:]}
   
   # (c) Write images/<key>.png (bytes copied verbatim), masks/<key>.png (remapped), splits/{train,val}.txt
   for sub in ("images", "masks", "splits"):
       (OUT / sub).mkdir(parents=True, exist_ok=True)
   for split, ks in splits.items():
       for key in ks:
           rgb, mask = pairs[key]
           shutil.copy2(rgb, OUT / "images" / f"{key}.png")
           Image.fromarray(remap(mask), mode="L").save(OUT / "masks" / f"{key}.png")
       (OUT / "splits" / f"{split}.txt").write_text("\n".join(ks) + "\n")
       print(f"{split}: {len(ks)} images + masks")
   

Then set DOWNLOAD_PATH (in the Parameters cell) to the folder that contains sugarbeets-2016/ — the notebook will pick up images/, masks/, and splits/ from there.

download_root = Path(DOWNLOAD_PATH).expanduser()
DATA_ROOT = download_root / "sugarbeets-2016"

if not (DATA_ROOT / "images").is_dir() or not (DATA_ROOT / "splits").is_dir():
    raise FileNotFoundError(
        f"Sugar Beets 2016 not found at {DATA_ROOT.resolve()}.\n"
        "Set DOWNLOAD_PATH to the folder that CONTAINS 'sugarbeets-2016' (with images/, masks/, splits/).\n"
        "See the 'Running this notebook yourself?' section above to prepare it from the raw dataset."
    )


def read_keys(split: str) -> list[str]:
    return (DATA_ROOT / "splits" / f"{split}.txt").read_text().split()


print(f"Dataset at {DATA_ROOT.resolve()} | train {len(read_keys('train'))} | val {len(read_keys('val'))}")

Build the 3LC Tables

One Table per split. Each row carries

• the RGB image,

• the 3-class segmentation, and

• the built-in per-row weight (sample weight) initialized to 1.0 and editable in the Dashboard.

def load_segmentation(stem: str) -> np.ndarray:
    """Dense (H, W) int32 label map. The RLE mask column stores the raw array directly."""
    return np.asarray(Image.open(DATA_ROOT / "masks" / f"{stem}.png")).astype(np.int32)


def write_table(split: str, table_name: str, n: int | None) -> tlc.Table:
    keys = read_keys(split)
    if n is not None:
        keys = random.Random(SEED).sample(keys, min(n, len(keys)))
    # RLE-backed semantic-segmentation mask. No `background=` is declared, so every class (incl.
    # soil, id 0) stays in the value map and renders in its own color in the Dashboard.
    segmentation_schema = SemanticSegmentationRleSchema(classes=CLASS_COLORS, display_name="segmentation")
    # Name the class sub-column "ground truth" so it does not read as a generic "label" in the Dashboard.
    segmentation_schema.values["instance_properties"].values["label"].display_name = "ground truth"
    schema = {
        "image": tlc.schemas.ImageSchema(),
        "segmentation": segmentation_schema,
        SAMPLE_WEIGHT: SampleWeightSchema(),  # built-in per-row sample weight (editable in the Dashboard)
    }
    writer = tlc.TableWriter(
        project_name=PROJECT_NAME,
        dataset_name=DATASET_NAME,
        table_name=table_name,
        schema=schema,
        if_exists="overwrite",
    )
    for stem in keys:
        writer.add_row(
            {
                "image": tlc.Url(DATA_ROOT / "images" / f"{stem}.png"),  # aliases auto-applied on write
                "segmentation": load_segmentation(stem),  # raw (H, W) array; RLE-encoded on write
                SAMPLE_WEIGHT: 1.0,
            }
        )
    table = writer.finalize()
    print(f"Wrote {table_name}: {len(table)} rows -> {table.url}")
    return table


train_table = write_table("train", TRAIN_TABLE_NAME, N_TRAIN)
val_table = write_table("val", VAL_TABLE_NAME, N_VAL)

sample = train_table[0]
seg = sample["segmentation"]  # SemanticSegmentation object; dense label map on `.mask`
fig, axes = plt.subplots(1, 2, figsize=(10, 4))
axes[0].imshow(sample["image"])
axes[0].set_title("image")
axes[1].imshow(sample["image"])
axes[1].imshow(seg.mask, alpha=0.5, cmap="viridis", interpolation="nearest")
axes[1].set_title("soil / sugarbeet / weed")
for ax in axes:
    ax.axis("off")
plt.tight_layout()
plt.show()
print("classes present in this sample:", sorted(set(np.unique(seg.mask).tolist())))

SegFormer model

A SegFormer encoder (ImageNet-pretrained mit-b0) with a fresh NUM_CLASSES decode head. The contiguous_views context manager is an MPS-safety shim — the MiT encoder does transpose(...).view(...) whose autograd backward fails on Apple MPS; routing view/reshape through .contiguous() is behaviour-preserving (a no-op on CUDA/CPU).

IMAGENET_MEAN = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
IMAGENET_STD = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)


@contextlib.contextmanager
def contiguous_views():
    orig_view, orig_reshape = torch.Tensor.view, torch.Tensor.reshape
    torch.Tensor.view = lambda s, *a, **k: orig_reshape(s.contiguous(), *a, **k)
    torch.Tensor.reshape = lambda s, *a, **k: orig_reshape(s.contiguous(), *a, **k)
    try:
        yield
    finally:
        torch.Tensor.view, torch.Tensor.reshape = orig_view, orig_reshape


def build_segformer(num_classes: int, backbone: str) -> nn.Module:
    id2label = {i: n for i, n in enumerate(CLASS_NAMES)}
    return SegformerForSemanticSegmentation.from_pretrained(
        backbone,
        num_labels=num_classes,
        id2label=id2label,
        label2id={n: i for i, n in id2label.items()},
        ignore_mismatched_sizes=True,
    )


def normalize(image01: torch.Tensor) -> torch.Tensor:
    return (image01 - IMAGENET_MEAN) / IMAGENET_STD


def preprocess_image(image: Image.Image, size: int) -> torch.Tensor:
    image = image.convert("RGB").resize((size, size))
    return normalize(torch.from_numpy(np.asarray(image).copy()).permute(2, 0, 1).float() / 255.0)

Per-class weights and the data pipeline

The training weight combines the global per-class weight and the per-row sample weight.

• The global per-class weight is derived from the current training table using the ENet inverse-log for each class \(k\): \(w_k = 1 / ln(c + p_k)\), where \(p_k\) is the pixel frequency of class \(k\) and \(c\) is a constant, INV_LOG_C.

• The per-row sample weight is the editable 3LC weight column that adapts to the current table revision. create_sampler drops zero-weight rows from sampling.

The effective per-pixel training weight is global per-class weight \(\times\) per-row sample weight.

def included_indices(table: tlc.Table) -> list[int]:
    """Row indices with sample weight > 0 (the rows actually used for training and evaluation)."""
    return [i for i, r in enumerate(table.table_rows) if float(r.get(SAMPLE_WEIGHT, 1.0)) > 0.0]


def compute_class_weights(table: tlc.Table, c: float) -> tuple[np.ndarray, np.ndarray]:
    """Per-class weights from the current table's per-class pixel frequencies (ENet inverse-log)."""
    counts = np.zeros(NUM_CLASSES, dtype=np.int64)
    for idx in tqdm(included_indices(table), desc="class-weights", leave=False):
        label_map = table[idx]["segmentation"].mask
        counts += np.bincount(label_map.ravel(), minlength=NUM_CLASSES)[:NUM_CLASSES]
    freq = counts / max(1, counts.sum())
    return (1.0 / np.log(c + freq)).astype(np.float32), freq


class SugarBeetsSegDataset(Dataset):
    """Table row -> (normalized image, label map, sample weight). Training augments with random
    scale -> crop -> flip -> photometric jitter; eval just resizes to ``size``."""

    def __init__(
        self, table: tlc.Table, size: int, *, augment: bool = False, scale_range: tuple[float, float] = (0.5, 2.0)
    ) -> None:
        self.table, self.size, self.augment, self.scale_range = table, size, augment, scale_range

    def __len__(self) -> int:
        return len(self.table)

    def __getitem__(self, idx: int):
        row = self.table[idx]
        image = row["image"].convert("RGB")
        label = Image.fromarray(row["segmentation"].mask.astype(np.uint8), mode="L")
        if self.augment:
            image, label = self._augment(image, label)
        else:
            image = image.resize((self.size, self.size))
            label = label.resize((self.size, self.size), Image.NEAREST)
        image_t = normalize(torch.from_numpy(np.asarray(image).copy()).permute(2, 0, 1).float() / 255.0)
        label_t = torch.from_numpy(np.asarray(label).copy()).long()
        weight = float(self.table.table_rows[idx].get(SAMPLE_WEIGHT, 1.0))
        return image_t, label_t, torch.tensor(weight, dtype=torch.float32)

    def _augment(self, image: Image.Image, label: Image.Image):
        s = random.uniform(*self.scale_range)
        w, h = image.size
        nw, nh = max(1, round(w * s)), max(1, round(h * s))
        image = image.resize((nw, nh), Image.BILINEAR)
        label = label.resize((nw, nh), Image.NEAREST)
        size = self.size
        pad_w, pad_h = max(0, size - nw), max(0, size - nh)
        if pad_w or pad_h:
            image = TF.pad(image, [0, 0, pad_w, pad_h], fill=0)
            label = TF.pad(label, [0, 0, pad_w, pad_h], fill=SOIL_ID)
        w, h = image.size
        x0, y0 = random.randint(0, w - size), random.randint(0, h - size)
        image, label = TF.crop(image, y0, x0, size, size), TF.crop(label, y0, x0, size, size)
        if random.random() < 0.5:
            image, label = TF.hflip(image), TF.hflip(label)
        image = TF.adjust_brightness(image, random.uniform(0.6, 1.4))
        image = TF.adjust_contrast(image, random.uniform(0.6, 1.4))
        image = TF.adjust_saturation(image, random.uniform(0.8, 1.2))
        return image, label

Optimizer, schedule and weighted loss

def make_optimizer(model: nn.Module, lr: float, head_lr_mult: float, weight_decay: float):
    """AdamW with a lower LR on the pretrained encoder and a higher LR on the fresh decode head."""
    encoder, head = [], []
    for name, param in model.named_parameters():
        if param.requires_grad:
            (head if "decode_head" in name else encoder).append(param)
    return torch.optim.AdamW(
        [{"params": encoder, "lr": lr}, {"params": head, "lr": lr * head_lr_mult}],
        weight_decay=weight_decay,
    )


def make_scheduler(optimizer, total_steps: int, warmup_steps: int):
    """Linear warmup then linear (poly power 1.0) decay to ~0, stepped per iteration."""

    def lr_lambda(step: int) -> float:
        if step < warmup_steps:
            return (step + 1) / max(1, warmup_steps)
        progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)
        return max(0.0, 1.0 - progress)

    return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)


def seg_loss(model, images, labels, sample_weights, class_weights):
    """Per-class x per-sample weighted cross-entropy at the label resolution."""
    with contiguous_views():
        logits = model(pixel_values=images).logits
    logits = F.interpolate(logits, size=labels.shape[-2:], mode="bilinear", align_corners=False)
    loss_map = F.cross_entropy(logits, labels, weight=class_weights, reduction="none")  # (B,H,W), class-weighted
    sw = sample_weights.view(-1, 1, 1)
    weight_map = class_weights[labels] * sw  # per-pixel applied weight (class weight x sample weight)
    return (loss_map * sw).sum() / weight_map.sum().clamp_min(1.0)


def seed_everything(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

Metrics collection

Evaluation runs each image at EVAL_SIZE, then upsamples the logits back to the native mask resolution for scoring. Every call returns the dataset-level mIoU and per-class IoU, which the training loop logs as dense per-epoch curves. On collection epochs it additionally writes the per-sample metrics — predicted segmentation, IoU, loss, entropy, and a pooled embedding — via run.add_metrics, keyed back to the foreign table by example_id.

class _EvalDataset(Dataset):
    def __init__(self, table, indices, eval_size):
        self.table, self.indices, self.eval_size = table, indices, eval_size

    def __len__(self):
        return len(self.indices)

    def __getitem__(self, i):
        idx = self.indices[i]
        row = self.table[idx]
        image = row["image"].convert("RGB")
        width, height = image.size
        return preprocess_image(image, self.eval_size), (width, height), row["segmentation"].mask, idx


def _eval_collate(batch):
    return (torch.stack([b[0] for b in batch]), [b[1] for b in batch], [b[2] for b in batch], [b[3] for b in batch])


def collect_metrics(run, table, model, device, eval_size, epoch, *, write_per_sample, eval_batch=16, num_workers=4):
    preds, ious, conf_mats, losses, entropies, embeddings, example_ids = [], [], [], [], [], [], []
    indices = included_indices(table)
    loader = DataLoader(
        _EvalDataset(table, indices, eval_size),
        batch_size=eval_batch,
        shuffle=False,
        num_workers=num_workers,
        collate_fn=_eval_collate,
    )
    class_ids = list(range(NUM_CLASSES))

    if torch.cuda.is_available():  # release training-time fragmentation before the eval forward
        torch.cuda.empty_cache()

    model.eval()
    with torch.no_grad():
        desc = "collect" if write_per_sample else "eval"
        for pixel_values, sizes, label_maps, row_indices in tqdm(loader, desc=desc, leave=False):
            pixel_values = pixel_values.to(device)
            with contiguous_views():
                outputs = model(pixel_values=pixel_values, output_hidden_states=True)
            logits_small, feats = outputs.logits, outputs.hidden_states[-1]
            for b in range(pixel_values.shape[0]):
                width, height = sizes[b]
                label_map = label_maps[b]
                logit = F.interpolate(
                    logits_small[b : b + 1], size=(height, width), mode="bilinear", align_corners=False
                )
                pred_map = logit.argmax(dim=1).squeeze(0).cpu().numpy().astype(np.int32)
                # All classes (incl. soil) are real classes in GT_CLASSES, so all are scored — there is
                # no background to exclude or re-add.
                m = semantic_segmentation_metrics(pred_map, label_map, GT_CLASSES)
                ious.append(m["mean_iou"])
                conf_mats.append([int(x) for r in m["confusion_matrix"] for x in r])
                class_ids = m["class_ids"]
                example_ids.append(row_indices[b])
                if write_per_sample:
                    embeddings.append(feats[b].mean(dim=(1, 2)).cpu().numpy().astype(np.float32))
                    target = torch.from_numpy(label_map).long()[None].to(device)
                    losses.append(F.cross_entropy(logit, target).item())
                    probs = logit.softmax(dim=1)
                    entropies.append(float((-(probs * probs.clamp_min(1e-12).log()).sum(dim=1)).mean()))
                    # RLE mask column takes the raw (H, W) label map directly.
                    preds.append(pred_map)

    if write_per_sample:
        # No `background=` -> soil (id 0) renders in its color, matching the GT mask column.
        predicted_schema = SemanticSegmentationRleSchema(classes=CLASS_COLORS, display_name="predicted segmentation")
        predicted_schema.values["instance_properties"].values["label"].display_name = "prediction"
        run.add_metrics(
            {
                "predicted_segmentation": preds,
                "iou": ious,
                "loss": losses,
                "entropy": entropies,
                "embedding": embeddings,
                "epoch": [epoch] * len(ious),
                EXAMPLE_ID: example_ids,
            },
            schema={
                "predicted_segmentation": predicted_schema,
                "embedding": tlc.schemas.EmbeddingSchema(shape=len(embeddings[0])),
            },
            foreign_table_url=table.url,
        )

    cm = np.asarray(conf_mats).reshape(-1, NUM_CLASSES, NUM_CLASSES).sum(axis=0)
    per_class = {}
    for i, class_id in enumerate(class_ids):
        tp = cm[i, i]
        union = cm[i, :].sum() + cm[:, i].sum() - tp
        per_class[CLASS_NAMES[class_id]] = float(tp / union) if union > 0 else float("nan")
    return {"miou": float(np.nanmean(list(per_class.values()))), **per_class}

Run #1: Training Run

Each epoch we train, then evaluate the validation set — logging a dense per-epoch val mIoU curve with tlc.log and checkpointing the best model. The full training-set evaluation is ~4x heavier, so we run it only on collection epochs (every COLLECT_FREQUENCY, and the final epoch); that is also when the heavy per-sample metrics (predictions, embeddings, …) are written for both splits. To evaluate only a subset for speed, set N_TRAIN / N_VAL. The best model (by val mIoU) is checkpointed locally so the diagnostics run can reload it.

seed_everything(SEED)
device = torch.device("cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu")
print("Using device:", device)

class_weights_np, train_freq = compute_class_weights(train_table, INV_LOG_C)
class_weights = torch.tensor(class_weights_np, device=device)
print("train class frequencies:", dict(zip(CLASS_NAMES, [round(float(f), 5) for f in train_freq])))
print("per-class weights:", dict(zip(CLASS_NAMES, [round(float(w), 3) for w in class_weights_np])))

train_loader = DataLoader(
    SugarBeetsSegDataset(train_table, CROP_SIZE, augment=AUGMENT),
    batch_size=BATCH_SIZE,
    sampler=create_sampler(train_table, weighted=False, shuffle=True),
    num_workers=NUM_WORKERS,
    drop_last=True,
)

model = build_segformer(NUM_CLASSES, BACKBONE).to(device)
optimizer = make_optimizer(model, LR, HEAD_LR_MULT, WEIGHT_DECAY)
total_steps = EPOCHS * max(1, len(train_loader))
scheduler = make_scheduler(optimizer, total_steps, warmup_steps=int(0.05 * total_steps))

run = tlc.init(
    PROJECT_NAME,
    run_name=RUN_NAME,
    description="SegFormer on Sugar Beets 2016 (crop-train / full-eval, per-class x per-sample weighted CE)",
    parameters={
        "backbone": BACKBONE,
        "epochs": EPOCHS,
        "batch_size": BATCH_SIZE,
        "lr": LR,
        "crop_size": CROP_SIZE,
        "eval_size": EVAL_SIZE,
        "seed": SEED,
        "train_table": train_table.url.to_str(),
    },
)

best_val_miou = float("-inf")
best_checkpoint = CHECKPOINT_DIR + "_best"
Path(CHECKPOINT_DIR).parent.mkdir(parents=True, exist_ok=True)

for epoch in range(EPOCHS):
    epoch_start = time.time()
    model.train()
    train_loss = 0.0
    print(f"[epoch {epoch}/{EPOCHS}] training...", flush=True)
    for images, labels, sample_weights in tqdm(train_loader, desc=f"epoch {epoch}", leave=False):
        images, labels, sample_weights = images.to(device), labels.to(device), sample_weights.to(device)
        optimizer.zero_grad()
        loss = seg_loss(model, images, labels, sample_weights, class_weights)
        loss.backward()
        optimizer.step()
        scheduler.step()
        train_loss += loss.item() * images.shape[0]
    train_loss /= len(train_loader.sampler)

    write_per_sample = (epoch + 1) % COLLECT_FREQUENCY == 0 or epoch == EPOCHS - 1

    # Val is evaluated EVERY epoch (dense curve + best-checkpoint signal). The full-train eval is ~4x
    # heavier, so we only run it on collection epochs to keep training time manageable.
    print(f"[epoch {epoch}] evaluating val ({len(val_table)} images)...", flush=True)
    val_stats = collect_metrics(
        run,
        val_table,
        model,
        device,
        EVAL_SIZE,
        epoch,
        write_per_sample=write_per_sample,
        eval_batch=EVAL_BATCH_SIZE,
        num_workers=NUM_WORKERS,
    )

    entry = {
        "epoch": epoch,
        "lr": optimizer.param_groups[0]["lr"],
        "train_loss": train_loss,
        "val_miou": val_stats["miou"],
    }
    for c in CLASS_NAMES:
        entry[f"val_iou_{c}"] = val_stats[c]

    if write_per_sample:
        print(f"[epoch {epoch}] evaluating train ({len(train_table)} images, collection epoch)...", flush=True)
        train_stats = collect_metrics(
            run,
            train_table,
            model,
            device,
            EVAL_SIZE,
            epoch,
            write_per_sample=True,
            eval_batch=EVAL_BATCH_SIZE,
            num_workers=NUM_WORKERS,
        )
        entry["train_miou"] = train_stats["miou"]
        for c in CLASS_NAMES:
            entry[f"train_iou_{c}"] = train_stats[c]

    epoch_secs = time.time() - epoch_start
    tlc.log(entry)
    msg = f"epoch {epoch}  {epoch_secs:.0f}s  train_loss={train_loss:.4f}  val_miou={val_stats['miou']:.4f}"
    if write_per_sample:
        msg += f"  train_miou={train_stats['miou']:.4f}"
    print(msg, flush=True)

    # Save the best model by val mIoU
    if val_stats["miou"] > best_val_miou:
        best_val_miou = val_stats["miou"]
        model.save_pretrained(best_checkpoint)
        print(f"  new best val_miou={best_val_miou:.4f} -> {best_checkpoint}")

print(f"Reducing embeddings ({REDUCTION_METHOD})...")
run.reduce_embeddings_by_foreign_table_url(val_table.url, method=REDUCTION_METHOD, delete_source_tables=True)
run.set_status_completed()
print("Training run:", run.url)

Run #2: Diagnostics Run on the best model

This second Run loads the best checkpoint and collects richer per-sample diagnostics on train and val, all in one shared embedding space:

• predicted segmentation

• an error map (correct/incorrect)

• a confidence map (binned max-softmax)

• per-class IoU

• a pooled per-sample embedding (as in training)

• per-class GT-masked and prediction-masked embeddings

# Error-map and confidence-map "class" palettes.
ERROR_CORRECT, ERROR_INCORRECT = 0, 1
# Full value map incl. the "correct" class. Passed to the RLE schema with ERROR_CORRECT as the
# background (transparent in the Dashboard), so only the incorrect pixels are highlighted.
ERROR_CLASS_MAP = {
    ERROR_CORRECT: MapElement("correct"),
    ERROR_INCORRECT: MapElement("incorrect", display_color="#e6194bff"),
}
N_BINS = 10


def turbo_bin_classes(n_bins: int) -> dict[int, MapElement]:
    turbo = colormaps["turbo"]
    classes = {}
    for i in range(n_bins):
        r, g, b, _ = turbo((i + 0.5) / n_bins)
        alpha = int(255 * min(1.0, (i + 1) / (n_bins / 3)))
        lo, hi = round(100 * i / n_bins), round(100 * (i + 1) / n_bins)
        classes[i] = MapElement(
            f"{lo}-{hi}%", display_color=f"#{int(r * 255):02x}{int(g * 255):02x}{int(b * 255):02x}{alpha:02x}"
        )
    return classes


CONFIDENCE_CLASSES = turbo_bin_classes(N_BINS)
PREDICTED_LAYER_IDS = sorted(PREDICTED_CLASSES)


def layer_embedding_schema(display_name: str) -> tlc.Schema:
    return tlc.Schema(
        display_name=display_name,
        value=Float32Value(),
        size0=DimensionNumericValue.fixed_size(2),
        size1=DimensionNumericValue(),
    )


def layer_scalar_schema(display_name: str) -> tlc.Schema:
    return tlc.Schema(display_name=display_name, value=Float32Value(), size0=DimensionNumericValue())


def masked_average_pool(features: torch.Tensor, mask: torch.Tensor) -> np.ndarray:
    mask = mask.to(features.dtype)
    denom = mask.sum()
    if denom < 1:
        return np.full(features.shape[0], np.nan, dtype=np.float32)
    return ((features * mask[None]).sum(dim=(1, 2)) / denom).cpu().numpy().astype(np.float32)


def _downsample_to(label_map, h, w, device):
    small = F.interpolate(torch.from_numpy(label_map.astype(np.float32))[None, None], size=(h, w), mode="nearest")
    return small.squeeze().to(device)

import umap


def collect_raw_diagnostics(table, model, device, image_size, *, light=False):
    # `light=True` collects ONLY the per-class masked embeddings (tiny) for the shared UMAP fit, holding
    # none of the full-resolution maps. The full pass (light=False) additionally builds the heavy
    # per-sample columns; the diagnostics cell writes them one table at a time so peak host RAM does not
    # scale with train+val combined.
    keys = (
        ("gt_masked", "pred_masked")
        if light
        else (
            "predicted_segmentation",
            "error_map",
            "error_fraction",
            "confidence_map",
            "class_iou",
            "embedding",
            "gt_masked",
            "pred_masked",
        )
    )
    out = {k: [] for k in keys}
    model.eval()
    with torch.no_grad():
        for idx in tqdm(range(len(table)), desc="diagnose", leave=False):
            row = table[idx]
            image = row["image"].convert("RGB")
            label_map = row["segmentation"].mask
            width, height = image.size

            pixel_values = preprocess_image(image, image_size)[None].to(device)
            with contiguous_views():
                outputs = model(pixel_values=pixel_values, output_hidden_states=True)
            logits = F.interpolate(outputs.logits, size=(height, width), mode="bilinear", align_corners=False)
            features = outputs.hidden_states[-1].squeeze(0)
            probs = logits.softmax(dim=1).squeeze(0).cpu().numpy()
            pred_map = probs.argmax(axis=0).astype(np.int32)  # int32 locally for metrics + masking

            # Per-class masked embeddings (tiny) — always collected; the only thing kept in light mode.
            fh, fw = features.shape[1], features.shape[2]
            gt_small = _downsample_to(label_map, fh, fw, device)
            pred_small = _downsample_to(pred_map, fh, fw, device)
            out["gt_masked"].append(
                np.stack([masked_average_pool(features, gt_small == c) for c in PREDICTED_LAYER_IDS])
            )
            out["pred_masked"].append(
                np.stack([masked_average_pool(features, pred_small == c) for c in PREDICTED_LAYER_IDS])
            )
            if light:
                continue

            # Heavy per-sample columns. Dense maps are stored as uint8 (class ids <= 255): 4x less host
            # RAM than int32, and byte-identical once RLE-encoded on write.
            # RLE mask columns take the raw (H, W) label maps directly.
            out["predicted_segmentation"].append(pred_map.astype(np.uint8))
            out["embedding"].append(features.mean(dim=(1, 2)).cpu().numpy().astype(np.float32))

            incorrect = pred_map != label_map
            err = np.where(incorrect, ERROR_INCORRECT, ERROR_CORRECT).astype(np.uint8)
            out["error_map"].append(err)
            out["error_fraction"].append(float(incorrect.mean()))

            bins = np.clip(probs.max(axis=0) * N_BINS, 0, N_BINS - 1).astype(np.uint8)
            out["confidence_map"].append(bins)

            # Soil is a normal class (no implicit background), so all classes are scored.
            m = semantic_segmentation_metrics(pred_map, label_map, GT_CLASSES)
            iou_by_class = dict(zip(m["class_ids"], m["per_class_iou"]))
            out["class_iou"].append([float(iou_by_class.get(c, float("nan"))) for c in PREDICTED_LAYER_IDS])
    return out


def reduce_groups_to_shared_2d(groups):
    names = list(groups)
    stacks = {n: np.stack(groups[n]) for n in names}
    n_rows, n_layers, dim = stacks[names[0]].shape
    matrix = np.concatenate([stacks[n].reshape(n_rows * n_layers, dim) for n in names]).astype(np.float32)
    present = np.isfinite(matrix).all(axis=1)
    reduced = np.full((matrix.shape[0], 2), np.nan, dtype=np.float32)
    reduced[present] = umap.UMAP(n_components=2).fit_transform(matrix[present]).astype(np.float32)
    out = {}
    for i, n in enumerate(names):
        block = reduced[i * n_rows * n_layers : (i + 1) * n_rows * n_layers]
        out[n] = list(block.reshape(n_rows, n_layers, 2).astype(np.float32))
    return out

import gc

# Free the training model + cached GPU memory before loading the diagnostics model (one model at a time).
for _name in ("model", "optimizer", "scheduler"):
    globals().pop(_name, None)
if torch.cuda.is_available():
    torch.cuda.empty_cache()

diag_model = SegformerForSemanticSegmentation.from_pretrained(best_checkpoint).to(device)
diag_run = tlc.init(
    PROJECT_NAME,
    run_name=DIAGNOSTICS_RUN_NAME,
    description="SegFormer diagnostics: error / confidence maps, per-class IoU and embeddings",
)

diag_tables = {"train": train_table, "val": val_table}
print("Starting diagnostics run...", flush=True)

# Pass 1 (cheap): collect ONLY the per-class masked embeddings from every table and fit ONE shared 2D
# space. The full-resolution maps are NOT held here, so this pass's RAM does not scale with train+val.
gt_all, pred_all, counts = [], [], {}
for name, t in diag_tables.items():
    print(f"Diagnosing '{name}' embeddings ({len(t)} images)...", flush=True)
    light = collect_raw_diagnostics(t, diag_model, device, EVAL_SIZE, light=True)
    counts[name] = len(light["gt_masked"])
    gt_all += light["gt_masked"]
    pred_all += light["pred_masked"]
    del light
    gc.collect()

print("Reducing per-class masked embeddings into one shared 2D space (UMAP)...")
shared = reduce_groups_to_shared_2d({"gt": gt_all, "pred": pred_all})
del gt_all, pred_all
gc.collect()

# Pass 2 (per table): rebuild the heavy per-sample maps for ONE table, write them with that table's
# slice of the shared coordinates, then free before the next table. Peak host RAM is now bounded by the
# larger single split rather than train+val combined. (Row order matches pass 1, so the offset slice of
# `shared` aligns with `raw`.)
offset = 0
for name, table in diag_tables.items():
    print(f"Writing diagnostics for '{name}' ({counts[name]} images)...", flush=True)
    raw = collect_raw_diagnostics(table, diag_model, device, EVAL_SIZE)
    n = counts[name]
    # Name each class sub-column so the columns do not all read as a generic "label" in the Dashboard.
    # No `background=` on predicted -> soil (id 0) renders in its color, matching the GT mask column.
    predicted_schema = SemanticSegmentationRleSchema(classes=CLASS_COLORS, display_name="predicted segmentation")
    predicted_schema.values["instance_properties"].values["label"].display_name = "prediction"
    # The error map is a derived column: `correct` (id 0) is kept as the transparent background
    error_schema = SemanticSegmentationRleSchema(
        classes=ERROR_CLASS_MAP, background=ERROR_CORRECT, display_name="error map"
    )
    error_schema.values["instance_properties"].values["label"].display_name = "error"
    confidence_schema = SemanticSegmentationRleSchema(classes=CONFIDENCE_CLASSES, display_name="prediction confidence")
    confidence_schema.values["instance_properties"].values["label"].display_name = "confidence"
    diag_run.add_metrics(
        {
            "predicted_segmentation": raw["predicted_segmentation"],
            "error_map": raw["error_map"],
            "error_fraction": raw["error_fraction"],
            "confidence_map": raw["confidence_map"],
            "class_iou": raw["class_iou"],
            "embedding": raw["embedding"],
            "gt_masked_embedding": [c.tolist() for c in shared["gt"][offset : offset + n]],
            "pred_masked_embedding": [c.tolist() for c in shared["pred"][offset : offset + n]],
        },
        schema={
            "predicted_segmentation": predicted_schema,
            "error_map": error_schema,
            "confidence_map": confidence_schema,
            "class_iou": layer_scalar_schema("class IoU"),
            "embedding": tlc.schemas.EmbeddingSchema(shape=len(raw["embedding"][0])),
            "gt_masked_embedding": layer_embedding_schema("GT-masked embedding"),
            "pred_masked_embedding": layer_embedding_schema("pred-masked embedding"),
        },
        foreign_table_url=table.url,
    )
    offset += n
    del raw
    gc.collect()
    print(f"  wrote diagnostics for '{name}' ({n} samples)")

print(f"Reducing the pooled per-sample embeddings ({REDUCTION_METHOD})...")
diag_run.reduce_embeddings_by_foreign_table_url(val_table.url, method=REDUCTION_METHOD, delete_source_tables=True)
diag_run.set_status_completed()
print("Diagnostics run:", diag_run.url)

Part 2: Data Curation Techniques

With the training and diagnostics runs from Part 1 in place, we now use the 3LC Dashboard to curate the dataset. This part is a screenshots-and-narrative walkthrough — there is no code to run. We cover two techniques:

• Data reduction — remove redundant, low-information samples so the model trains on less data with little to no loss in accuracy.

• Correcting annotations — surface and fix ground-truth labels that are likely wrong.

Run #3: Full Diagnostics

For a more thorough investigation of your dataset, run the notebook again on the full dataset by setting both row caps to None in the Parameters cell:

N_TRAIN = None
N_VAL = None

You can also train longer and at higher resolution — increase EPOCHS, CROP_SIZE, and EVAL_SIZE for better convergence — and lower COLLECT_FREQUENCY to capture per-sample metrics more often. Once the model is well converged, apply the curation techniques below to the resulting run.

1. Data Reduction in Semantic Segmentation

We start from the diagnostics run and use the traversal index to find the least informative samples. The traversal index orders samples to maximize coverage of an embedding space: the earliest indices are the most diverse, representative samples, while the highest indices are the most redundant. Ranking by it therefore tells us which samples add little new information.

In classification, a single whole-image embedding is enough. Semantic segmentation is different — within one image some classes can be highly distinctive while others are repetitive — so we rank per class, using the per-class embeddings masked by ground truth (the GT-masked embedding column), and target samples that are redundant across several classes.

1. Open the GT-masked embedding column as a chart: select it and press 2.

2. Color the embeddings by ground-truth label by dragging the ground-truth column onto the chart.

3. Compute a traversal index for each class: right-click the GT-masked embedding column → Pick[0] to isolate the first class (soil), then right-click the isolated embedding → Traversal index. Repeat with Pick[1] (sugarbeet) and Pick[2] (weed).

4. Each per-class traversal index is a numeric column where low = unique, high = redundant, so the least-unique samples sit on the right. For each class, filter its traversal-index column to a band at the right-hand (high) end: leave the range’s max at the far right and drag the min handle inward. The further left you drag min, the wider the band is and the more of that class’s samples you mark for removal.

   Match each band’s width to the class’s difficulty — prune the hard classes conservatively and the easy ones aggressively:

   • weed (hardest): a narrow band at the far right — mark only its most redundant samples.

   • sugarbeet: a medium band.

   • soil (easiest): a wide band — redundant soil adds the least value, so it is the cheapest to drop.

   With all three filters active at once, the Dashboard selects their intersection, which are the samples that are redundant across every class. Because the weed band is the narrowest, a sample carrying any distinctive weed (or sugarbeet) content is never selected. In fact, the samples considered redundant here are all samples that contain only soil or soil with sugarbeet.

   Once the bands are determined, set the weight of the selected samples to 0 to exclude them from training. Click the select all checkbox and edit one row for all selected rows to sync to weight=0. As a starting point, you can remove 10% of the training set; tune the per-class bands to taste.

Run #4: Train on Reduced Data

Retrain the model on the reduced table — in this example, about ~10% of the training samples were excluded. Overlaying the two runs’ curves, validation mIoU holds essentially flat (0.7892 → 0.7857): the removed samples were redundant, so the model loses little by not seeing them. The result is a smaller, cheaper training set at nearly the same accuracy.

2. Correcting Annotations

Model errors are a strong signal for label problems: where a well-trained model confidently disagrees with the ground truth, the ground truth is often at fault. Two diagnostics point to likely-mislabeled samples:

• many incorrect pixels — right-click the error map column and take its pixel count (the number of mispredicted pixels), then rank by it, and

• a low class IoU — poor agreement on a class the model otherwise handles well.

Sort by one of these to bring the worst samples to the top: hold shift and click the column header to sort by its values. Inspect the image, ground-truth mask, and predicted mask together. When the prediction is clearly right and the ground truth is wrong, adopt the prediction — right-click the mask → Accept selected predictions (or press I on the row) to overwrite the ground truth with the model’s output.

Run #5: Train on Refined Data

Retrain on the corrected table, then overlay this run against the earlier ones in the Dashboard. Comparing the curated and uncurated validation-mIoU curves quantifies the accuracy gained from fixing the labels.

Conclusion

Using only the model’s own diagnostics, we curated the dataset two ways: pruning redundant, low-information samples with the per-class traversal index, and fixing mislabeled ground truth found via the error map’s pixel count and per-class IoU. Both are iterative — after curating, re-run the notebook on the new table revision (the per-class weights recompute automatically) and compare the runs in the Dashboard to confirm the change helped.