{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Create a Semantic Segmentation Table from Oxford-IIIT Pets\n",
    "\n",
    "Ingest the Oxford-IIIT Pets dataset into 3LC tables for semantic segmentation. Each image has a trimap giving three classes — pet, background, and border — stored as run-length-encoded layers.\n",
    "\n",
    "![](../../images/oxford-pets-semseg.png)\n",
    "\n",
    "<!-- Tags: [\"semantic-segmentation\", \"oxford-pets\"] -->\n",
    "\n",
    "Semantic segmentation assigns every pixel a class label, producing a dense `(H, W)` map that exhaustively\n",
    "partitions the image. We ingest [Oxford-IIIT Pets](https://www.robots.ox.ac.uk/~vgg/data/pets/) through the\n",
    "documented front door `Table.from_semantic_segmentation`, which takes the images and the raw integer masks\n",
    "directly and writes an `image` column plus an RLE-backed `mask` column. `background` is implicit — omitted\n",
    "from both per-row storage and the value map (its id rides in the column schema's metadata) and recovered\n",
    "on read — keeping storage compact. We then attach the per-image `species` and `breed` labels with\n",
    "`Table.add_column`.\n",
    "\n",
    "This is the *ingest* half of a pair: the companion notebook `pytorch-oxford-pets-unet-training` trains\n",
    "a UNet on the tables created here and collects per-sample metrics back into 3LC."
   ]
  },
  {
   "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",
    "TABLE_NAME = \"train\"\n",
    "VAL_TABLE_NAME = \"val\"\n",
    "DOWNLOAD_PATH = \"../../transient_data\"  # dataset root; point at an existing copy to skip the download\n",
    "VAL_FRACTION = 0.1  # held-out val slice of the official trainval split (stratified by breed); test stays out\n",
    "N_TRAIN = None  # cap the train rows for a quick run; None keeps the whole train split\n",
    "N_VAL = None  # cap the val rows for a quick run; None keeps the whole val split\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 torchvision matplotlib"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "from collections.abc import Sequence\n",
    "from pathlib import Path\n",
    "\n",
    "import numpy as np\n",
    "import tlc\n",
    "from PIL import Image\n",
    "from torchvision.datasets import OxfordIIITPet"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7",
   "metadata": {},
   "source": [
    "## Download the dataset\n",
    "\n",
    "`torchvision` handles the download — one call fetches the whole dataset, including both the `trainval` and\n",
    "`test` annotation lists. We only need it to fetch the files; afterwards we read the images and trimaps\n",
    "straight from disk so the `Table` references the original image files by URL (rather than re-encoding them).\n",
    "\n",
    "The download only runs when the data is missing: if `DOWNLOAD_PATH/oxford-iiit-pet/` already exists,\n",
    "torchvision uses it as-is and never writes (so a pre-seeded read-only cache is safe). To reuse an existing\n",
    "copy, point `DOWNLOAD_PATH` at its parent — e.g. `DOWNLOAD_PATH = \"~/data\"` when the dataset lives at\n",
    "`~/data/oxford-iiit-pet`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "download_root = Path(DOWNLOAD_PATH).expanduser()\n",
    "data_root = download_root / \"oxford-iiit-pet\"\n",
    "# Only download when missing — torchvision needs both images/ and annotations/ to consider it present.\n",
    "have_data = (data_root / \"images\").is_dir() and (data_root / \"annotations\").is_dir()\n",
    "print(\n",
    "    f\"Dataset cache {'HIT' if have_data else 'MISS'} at {data_root.resolve()} \"\n",
    "    f\"({'using existing data' if have_data else 'downloading'})\"\n",
    ")\n",
    "OxfordIIITPet(root=download_root, split=\"trainval\", target_types=\"segmentation\", download=not have_data)\n",
    "print(\"Dataset at\", data_root.resolve())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9",
   "metadata": {},
   "source": [
    "## Define the class universe\n",
    "\n",
    "The trimap stores three pixel values, and we keep them **exactly as they appear in the source masks** —\n",
    "`1 = pet`, `2 = background`, `3 = border` — no remapping. Two classes get special roles, declared via the\n",
    "`background` / `void` arguments to the front door:\n",
    "\n",
    "- **background** (`id 2`) — not a class in the value map. Its id is recorded in the column schema's\n",
    "  metadata, and the background is omitted from per-row storage (the implicit fill) and recovered on read.\n",
    "  Rendered transparent in the Dashboard.\n",
    "- **border** (`id 3`) — tagged **void** (a \"don't care\" ring around each pet), so it can be excluded\n",
    "  from the loss and from metrics like mean IoU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The trimap PNG already holds the class ids: 1 = pet, 2 = background, 3 = border. We store them as-is.\n",
    "SEGMENTATION_CLASSES = {1: \"pet\", 2: \"background\", 3: \"border\"}\n",
    "BACKGROUND_ID = 2\n",
    "VOID_ID = 3\n",
    "SPECIES_CLASSES = [\"cat\", \"dog\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11",
   "metadata": {},
   "source": [
    "## Parse the annotations and split\n",
    "\n",
    "Oxford-IIIT Pets ships an official partition: `trainval.txt` (the train pool) and `test.txt` (the held-out\n",
    "benchmark set). There is no canonical sub-split of `trainval`, so we make our own — parse `trainval` and\n",
    "carve off a `VAL_FRACTION` validation slice, **stratified by breed** so all 37 breeds are represented in\n",
    "both tables. `test.txt` is deliberately left untouched as the final benchmark. Each list line is\n",
    "`IMAGE CLASS-ID SPECIES BREED-ID`, giving breed (37 classes) and species (cat/dog) and locating the trimap."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "def parse_annotations(root: Path, split: str) -> list[dict]:\n",
    "    # Parse annotations/<split>.txt (split is \"trainval\" or \"test\") into per-sample dicts.\n",
    "    samples = []\n",
    "    for line in (root / \"annotations\" / f\"{split}.txt\").read_text().splitlines():\n",
    "        if line.startswith(\"#\") or not line.strip():\n",
    "            continue\n",
    "        name, class_id, species, _breed_id = line.split()\n",
    "        image_path = root / \"images\" / f\"{name}.jpg\"\n",
    "        trimap_path = root / \"annotations\" / \"trimaps\" / f\"{name}.png\"\n",
    "        if not image_path.exists() or not trimap_path.exists():\n",
    "            continue\n",
    "        samples.append(\n",
    "            {\n",
    "                \"name\": name,\n",
    "                \"image_path\": image_path.resolve(),\n",
    "                \"trimap_path\": trimap_path.resolve(),\n",
    "                \"breed\": int(class_id) - 1,  # 0-based\n",
    "                \"species\": int(species) - 1,  # 0: cat, 1: dog\n",
    "            }\n",
    "        )\n",
    "    return samples\n",
    "\n",
    "\n",
    "def breed_names(samples: list[dict]) -> list[str]:\n",
    "    # Derive 0-indexed breed names from the image file names.\n",
    "    names: dict[int, str] = {}\n",
    "    for s in samples:\n",
    "        names[s[\"breed\"]] = s[\"name\"].rsplit(\"_\", 1)[0].replace(\"_\", \" \").lower()\n",
    "    return [names[i] for i in range(max(names) + 1)]\n",
    "\n",
    "\n",
    "def stratified_split(samples: list[dict], val_fraction: float, seed: int) -> tuple[list[dict], list[dict]]:\n",
    "    # Deterministic per-breed split so train and val share the same 37-breed distribution.\n",
    "    by_breed: dict[int, list[dict]] = {}\n",
    "    for s in samples:\n",
    "        by_breed.setdefault(s[\"breed\"], []).append(s)\n",
    "    rng = random.Random(seed)\n",
    "    train, val = [], []\n",
    "    for breed in sorted(by_breed):\n",
    "        items = by_breed[breed]\n",
    "        rng.shuffle(items)\n",
    "        n_val = max(1, round(len(items) * val_fraction))  # at least one val image per breed\n",
    "        val.extend(items[:n_val])\n",
    "        train.extend(items[n_val:])\n",
    "    return train, val\n",
    "\n",
    "\n",
    "# Split the official trainval pool into train/val (stratified by breed); test.txt stays held out.\n",
    "trainval_samples = parse_annotations(data_root, \"trainval\")\n",
    "train_samples, val_samples = stratified_split(trainval_samples, VAL_FRACTION, SEED)\n",
    "breeds = breed_names(trainval_samples)\n",
    "print(\n",
    "    f\"Split {len(trainval_samples)} trainval images -> {len(train_samples)} train + {len(val_samples)} val \"\n",
    "    f\"across {len(breeds)} breeds (test.txt held out)\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13",
   "metadata": {},
   "source": [
    "## Decode trimaps into dense masks\n",
    "\n",
    "The front door takes the masks as a sequence of dense `(H, W)` integer arrays. `TrimapMasks` decodes each\n",
    "PNG on indexed access — the pixel values are already the class ids — so we never hold every mask in memory\n",
    "at once."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TrimapMasks(Sequence):\n",
    "    def __init__(self, trimap_paths: list[Path]) -> None:\n",
    "        self._trimap_paths = trimap_paths\n",
    "\n",
    "    def __len__(self) -> int:\n",
    "        return len(self._trimap_paths)\n",
    "\n",
    "    def __getitem__(self, index: int) -> np.ndarray:\n",
    "        # The trimap PNG's pixel values are the class ids (1 = pet, 2 = background, 3 = border).\n",
    "        return np.asarray(Image.open(self._trimap_paths[index]), dtype=np.int32)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15",
   "metadata": {},
   "source": [
    "## Write the tables\n",
    "\n",
    "`Table.from_semantic_segmentation` writes the `image` and `mask` columns in one call, recording the\n",
    "`background` id in the schema metadata and tagging the `void` class. We then attach the per-image\n",
    "`species` and `breed` labels with `Table.add_column` (each returns a new revision of the table). The\n",
    "`train` and `val` tables are the breed-stratified split of `trainval` from the previous cell; by default\n",
    "all rows are used, with `N_TRAIN` / `N_VAL` capping each for a quick run.\n",
    "\n",
    "> The front door covers the image + mask case; `add_column` is the general way to attach extra per-row\n",
    "> columns. For very large datasets you may prefer to author all columns in a single `TableWriter` pass."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {},
   "outputs": [],
   "source": [
    "def write_table(rows: list[dict], table_name: str) -> tlc.Table:\n",
    "    table = tlc.Table.from_semantic_segmentation(\n",
    "        images=[s[\"image_path\"] for s in rows],\n",
    "        masks=TrimapMasks([s[\"trimap_path\"] for s in rows]),\n",
    "        classes=SEGMENTATION_CLASSES,\n",
    "        background=BACKGROUND_ID,\n",
    "        void=VOID_ID,\n",
    "        project_name=PROJECT_NAME,\n",
    "        dataset_name=DATASET_NAME,\n",
    "        table_name=table_name,\n",
    "        if_exists=\"overwrite\",\n",
    "    )\n",
    "    table = table.add_column(\n",
    "        \"species\", [s[\"species\"] for s in rows], schema=tlc.schemas.CategoricalLabelSchema(SPECIES_CLASSES)\n",
    "    )\n",
    "    table = table.add_column(\"breed\", [s[\"breed\"] for s in rows], schema=tlc.schemas.CategoricalLabelSchema(breeds))\n",
    "    print(f\"Wrote {table_name}: {len(table)} rows -> {table.url}\")\n",
    "    return table\n",
    "\n",
    "\n",
    "def select(rows: list[dict], n: int | None) -> list[dict]:\n",
    "    # None -> use every row; otherwise take a deterministic random subset of n rows.\n",
    "    if n is None:\n",
    "        return rows\n",
    "    return random.Random(SEED).sample(rows, min(n, len(rows)))\n",
    "\n",
    "\n",
    "train_table = write_table(select(train_samples, N_TRAIN), TABLE_NAME)\n",
    "val_table = write_table(select(val_samples, N_VAL), VAL_TABLE_NAME)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17",
   "metadata": {},
   "source": [
    "## Inspect a sample\n",
    "\n",
    "Reading a row back hands you the dense mask plus the categorical labels. We overlay the mask on the image\n",
    "to sanity-check the alignment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "sample = train_table[0]\n",
    "seg = sample[\"mask\"]\n",
    "print(f\"Classes present: {seg.present_class_ids.tolist()}\")\n",
    "\n",
    "image = sample[\"image\"]\n",
    "mask = seg.mask  # dense (H, W) label map\n",
    "\n",
    "fig, axes = plt.subplots(1, 2, figsize=(9, 4))\n",
    "axes[0].imshow(image)\n",
    "axes[0].set_title(f\"{breeds[sample['breed']]} ({SPECIES_CLASSES[sample['species']]})\")\n",
    "axes[1].imshow(image)\n",
    "axes[1].imshow(mask, alpha=0.5, cmap=\"viridis\", interpolation=\"nearest\")\n",
    "axes[1].set_title(\"background / pet / border\")\n",
    "for ax in axes:\n",
    "    ax.axis(\"off\")\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19",
   "metadata": {},
   "source": [
    "## Next steps\n",
    "\n",
    "The `train` and `val` tables are now in your 3LC project — open the project in the\n",
    "[3LC Dashboard](https://dashboard.3lc.ai) to browse the images and their segmentations.\n",
    "\n",
    "Continue with `pytorch-oxford-pets-unet-training` to train a UNet on these tables and collect\n",
    "per-sample predictions, IoU, and embeddings back into a 3LC `Run`."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.12.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}