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