{ "cells": [ { "cell_type": "markdown", "id": "0", "metadata": {}, "source": [ "# Fine-tune Mask2Former on Pascal VOC and collect metrics\n", "\n", "Fine-tune a Mask2Former model for semantic segmentation on the Pascal VOC 2012 tables, and collect per-sample predictions and IoU back into a 3LC run.\n", "\n", "![](../images/pascal-voc-mask2former.png)\n", "\n", "\n", "\n", "We fine-tune [Mask2Former](https://huggingface.co/docs/transformers/en/model_doc/mask2former) on the tables\n", "from `create-pascal-voc-semseg-table`. Mask2Former is a *universal* segmentation model — it predicts a\n", "set of binary masks plus a class for each. We use it in **semantic mode**: starting from an ADE20k-semantic\n", "checkpoint (its Swin backbone and decoders already speak segmentation), we swap in a 21-class head for VOC\n", "and fine-tune, then collapse the mask set back to a dense per-pixel label map with the processor's\n", "`post_process_semantic_segmentation`. Only the classification head is reinitialized, so we fine-tune it at\n", "a higher learning rate than the pretrained weights. The void boundary (id 255) is excluded from both the\n", "loss and the metrics. We track train/val loss every epoch and, after training, collect per sample on both\n", "splits: the predicted segmentation (as RLE), mean IoU, the per-image confusion matrix (via the core\n", "metrics helper), and a pooled decoder embedding (reduced to 2D with PaCMAP)." ] }, { "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 - Pascal VOC 2012\"\n", "DATASET_NAME = \"pascal-voc-2012\"\n", "RUN_NAME = \"Fine-tune Mask2Former\"\n", "RUN_DESCRIPTION = \"Mask2Former (semantic mode) fine-tuned on Pascal VOC 2012\"\n", "CHECKPOINT = \"facebook/mask2former-swin-tiny-ade-semantic\"\n", "EPOCHS = 10\n", "BATCH_SIZE = 4\n", "LR = 5e-5 # learning rate for the pretrained weights (backbone + decoders)\n", "HEAD_LR_MULT = 10 # the reinitialized classification head trains at LR * HEAD_LR_MULT\n", "IMAGE_SIZE = 384\n", "NUM_WORKERS = 0 # 0 is safest in notebooks: notebook-defined transforms can't be pickled to\n", "# spawned DataLoader workers. Bump it only when running as an importable module.\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 transformers 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", "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\n", "from transformers import AutoImageProcessor, Mask2FormerForUniversalSegmentation" ] }, { "cell_type": "markdown", "id": "7", "metadata": {}, "source": [ "## Class universe\n", "\n", "VOC has 21 classes (background + 20 objects); `void` (id 255) is the ignore boundary. `background` (id 0)\n", "is not a value-map class — it rides in the column schema's metadata — so both the GT and predicted value\n", "maps show the 20 objects, with the GT map additionally carrying the tagged `void` class. `id2label` /\n", "`label2id` are what we hand to the model's reinitialized classification head (the model still predicts\n", "background as id 0)." ] }, { "cell_type": "code", "execution_count": null, "id": "8", "metadata": {}, "outputs": [], "source": [ "VOC_CLASS_NAMES = [\n", " \"background\",\n", " \"aeroplane\",\n", " \"bicycle\",\n", " \"bird\",\n", " \"boat\",\n", " \"bottle\",\n", " \"bus\",\n", " \"car\",\n", " \"cat\",\n", " \"chair\",\n", " \"cow\",\n", " \"diningtable\",\n", " \"dog\",\n", " \"horse\",\n", " \"motorbike\",\n", " \"person\",\n", " \"pottedplant\",\n", " \"sheep\",\n", " \"sofa\",\n", " \"train\",\n", " \"tvmonitor\",\n", "]\n", "BACKGROUND_ID = 0\n", "VOID_ID = 255\n", "NUM_CLASSES = len(VOC_CLASS_NAMES) # 21\n", "\n", "GT_CLASSES = {**{i: n for i, n in enumerate(VOC_CLASS_NAMES)}, VOID_ID: \"void\"}\n", "PREDICTED_CLASSES = {i: n for i, n in enumerate(VOC_CLASS_NAMES)}\n", "\n", "id2label = {i: n for i, n in enumerate(VOC_CLASS_NAMES)}\n", "label2id = {n: i for i, n in enumerate(VOC_CLASS_NAMES)}" ] }, { "cell_type": "markdown", "id": "9", "metadata": {}, "source": [ "## Load the processor and model\n", "\n", "The processor turns images (and, for training, segmentation maps) into the mask-classification targets\n", "Mask2Former expects, and post-processes predictions back into dense label maps. We set `ignore_index` to\n", "the VOC void id and keep `do_reduce_labels=False` (VOC's background is a real class, not an ignore label).\n", "\n", "The model loads from an ADE20k-semantic checkpoint; `num_labels=21` with `ignore_mismatched_sizes=True`\n", "reinitializes the classification head for VOC while keeping the pretrained backbone and pixel decoder." ] }, { "cell_type": "code", "execution_count": null, "id": "10", "metadata": {}, "outputs": [], "source": [ "processor = AutoImageProcessor.from_pretrained(\n", " CHECKPOINT,\n", " size={\"height\": IMAGE_SIZE, \"width\": IMAGE_SIZE},\n", " ignore_index=VOID_ID,\n", " do_reduce_labels=False,\n", ")\n", "\n", "model = Mask2FormerForUniversalSegmentation.from_pretrained(\n", " CHECKPOINT,\n", " num_labels=NUM_CLASSES,\n", " id2label=id2label,\n", " label2id=label2id,\n", " ignore_mismatched_sizes=True,\n", ")" ] }, { "cell_type": "markdown", "id": "11", "metadata": {}, "source": [ "## Sample transform\n", "\n", "We feed the table to the DataLoader via `Table.with_transform` rather than a custom `Dataset`. The\n", "transform returns a `(PIL image, dense label map)` pair per row; the Mask2Former processor handles\n", "resizing and target construction in the collate function. Curation is handled by the sampler\n", "(`create_sampler` drops zero-weight rows), so Dashboard exclusions take effect on the next run." ] }, { "cell_type": "code", "execution_count": null, "id": "12", "metadata": {}, "outputs": [], "source": [ "def voc_sample_transform(sample: dict) -> tuple:\n", " seg: SemanticSegmentation = sample[\"mask\"]\n", " return sample[\"image\"].convert(\"RGB\"), seg.mask.astype(np.int64)\n", "\n", "\n", "def collate_fn(batch: list) -> dict:\n", " images = [image for image, _ in batch]\n", " masks = [mask for _, mask in batch]\n", " return processor(images=images, segmentation_maps=masks, return_tensors=\"pt\")" ] }, { "cell_type": "markdown", "id": "13", "metadata": {}, "source": [ "## Metrics collection\n", "\n", "For each sample we predict at the original image size, then derive the dense label map with\n", "`post_process_semantic_segmentation`. We store the predicted segmentation (as RLE), the mean IoU and the\n", "per-image confusion matrix (both from the core helper, void excluded), and a pooled embedding from the\n", "transformer decoder's per-query hidden states." ] }, { "cell_type": "code", "execution_count": null, "id": "14", "metadata": {}, "outputs": [], "source": [ "def collect_metrics(run: tlc.Run, table: tlc.Table, model, device: torch.device, epoch: int) -> float:\n", " predictions: list[np.ndarray] = []\n", " ious: list[float] = []\n", " confusion_matrices: list[list[int]] = []\n", " embeddings: list[np.ndarray] = []\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", " inputs = processor(images=image, return_tensors=\"pt\").to(device)\n", " outputs = model(**inputs)\n", "\n", " pred_map = (\n", " processor.post_process_semantic_segmentation(outputs, target_sizes=[(height, width)])[0]\n", " .cpu()\n", " .numpy()\n", " .astype(np.int32)\n", " )\n", "\n", " # Pooled decoder embedding: mean over the object queries -> one vector per image.\n", " decoder_hidden = outputs.transformer_decoder_last_hidden_state # (1, num_queries, hidden)\n", " embeddings.append(decoder_hidden.mean(dim=1).squeeze(0).cpu().numpy().astype(np.float32))\n", "\n", " # Bare label map; the metrics writer serializes it via the predicted_segmentation\n", " # schema below (which records background id 0 in its metadata), so no wrapper is needed.\n", " predictions.append(pred_map)\n", "\n", " m = semantic_segmentation_metrics(\n", " pred_map, seg.mask, GT_CLASSES, background=BACKGROUND_ID, void=VOID_ID, include_background=True\n", " )\n", " ious.append(m[\"mean_iou\"])\n", " confusion_matrices.append([int(x) for cm_row in m[\"confusion_matrix\"] for x in cm_row])\n", "\n", " run.add_metrics(\n", " {\n", " \"predicted_segmentation\": predictions,\n", " \"iou\": ious,\n", " \"confusion_matrix\": confusion_matrices,\n", " \"embedding\": embeddings,\n", " \"epoch\": [epoch] * len(ious),\n", " },\n", " schema={\n", " # Background (id 0) rides in the schema metadata, not the value map.\n", " \"predicted_segmentation\": SemanticSegmentationRleSchema(\n", " classes=PREDICTED_CLASSES, background=BACKGROUND_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 helpers" ] }, { "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)" ] }, { "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", "without code changes." ] }, { "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", "model = model.to(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 DataLoader. create_sampler reads the\n", "# weight column: zero-weight rows (excluded in the Dashboard) are dropped, train is shuffled, val is\n", "# sequential.\n", "train_view = train_table.with_transform(voc_sample_transform)\n", "val_view = val_table.with_transform(voc_sample_transform)\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={\n", " \"checkpoint\": CHECKPOINT,\n", " \"epochs\": EPOCHS,\n", " \"batch_size\": BATCH_SIZE,\n", " \"lr\": LR,\n", " \"image_size\": IMAGE_SIZE,\n", " \"seed\": SEED,\n", " },\n", ")\n", "\n", "train_loader = DataLoader(\n", " train_view,\n", " batch_size=BATCH_SIZE,\n", " sampler=train_sampler,\n", " num_workers=NUM_WORKERS,\n", " collate_fn=collate_fn,\n", ")\n", "val_loader = DataLoader(\n", " val_view,\n", " batch_size=BATCH_SIZE,\n", " sampler=val_sampler,\n", " num_workers=NUM_WORKERS,\n", " collate_fn=collate_fn,\n", ")" ] }, { "cell_type": "markdown", "id": "19", "metadata": {}, "source": [ "## Train\n", "\n", "Two parameter groups give a **discriminative learning rate**: the pretrained backbone and decoders update\n", "gently at `LR`, while the freshly-initialized classification head trains at `LR * HEAD_LR_MULT` so it can\n", "catch up to weights that already learned good segmentation features on ADE20k. Each epoch logs train and\n", "val loss; the full per-sample metrics collection runs once after the loop." ] }, { "cell_type": "code", "execution_count": null, "id": "20", "metadata": {}, "outputs": [], "source": [ "# Discriminative LR: the reinitialized class head learns faster than the pretrained backbone/decoders.\n", "head_params = list(model.class_predictor.parameters())\n", "head_param_ids = {id(p) for p in head_params}\n", "pretrained_params = [p for p in model.parameters() if id(p) not in head_param_ids]\n", "optimizer = torch.optim.AdamW(\n", " [\n", " {\"params\": pretrained_params, \"lr\": LR}, # pretrained backbone + pixel/transformer decoders\n", " {\"params\": head_params, \"lr\": LR * HEAD_LR_MULT}, # freshly initialized for VOC's 21 classes\n", " ]\n", ")\n", "\n", "for epoch in range(EPOCHS):\n", " model.train()\n", " train_loss = 0.0\n", " n_batches = 0\n", " for batch in tqdm(train_loader, desc=f\"epoch {epoch}\", leave=False):\n", " optimizer.zero_grad()\n", " outputs = model(\n", " pixel_values=batch[\"pixel_values\"].to(device),\n", " mask_labels=[m.to(device) for m in batch[\"mask_labels\"]],\n", " class_labels=[c.to(device) for c in batch[\"class_labels\"]],\n", " )\n", " loss = outputs.loss\n", " loss.backward()\n", " optimizer.step()\n", " train_loss += loss.item()\n", " n_batches += 1\n", " train_loss /= max(n_batches, 1)\n", "\n", " # Cheap val-loss pass every epoch (same loss the model trains on); the expensive per-sample metrics\n", " # collection runs once after training.\n", " model.eval()\n", " val_loss = 0.0\n", " n_val_batches = 0\n", " with torch.no_grad():\n", " for batch in val_loader:\n", " outputs = model(\n", " pixel_values=batch[\"pixel_values\"].to(device),\n", " mask_labels=[m.to(device) for m in batch[\"mask_labels\"]],\n", " class_labels=[c.to(device) for c in batch[\"class_labels\"]],\n", " )\n", " val_loss += outputs.loss.item()\n", " n_val_batches += 1\n", " val_loss /= max(n_val_batches, 1)\n", "\n", " log_entry = {\"epoch\": epoch, \"train_loss\": train_loss, \"val_loss\": val_loss}\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": [ "## Collect metrics and reduce embeddings\n", "\n", "With training finished, we run the full per-sample collection once on both the val and train splits, then\n", "fit one PaCMAP model and apply it across every metrics table so they share a single 2D space." ] }, { "cell_type": "code", "execution_count": null, "id": "22", "metadata": {}, "outputs": [], "source": [ "print(\"Collecting per-sample metrics on val and train...\")\n", "collect_metrics(run, val_table, model, device, epoch=EPOCHS - 1)\n", "collect_metrics(run, train_table, model, device, epoch=EPOCHS - 1)\n", "\n", "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 overlay predictions on the images,\n", "sort by IoU to find the hardest classes and images, and use the per-image confusion matrix to see which\n", "VOC classes get confused. Exclude or relabel samples in the Dashboard, then re-run — `.latest()` picks up\n", "the curated revision automatically." ] } ], "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 }