Prerequisites
This notebook reuses tables created by other example notebooks. Run them first:
Train a UNet on Oxford-IIIT Pets and collect metrics¶
Train a small UNet for semantic segmentation on the Oxford-IIIT Pets tables, and collect rich per-sample metrics — predictions, IoU, loss, entropy, and embeddings — back into a 3LC run.
These are the tables created by create-oxford-pets-semseg-table. Their ground truth carries three classes — background, pet, and border — but border is an ignore region, not a prediction target. The model outputs only background/pet; border pixels are excluded from the loss (ignore_index) and from IoU. Every few epochs we collect, per sample: the predicted segmentation (stored as RLE via the semantic_segmentation sample type), mean and per-class IoU,
cross-entropy loss, prediction entropy, and a pooled bottleneck embedding (reduced to 2D with PaCMAP after training).
Project setup¶
[ ]:
# Parameters
PROJECT_NAME = "3LC Tutorials - Oxford-IIIT Pets"
DATASET_NAME = "oxford-iiit-pets"
RUN_NAME = "Train TinyUNet"
RUN_DESCRIPTION = "TinyUNet semantic segmentation on Oxford-IIIT Pets"
EPOCHS = 40
BATCH_SIZE = 32
LR = 1e-3
IMAGE_SIZE = 128
COLLECT_FREQUENCY = 10 # collect per-sample metrics every N epochs (and on the final epoch)
NUM_WORKERS = 0 # safest in notebooks: a notebook-defined transform can't be pickled to
# spawned DataLoader workers. Bump it only when running as an importable module.
AUGMENT = True
SEED = 42
INSTALL_DEPENDENCIES = True
Install dependencies¶
[ ]:
if INSTALL_DEPENDENCIES:
%pip install -q 3lc torch torchvision tqdm pacmap
Imports¶
[ ]:
import random
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 PIL import Image
from tlc.data_types import SemanticSegmentation
from tlc.integration.torch.samplers import create_sampler
from tlc.metrics.semantic_segmentation import semantic_segmentation_metrics
from tlc.schemas import SemanticSegmentationRleSchema
from torch.utils.data import DataLoader
from tqdm.auto import tqdm
Class universe¶
The GT tables carry the Oxford trimap classes verbatim — pet = 1, background = 2, border = 3. background is not a value-map class: it rides in the column schema’s metadata (rendered as the implicit fill), so the GT value map shows pet + border. border is the void/ignore class. The model predicts pet and background, so the predicted value map shows only pet.
Those GT ids are not a contiguous 0-based range, so — as in any segmentation training loop — we map GT class ids to model output indices (and back). This mapping lives here, in the training code, not in the ground-truth table.
[ ]:
# GT classes for the metrics helper, as a plain id -> name map. Background (id 2) and void
# (id 3, the border/ignore region) are passed to the helper as ids below: background is folded
# into the matrix, void is excluded.
BACKGROUND_CLASS_ID = 2
VOID_CLASS_ID = 3
FOREGROUND_CLASS_ID = 1 # "pet" — the GT class id that actually matters here
GT_CLASSES = {1: "pet", 2: "background", 3: "border"}
# The classes the model outputs, in GT-id space: pet + background (border/void is never predicted).
# These are an explicit list — background is a real output channel even though it is not a value-map class.
MODEL_CLASS_IDS = [FOREGROUND_CLASS_ID, BACKGROUND_CLASS_ID] # [1, 2]
ID_TO_INDEX = {cid: i for i, cid in enumerate(MODEL_CLASS_IDS)} # {1: 0, 2: 1}
INDEX_TO_ID = {i: cid for cid, i in ID_TO_INDEX.items()} # {0: 1, 1: 2}
NUM_CLASSES = len(MODEL_CLASS_IDS) # 2: pet, background
IGNORE_INDEX = 255 # the value torch's CrossEntropyLoss ignores (border maps here)
def gt_to_model_indices(gt_mask: np.ndarray) -> np.ndarray:
# GT class ids -> 0-based model indices; void/border -> IGNORE_INDEX.
out = np.full(gt_mask.shape, IGNORE_INDEX, dtype=np.int64)
for class_id, index in ID_TO_INDEX.items():
out[gt_mask == class_id] = index
return out
def model_indices_to_ids(index_map: np.ndarray) -> np.ndarray:
# 0-based model indices -> GT class ids (for metrics and storage).
out = np.empty(index_map.shape, dtype=np.int32)
for index, class_id in INDEX_TO_ID.items():
out[index_map == index] = class_id
return out
The model — a tiny 3-level UNet¶
[ ]:
class DoubleConv(nn.Module):
def __init__(self, in_ch: int, out_ch: int) -> None:
super().__init__()
self.block = nn.Sequential(
nn.Conv2d(in_ch, out_ch, 3, padding=1, bias=False),
nn.BatchNorm2d(out_ch),
nn.ReLU(inplace=True),
nn.Conv2d(out_ch, out_ch, 3, padding=1, bias=False),
nn.BatchNorm2d(out_ch),
nn.ReLU(inplace=True),
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.block(x)
class TinyUNet(nn.Module):
# A small 3-level UNet.
def __init__(self, num_classes: int, base: int = 16) -> None:
super().__init__()
self.enc1 = DoubleConv(3, base)
self.enc2 = DoubleConv(base, base * 2)
self.enc3 = DoubleConv(base * 2, base * 4)
self.bottleneck = DoubleConv(base * 4, base * 8)
self.pool = nn.MaxPool2d(2)
self.up3 = nn.ConvTranspose2d(base * 8, base * 4, 2, stride=2)
self.dec3 = DoubleConv(base * 8, base * 4)
self.up2 = nn.ConvTranspose2d(base * 4, base * 2, 2, stride=2)
self.dec2 = DoubleConv(base * 4, base * 2)
self.up1 = nn.ConvTranspose2d(base * 2, base, 2, stride=2)
self.dec1 = DoubleConv(base * 2, base)
self.head = nn.Conv2d(base, num_classes, 1)
def forward(self, x: torch.Tensor) -> torch.Tensor:
e1 = self.enc1(x)
e2 = self.enc2(self.pool(e1))
e3 = self.enc3(self.pool(e2))
b = self.bottleneck(self.pool(e3))
d3 = self.dec3(torch.cat([self.up3(b), e3], dim=1))
d2 = self.dec2(torch.cat([self.up2(d3), e2], dim=1))
d1 = self.dec1(torch.cat([self.up1(d2), e1], dim=1))
return self.head(d1)
Sample transform¶
Rather than wrap the table in a custom Dataset, we hand Table.with_transform a single callable that turns one sample into an (image_tensor, label_tensor) pair — augment is just a flag on it. It resizes image and label to IMAGE_SIZE and maps GT class ids to model indices; with augment=True it adds a random horizontal flip and a small rotation/scale (applied jointly to image and label — bilinear for the image, nearest for the label, out-of-frame pixels → IGNORE_INDEX)
plus image-only color jitter.
Curation is handled separately by the sampler (next cell): create_sampler drops zero-weight rows, so samples you exclude in the Dashboard leave training with no code changes.
[ ]:
class SemSegSampleTransform:
# A table sample -> (image_tensor, label_tensor); augment is a flag.
def __init__(self, image_size: int, *, augment: bool = False) -> None:
self.image_size = image_size
self.augment = augment
def __call__(self, sample: dict) -> tuple[torch.Tensor, torch.Tensor]:
image = sample["image"].convert("RGB").resize((self.image_size, self.image_size))
seg: SemanticSegmentation = sample["mask"]
label_map = gt_to_model_indices(seg.mask) # GT ids -> model indices; border -> IGNORE_INDEX
label = Image.fromarray(label_map.astype(np.uint8), mode="L").resize(
(self.image_size, self.image_size), Image.NEAREST
)
if self.augment:
image, label = self._augment(image, label)
image_tensor = torch.from_numpy(np.asarray(image).copy()).permute(2, 0, 1).float() / 255.0
label_tensor = torch.from_numpy(np.asarray(label).copy()).long()
return image_tensor, label_tensor
def _augment(self, image: Image.Image, label: Image.Image) -> tuple[Image.Image, Image.Image]:
if torch.rand(1).item() < 0.5:
image, label = TF.hflip(image), TF.hflip(label)
angle = torch.empty(1).uniform_(-15, 15).item()
scale = torch.empty(1).uniform_(0.9, 1.1).item()
image = TF.affine(
image,
angle=angle,
translate=(0, 0),
scale=scale,
shear=0,
interpolation=TF.InterpolationMode.BILINEAR,
fill=0,
)
label = TF.affine(
label,
angle=angle,
translate=(0, 0),
scale=scale,
shear=0,
interpolation=TF.InterpolationMode.NEAREST,
fill=IGNORE_INDEX,
)
image = TF.adjust_brightness(image, torch.empty(1).uniform_(0.8, 1.2).item())
image = TF.adjust_contrast(image, torch.empty(1).uniform_(0.8, 1.2).item())
image = TF.adjust_saturation(image, torch.empty(1).uniform_(0.8, 1.2).item())
return image, label
Metrics collection¶
For each sample we predict at the original image size and write per-sample metrics to the run: the predicted segmentation (as RLE), mean IoU and pet IoU, the cross-entropy loss and mean prediction entropy (border excluded), and a pooled bottleneck embedding.
loss/entropy are deliberately not proportional to IoU — they surface confidently-wrong and uncertain samples that hard-label IoU misses. The IoU readouts come from the core helper semantic_segmentation_metrics.
[ ]:
def collect_metrics(
run: tlc.Run,
table: tlc.Table,
model: nn.Module,
device: torch.device,
image_size: int,
epoch: int,
) -> float:
predictions: list[np.ndarray] = []
ious: list[float] = []
pet_ious: list[float] = []
losses: list[float] = []
entropies: list[float] = []
embeddings: list[np.ndarray] = []
# Tap the bottleneck activations; one pooled vector per sample becomes the embedding.
captured: dict[str, torch.Tensor] = {}
handle = model.bottleneck.register_forward_hook(lambda _m, _i, out: captured.__setitem__("emb", out))
model.eval()
with torch.no_grad():
for idx in tqdm(range(len(table)), desc="collect", leave=False):
row = table[idx]
image = row["image"].convert("RGB")
seg: SemanticSegmentation = row["mask"]
width, height = image.size
image_tensor = (
torch.from_numpy(np.array(image.resize((image_size, image_size)))).permute(2, 0, 1).float() / 255.0
)
logits = model(image_tensor[None].to(device))
embeddings.append(captured["emb"].mean(dim=(2, 3)).squeeze(0).cpu().numpy().astype(np.float32))
logits = F.interpolate(logits, size=(height, width), mode="bilinear", align_corners=False)
pred_index_map = logits.argmax(dim=1).squeeze(0).cpu().numpy()
pred_map = model_indices_to_ids(pred_index_map) # back to GT class ids for metrics + storage
target = gt_to_model_indices(seg.mask) # GT ids -> model indices; border -> IGNORE_INDEX
target_tensor = torch.from_numpy(target).long()[None].to(device)
losses.append(F.cross_entropy(logits, target_tensor, ignore_index=IGNORE_INDEX).item())
probs = logits.softmax(dim=1)
per_pixel_entropy = -(probs * probs.clamp_min(1e-12).log()).sum(dim=1) # (1, H, W)
valid = target_tensor != IGNORE_INDEX # exclude the border/void ring, as loss and IoU do
entropies.append(float(per_pixel_entropy[valid].mean()) if valid.any() else float("nan"))
# The prediction is a bare (H, W) label map in GT-class-id space. The metrics writer
# serializes it via the predicted_segmentation schema below (which records the background
# in its metadata), so no SemanticSegmentation wrapper is needed.
predictions.append(pred_map)
m = semantic_segmentation_metrics(
pred_map,
seg.mask,
GT_CLASSES,
background=BACKGROUND_CLASS_ID,
void=VOID_CLASS_ID,
include_background=True,
)
ious.append(m["mean_iou"])
pet_ious.append(m["per_class_iou"][m["class_ids"].index(FOREGROUND_CLASS_ID)])
handle.remove()
run.add_metrics(
{
"predicted_segmentation": predictions,
"iou": ious,
"pet_iou": pet_ious,
"loss": losses,
"entropy": entropies,
"embedding": embeddings,
"epoch": [epoch] * len(ious),
},
schema={
# Predicted value map is pet only; background rides in the schema metadata (not the map).
"predicted_segmentation": SemanticSegmentationRleSchema(
classes={1: "pet", 2: "background"}, background=BACKGROUND_CLASS_ID
),
"embedding": tlc.schemas.EmbeddingSchema(shape=len(embeddings[0])),
},
foreign_table_url=table.url,
)
# Average only over images with a defined IoU; a degenerate image yields nan, and a nan headline
# would be logged into the run object as an invalid JSON token.
finite = [v for v in ious if np.isfinite(v)]
return float(np.mean(finite)) if finite else float("nan")
Reproducibility helper¶
[ ]:
def seed_everything(seed: int) -> None:
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True
Load the tables and initialize the Run¶
.latest() picks up the newest revision of each table, so retraining consumes any Dashboard curation (excluded/relabeled samples) without code changes. The val table is the fixed ruler.
[ ]:
seed_everything(SEED)
device = torch.device("cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu")
print(f"Using device: {device}")
train_table = tlc.Table.from_names(table_name="train", dataset_name=DATASET_NAME, project_name=PROJECT_NAME).latest()
val_table = tlc.Table.from_names(table_name="val", dataset_name=DATASET_NAME, project_name=PROJECT_NAME).latest()
print(f"Using train table {train_table.url}")
print(f"Using val table {val_table.url}")
# with_transform returns a map-style view — pass it straight to a
# DataLoader. create_sampler reads the weight column: zero-weight rows are
# dropped, train is shuffled, val is sequential.
train_view = train_table.with_transform(SemSegSampleTransform(IMAGE_SIZE, augment=AUGMENT))
val_view = val_table.with_transform(SemSegSampleTransform(IMAGE_SIZE))
train_sampler = create_sampler(train_table, weighted=False, shuffle=True)
val_sampler = create_sampler(val_table, weighted=False, shuffle=False)
print(f"Train: {len(train_sampler)} of {len(train_table)} rows after weight filtering | Val: {len(val_sampler)}")
run = tlc.init(
PROJECT_NAME,
run_name=RUN_NAME,
description=RUN_DESCRIPTION,
parameters={"epochs": EPOCHS, "batch_size": BATCH_SIZE, "lr": LR, "image_size": IMAGE_SIZE, "seed": SEED},
)
train_loader = DataLoader(train_view, batch_size=BATCH_SIZE, sampler=train_sampler, num_workers=NUM_WORKERS)
val_loader = DataLoader(val_view, batch_size=BATCH_SIZE, sampler=val_sampler, num_workers=NUM_WORKERS)
Train¶
[ ]:
model = TinyUNet(NUM_CLASSES).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=LR)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=EPOCHS)
criterion = nn.CrossEntropyLoss(ignore_index=IGNORE_INDEX)
for epoch in range(EPOCHS):
model.train()
train_loss = 0.0
for images, labels in tqdm(train_loader, desc=f"epoch {epoch}", leave=False):
images, labels = images.to(device), labels.to(device)
optimizer.zero_grad()
loss = criterion(model(images), labels)
loss.backward()
optimizer.step()
train_loss += loss.item() * images.shape[0]
train_loss /= len(train_loader.dataset)
model.eval()
val_loss = 0.0
with torch.no_grad():
for images, labels in val_loader:
images, labels = images.to(device), labels.to(device)
val_loss += criterion(model(images), labels).item() * images.shape[0]
val_loss /= len(val_loader.dataset)
log_entry = {"epoch": epoch, "lr": optimizer.param_groups[0]["lr"], "train_loss": train_loss, "val_loss": val_loss}
scheduler.step()
is_final = epoch == EPOCHS - 1
if (epoch + 1) % COLLECT_FREQUENCY == 0 or is_final:
train_miou = collect_metrics(run, train_table, model, device, IMAGE_SIZE, epoch=epoch)
val_miou = collect_metrics(run, val_table, model, device, IMAGE_SIZE, epoch=epoch)
log_entry |= {"train_miou": train_miou, "val_miou": val_miou}
tlc.log(log_entry)
print(" ".join(f"{k}={v:.4f}" if k != "epoch" else f"epoch {v}" for k, v in log_entry.items()))
Reduce embeddings and finish¶
Fit one PaCMAP model on the final-epoch val embeddings and apply it to every metrics table (both splits, all epochs) so they share a single, stable 2D space. delete_source_tables drops the raw high-dim vectors afterwards — only the 2D reduction is kept.
[ ]:
print("Reducing embeddings (PaCMAP)...")
run.reduce_embeddings_by_foreign_table_url(val_table.url, method="pacmap", delete_source_tables=True)
run.set_status_completed()
print(f"Run: {run.url}")
Next steps¶
Open the Run in the 3LC Dashboard to explore predictions overlaid on the images, sort by IoU / loss / entropy to find failure cases, and use the 2D embedding to spot clusters of hard samples. Exclude or relabel samples in the Dashboard, then re-run this notebook — .latest() picks up the curated revision automatically.
For a larger, multi-class example, see the Pascal VOC pair (create-pascal-voc-semseg-table + huggingface-pascal-voc-mask2former-finetuning).