roboflow--rf-detr/docs/cookbooks/fine-tune_segmentation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e6fd4da4",
   "metadata": {},
   "source": [
    "# RF-DETR instance segmentation fine-tuning on Roboflow Universe datasets\n",
    "\n",
    "Select one dataset with `DATASET_KEY`, then run the same fine-tuning, plotting, checkpoint loading, and inference cells.\n",
    "\n",
    "RF-DETR extends bounding-box detection with a sparse segmentation head that predicts per-object binary masks — useful\n",
    "for facade analysis, infrastructure damage inspection, vehicle part assessment, dental imaging, litter mapping, and\n",
    "any task where *what shape* matters as much as *what class*.\n",
    "\n",
    "**Datasets** — set `DATASET_KEY` to one of:\n",
    "\n",
    "| Key | Dataset | Classes |\n",
    "|-----|---------|---------|\n",
    "| `\"building_facade\"` | Building Facade Segmentation | facade, window, street, vegetation, car, fence, ... |\n",
    "| `\"spalling_rebar\"` | Spalling and Exposed Rebar | exposed_rebar, spalling |\n",
    "| `\"car_damage_parts\"` | Car Part Damage Segmentation | bumper, wheel, door, headlight, dent, scratch, ... |\n",
    "| `\"teeth_numbering\"` | Teeth Numbering Segmentation | tooth IDs 11-48 |\n",
    "| `\"taco_trash\"` | TACO Trash | bottle, can, carton, cup, lid, straw, wrapper, ... |\n",
    "\n",
    "Every subsequent cell (fine-tuning, plotting, checkpoint loading, inference) adapts automatically.\n",
    "By the end you will have a fine-tuned model, training curves, and annotated inference images with segmentation masks."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6eaacf22",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "Install `rfdetr` with the `train` and `visual` extras. The `train` extra pulls in the training loop dependencies\n",
    "(PyTorch Lightning, COCO evaluation tools, and data augmentation libraries), while `visual` adds the visualisation\n",
    "helpers used later in the notebook. The `roboflow` package handles dataset download; `pandas` and `seaborn` are needed\n",
    "for the metrics table and optional plot styling. If you hit an `ImportError` after running this cell, the most likely\n",
    "cause is a stale in-memory import — restart the Python runtime once and re-run from the top."
   ]
  },
  {
   "cell_type": "code",
   "id": "64736ca1",
   "metadata": {},
   "source": "!pip install -q \"rfdetr[train,visual]>=1.8.2\" roboflow pandas seaborn",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b4797e9",
   "metadata": {
    "title": "[bash]"
   },
   "outputs": [],
   "source": [
    "\"\"\"Shared RF-DETR instance segmentation fine-tuning demo for several Roboflow COCO segmentation exports.\"\"\"\n",
    "\n",
    "import json\n",
    "import os\n",
    "from pathlib import Path\n",
    "from typing import Any\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import supervision as sv\n",
    "import torch\n",
    "from IPython.display import display\n",
    "from matplotlib import pyplot as plt\n",
    "from matplotlib.figure import Figure\n",
    "from PIL import Image\n",
    "from roboflow import Roboflow\n",
    "\n",
    "from rfdetr import RFDETRSegSmall\n",
    "from rfdetr.config import SegmentationTrainConfig\n",
    "from rfdetr.training import RFDETRDataModule, RFDETRModelModule, build_trainer\n",
    "from rfdetr.training.callbacks.best_model import BestModelCallback\n",
    "from rfdetr.utilities.reproducibility import seed_all\n",
    "from rfdetr.visualize.training import plot_loss_metrics, plot_map_metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99f1f539",
   "metadata": {},
   "outputs": [],
   "source": [
    "PROJECT_ROOT = Path(__file__).resolve().parent if \"__file__\" in globals() else Path.cwd()\n",
    "DATASETS_DIR = PROJECT_ROOT / \"datasets\"\n",
    "\n",
    "# Resolution must be divisible by patch_size × num_windows (24 for RFDETRSegSmall).\n",
    "# Larger resolutions capture fine-grained mask boundaries at the cost of GPU memory and time.\n",
    "DATASETS: dict[str, dict[str, Any]] = {\n",
    "    \"building_facade\": {\n",
    "        \"name\": \"Building Facade Segmentation\",\n",
    "        \"workspace\": \"building-facade\",\n",
    "        \"project\": \"building-facade-segmentation-instance\",\n",
    "        \"version\": 4,\n",
    "        \"source_url\": \"https://universe.roboflow.com/building-facade/building-facade-segmentation-instance\",\n",
    "        \"output_name\": \"seg_building_facade\",\n",
    "        # 432 — structured street scenes with medium-to-large facade regions plus\n",
    "        # smaller windows, cars, vegetation, and infrastructure boundaries.\n",
    "        \"resolution\": 432,\n",
    "    },\n",
    "    \"spalling_rebar\": {\n",
    "        \"name\": \"Spalling and Exposed Rebar\",\n",
    "        \"workspace\": \"uni-of-birmingham\",\n",
    "        \"project\": \"spalling-and-exposed-rebar\",\n",
    "        \"version\": 7,\n",
    "        \"source_url\": \"https://universe.roboflow.com/uni-of-birmingham/spalling-and-exposed-rebar\",\n",
    "        \"output_name\": \"seg_spalling_rebar\",\n",
    "        # 432 — civil-infrastructure damage masks with irregular thin regions and\n",
    "        # exposed-rebar details that stress boundary quality.\n",
    "        \"resolution\": 432,\n",
    "    },\n",
    "    \"car_damage_parts\": {\n",
    "        \"name\": \"Car Part Damage Segmentation\",\n",
    "        \"workspace\": \"car-damaged-detection-e66m0\",\n",
    "        \"project\": \"car-part-detection-with-damage-part\",\n",
    "        \"version\": 1,\n",
    "        \"source_url\": \"https://universe.roboflow.com/car-damaged-detection-e66m0/car-part-detection-with-damage-part\",\n",
    "        \"output_name\": \"seg_car_damage_parts\",\n",
    "        # 432 — fine-grained automotive part masks with many adjacent classes and\n",
    "        # damage regions, useful for checking class-specific mask boundaries.\n",
    "        \"resolution\": 432,\n",
    "    },\n",
    "    \"teeth_numbering\": {\n",
    "        \"name\": \"Teeth Numbering Segmentation\",\n",
    "        \"workspace\": \"godento2\",\n",
    "        \"project\": \"teeth-seg-3537-iaky1\",\n",
    "        \"version\": 1,\n",
    "        \"source_url\": \"https://universe.roboflow.com/godento2/teeth-seg-3537-iaky1\",\n",
    "        \"output_name\": \"seg_teeth_numbering\",\n",
    "        # 432 — dental images with many repeated small instances whose labels are\n",
    "        # tied to position, stressing class confusion and small-mask quality.\n",
    "        \"resolution\": 432,\n",
    "    },\n",
    "    \"taco_trash\": {\n",
    "        \"name\": \"TACO Trash\",\n",
    "        \"workspace\": \"mohamed-traore-2ekkp\",\n",
    "        \"project\": \"taco-trash-annotations-in-context\",\n",
    "        \"version\": 13,\n",
    "        \"source_url\": \"https://universe.roboflow.com/mohamed-traore-2ekkp/taco-trash-annotations-in-context\",\n",
    "        \"output_name\": \"seg_taco_trash\",\n",
    "        # 432 — long-tail litter segmentation with many small, cluttered object\n",
    "        # categories. Use as a harder stress test after the simpler datasets.\n",
    "        \"resolution\": 432,\n",
    "    },\n",
    "}\n",
    "\n",
    "DATASET_KEY = \"building_facade\"\n",
    "DATASET_INFO = DATASETS[DATASET_KEY]\n",
    "\n",
    "OUTPUT_DIR = PROJECT_ROOT / \"output\" / str(DATASET_INFO[\"output_name\"])\n",
    "METRICS_CSV = OUTPUT_DIR / \"metrics.csv\"\n",
    "VALIDATION_METRICS_JSON = OUTPUT_DIR / \"validation_metrics.json\"\n",
    "FINAL_CHECKPOINT_PATH = OUTPUT_DIR / \"checkpoint_final_demo.pth\"\n",
    "\n",
    "RESOLUTION = int(DATASET_INFO[\"resolution\"])\n",
    "\n",
    "SEED = 7\n",
    "EPOCHS = 50\n",
    "BATCH_SIZE = 8\n",
    "GRAD_ACCUM_STEPS = 2\n",
    "NUM_WORKERS = 8\n",
    "LR = 1e-4\n",
    "LR_ENCODER = 1e-4\n",
    "SAMPLE_PREVIEW_COUNT = 6\n",
    "SAMPLE_PREVIEW_COLUMNS = 3\n",
    "SAMPLE_PREVIEW_FIGURE_SIZE: tuple[float, float] = (15.0, 10.0)\n",
    "INFERENCE_COUNT = 6\n",
    "INFERENCE_COLUMNS = 3\n",
    "INFERENCE_THRESHOLD = 0.3\n",
    "PLOT_LOSS_LOG_SCALE = False\n",
    "\n",
    "print(f\"dataset_key={DATASET_KEY}\")\n",
    "print(f\"dataset={DATASET_INFO['name']}\")\n",
    "print(f\"source_url={DATASET_INFO['source_url']}\")\n",
    "print(f\"resolution={RESOLUTION}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b60971f9",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "## Notebook display\n",
    "\n",
    "This helper registers the `%matplotlib inline` magic so figures render directly beneath each cell when you run the\n",
    "notebook in Jupyter or Google Colab. If you are running this file as a plain Python script the function detects the\n",
    "absence of an IPython kernel and exits silently, so it is always safe to call."
   ]
  },
  {
   "cell_type": "code",
   "id": "1c40c67d",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "def _enable_notebook_inline_matplotlib() -> None:\n",
    "    \"\"\"Enable inline matplotlib figures when running in IPython.\"\"\"\n",
    "    get_ipython_func = globals().get(\"get_ipython\")\n",
    "    if not callable(get_ipython_func):\n",
    "        return\n",
    "    ipython = get_ipython_func()\n",
    "    if ipython is not None:\n",
    "        ipython.run_line_magic(\"matplotlib\", \"inline\")\n",
    "        ipython.run_line_magic(\"config\", \"InlineBackend.close_figures = True\")\n",
    "\n",
    "\n",
    "_enable_notebook_inline_matplotlib()"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "7f99c55d",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "## 1 - Download dataset\n",
    "\n",
    "Roboflow exports instance segmentation datasets in COCO format: a JSON file where each annotation contains a\n",
    "`segmentation` field with polygon coordinates that trace the precise object boundary.  The download cell fetches the\n",
    "exact dataset version listed in `DATASETS` and places it under `datasets/<DATASET_KEY>/`. You need a Roboflow API key\n",
    "— get one for free at app.roboflow.com/settings/api and set it as the `ROBOFLOW_API_KEY` environment variable (or as\n",
    "a Colab secret with the same name). The download is idempotent: if the target directory already exists Roboflow skips\n",
    "the network transfer, so re-running this cell after a successful download is fast."
   ]
  },
  {
   "cell_type": "code",
   "id": "ab8fdeb6",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "try:\n",
    "    from google.colab import userdata\n",
    "\n",
    "    try:\n",
    "        ROBOFLOW_API_KEY = userdata.get(\"ROBOFLOW_API_KEY\") or \"\"\n",
    "    except Exception:\n",
    "        ROBOFLOW_API_KEY = \"\"\n",
    "except ImportError:\n",
    "    ROBOFLOW_API_KEY = \"\"\n",
    "\n",
    "if not ROBOFLOW_API_KEY:\n",
    "    ROBOFLOW_API_KEY = os.environ.get(\"ROBOFLOW_API_KEY\", \"\")\n",
    "if not ROBOFLOW_API_KEY:\n",
    "    raise RuntimeError(\n",
    "        \"ROBOFLOW_API_KEY not found. \"\n",
    "        \"In Colab: add it via Secrets (key icon). \"\n",
    "        \"Locally: set the environment variable before running.\"\n",
    "    )\n",
    "\n",
    "rf = Roboflow(api_key=ROBOFLOW_API_KEY)\n",
    "dataset = (\n",
    "    rf.workspace(str(DATASET_INFO[\"workspace\"]))\n",
    "    .project(str(DATASET_INFO[\"project\"]))\n",
    "    .version(int(DATASET_INFO[\"version\"]))\n",
    "    .download(\"coco\", location=str(DATASETS_DIR / DATASET_KEY))\n",
    ")\n",
    "DATASET_DIR = Path(dataset.location)\n",
    "TRAIN_ANNOTATIONS = DATASET_DIR / \"train\" / \"_annotations.coco.json\"\n",
    "print(f\"dataset_dir={DATASET_DIR}\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "e38db150",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "## 2 - Infer class names\n",
    "\n",
    "Read the class names directly from the COCO annotation file.  RF-DETR uses 0-based class indices internally; the\n",
    "`class_names` list maps each index to a human-readable label so the inference visualisation shows category names\n",
    "rather than numbers.  If your dataset has a single category the list will have one element — that is expected."
   ]
  },
  {
   "cell_type": "code",
   "id": "2c6b8dc4",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "with TRAIN_ANNOTATIONS.open() as _f:\n",
    "    _coco = json.load(_f)\n",
    "\n",
    "CLASS_NAMES: list[str] = [cat[\"name\"] for cat in sorted(_coco[\"categories\"], key=lambda c: c[\"id\"])]\n",
    "NUM_CLASSES = len(CLASS_NAMES)\n",
    "\n",
    "print(f\"class_names={CLASS_NAMES}\")\n",
    "print(f\"num_classes={NUM_CLASSES}\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "945822fa",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "`_save_final_checkpoint` writes a self-contained `.pth` file that bundles model weights with the full training and\n",
    "model config. This is separate from the PTL checkpoint because it can be loaded with a single `from_checkpoint` call\n",
    "on any machine, without reconstructing the original config."
   ]
  },
  {
   "cell_type": "code",
   "id": "ed36df66",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "def _save_final_checkpoint(\n",
    "    module: RFDETRModelModule,\n",
    "    trainer: Any,\n",
    "    train_config: SegmentationTrainConfig,\n",
    "    model_config: Any,\n",
    "    output_path: Path,\n",
    ") -> Path:\n",
    "    \"\"\"Save a final RF-DETR .pth checkpoint that can be loaded with ``from_checkpoint``.\"\"\"\n",
    "    raw_model: Any = getattr(module.model, \"_orig_mod\", module.model)\n",
    "    output_path.parent.mkdir(parents=True, exist_ok=True)\n",
    "    torch.save(\n",
    "        BestModelCallback._build_checkpoint_payload(\n",
    "            raw_model.state_dict(),\n",
    "            train_config.model_dump(),\n",
    "            trainer,\n",
    "            model_name=\"RFDETRSegSmall\",\n",
    "            model_config_dict=model_config.model_dump(),\n",
    "        ),\n",
    "        output_path,\n",
    "    )\n",
    "    return output_path"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "7000d9a4",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "`_segmentation_grid_figure` arranges multiple annotated images into a fixed-column matplotlib grid.\n",
    "Masks are drawn with `sv.MaskAnnotator` (translucent colour fill) and bounding boxes with\n",
    "`sv.BoxAnnotator`.  Labels show the class name and confidence score."
   ]
  },
  {
   "cell_type": "code",
   "id": "477cd2f6",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "def _segmentation_grid_figure(\n",
    "    items: list[tuple[str, np.ndarray, sv.Detections]],\n",
    "    columns: int,\n",
    "    class_names: list[str],\n",
    ") -> Figure:\n",
    "    \"\"\"Render segmentation masks and bounding boxes in a fixed-column subplot grid.\"\"\"\n",
    "    if columns <= 0:\n",
    "        raise ValueError(f\"columns must be positive, got {columns}.\")\n",
    "\n",
    "    mask_annotator = sv.MaskAnnotator()\n",
    "    box_annotator = sv.BoxAnnotator(thickness=2)\n",
    "    label_annotator = sv.LabelAnnotator(text_scale=0.5, text_thickness=1, text_padding=4)\n",
    "\n",
    "    rows = max(1, (len(items) + columns - 1) // columns)\n",
    "    figure, axes = plt.subplots(rows, columns, figsize=(5 * columns, 5 * rows))\n",
    "    axes_array = np.asarray(axes, dtype=object).reshape(-1)\n",
    "    for axis in axes_array:\n",
    "        axis.axis(\"off\")\n",
    "\n",
    "    for axis, (title, image, detections) in zip(axes_array, items, strict=False):\n",
    "        scene = image.copy()\n",
    "        if detections.mask is not None and len(detections) > 0:\n",
    "            scene = mask_annotator.annotate(scene=scene, detections=detections)\n",
    "        scene = box_annotator.annotate(scene=scene, detections=detections)\n",
    "        if len(detections) > 0 and detections.class_id is not None:\n",
    "            conf_list = (\n",
    "                detections.confidence.tolist() if detections.confidence is not None else [None] * len(detections)\n",
    "            )\n",
    "            labels = [\n",
    "                f\"{class_names[cid] if cid < len(class_names) else cid}\" + (f\" {conf:.2f}\" if conf is not None else \"\")\n",
    "                for cid, conf in zip(detections.class_id.tolist(), conf_list)\n",
    "            ]\n",
    "            scene = label_annotator.annotate(scene=scene, detections=detections, labels=labels)\n",
    "        axis.imshow(scene)\n",
    "        axis.set_title(title, fontsize=10)\n",
    "        axis.axis(\"off\")\n",
    "\n",
    "    figure.tight_layout()\n",
    "    return figure"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "8e53a5a5",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "## 3 - Configure training\n",
    "\n",
    "`SegmentationTrainConfig` centralises every hyperparameter the training loop needs, including mask-specific loss\n",
    "coefficients (`mask_ce_loss_coef` and `mask_dice_loss_coef`).  The defaults (5.0 each) work well for most datasets;\n",
    "if the segmentation head overfits early you can reduce them towards 2.0 while keeping `cls_loss_coef=5.0`.\n",
    "\n",
    "`lr` and `lr_encoder` control the head and backbone learning rates respectively.  For fine-tuning both are usually\n",
    "set to the same value; if you notice the backbone overfitting early, halve `lr_encoder` to protect the pre-trained\n",
    "features.  `grad_accum_steps=2` doubles the effective batch size without extra GPU memory.\n",
    "\n",
    "`multi_scale=False` and `expanded_scales=False` disable multi-resolution augmentation to shorten epoch time during\n",
    "fine-tuning — they rarely help on small custom datasets.  The `notes` dict is stored alongside the weights in the\n",
    "checkpoint, giving you a lightweight experiment log that travels with the model file."
   ]
  },
  {
   "cell_type": "code",
   "id": "3e71bc58",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "seed_all(SEED)\n",
    "variant = RFDETRSegSmall(  # type: ignore[no-untyped-call]\n",
    "    num_classes=NUM_CLASSES,\n",
    "    resolution=RESOLUTION,\n",
    ")\n",
    "variant.model_config.model_name = type(variant).__name__\n",
    "\n",
    "train_config = SegmentationTrainConfig(\n",
    "    dataset_file=\"roboflow\",\n",
    "    dataset_dir=str(DATASET_DIR),\n",
    "    output_dir=str(OUTPUT_DIR),\n",
    "    epochs=EPOCHS,\n",
    "    batch_size=BATCH_SIZE,\n",
    "    grad_accum_steps=GRAD_ACCUM_STEPS,\n",
    "    num_workers=NUM_WORKERS,\n",
    "    lr=LR,\n",
    "    lr_encoder=LR_ENCODER,\n",
    "    use_ema=False,\n",
    "    run_test=False,\n",
    "    compute_train_metrics=True,\n",
    "    compute_val_loss=True,\n",
    "    multi_scale=False,\n",
    "    expanded_scales=False,\n",
    "    do_random_resize_via_padding=False,\n",
    "    tensorboard=False,\n",
    "    wandb=False,\n",
    "    mlflow=False,\n",
    "    clearml=False,\n",
    "    class_names=CLASS_NAMES,\n",
    "    notes={\n",
    "        \"demo\": f\"segmentation PTL fine-tune on Roboflow Universe {DATASET_INFO['project']}\",\n",
    "        \"source_url\": DATASET_INFO[\"source_url\"],\n",
    "        \"roboflow_workspace\": DATASET_INFO[\"workspace\"],\n",
    "        \"roboflow_project\": DATASET_INFO[\"project\"],\n",
    "        \"roboflow_version\": DATASET_INFO[\"version\"],\n",
    "        \"num_classes\": NUM_CLASSES,\n",
    "        \"class_names\": CLASS_NAMES,\n",
    "    },\n",
    "    progress_bar=\"tqdm\",\n",
    ")\n",
    "\n",
    "datamodule = RFDETRDataModule(variant.model_config, train_config)\n",
    "model = RFDETRModelModule(variant.model_config, train_config)\n",
    "trainer = build_trainer(train_config, variant.model_config)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "5f0ccbd3",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "## 4 - Preview dataset inputs\n",
    "\n",
    "Always look at your data before you start a long training run. This cell renders a grid of annotated training images\n",
    "so you can confirm that segmentation masks align with objects, colour jitter and flips look reasonable, and there are\n",
    "no systematic labelling errors such as masks placed on the wrong object or masks that bleed outside the object\n",
    "boundary. Catching label noise here costs a few seconds; catching it after 50 epochs costs much more."
   ]
  },
  {
   "cell_type": "code",
   "id": "71e1a7db",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "sample_figure = datamodule._show_samples(\n",
    "    SAMPLE_PREVIEW_COUNT,\n",
    "    split=\"train\",\n",
    "    columns=SAMPLE_PREVIEW_COLUMNS,\n",
    "    figure_size=SAMPLE_PREVIEW_FIGURE_SIZE,\n",
    ")\n",
    "display(sample_figure)\n",
    "plt.close(sample_figure)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "993b1cd1",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "## 5 - Fine-tune\n",
    "\n",
    "`trainer.fit` hands control to PyTorch Lightning, which drives the full training loop: forward pass, loss\n",
    "computation (cross-entropy + Dice on sampled mask points), gradient accumulation, weight updates,\n",
    "learning-rate scheduling, and periodic validation with COCO mask mAP.  After each epoch the trainer appends a row to\n",
    "`metrics.csv` and, if validation mAP improves, saves a new best checkpoint via `BestModelCallback`. On a single\n",
    "consumer GPU (RTX 3090 or similar) 50 epochs on the crack dataset takes roughly 20–35 minutes depending on\n",
    "`BATCH_SIZE` and `NUM_WORKERS`."
   ]
  },
  {
   "cell_type": "code",
   "id": "0564b2c6",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "trainer.fit(model, datamodule=datamodule, ckpt_path=train_config.resume or None)\n",
    "print(f\"output_dir={OUTPUT_DIR}\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "c746e44d",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "## 6 - Save checkpoint/model\n",
    "\n",
    "PTL saves its own checkpoints during training (optimizer state, scheduler state, epoch counter), but those files are\n",
    "not directly portable — they require the same class hierarchy to load.  The `.pth` file written here is a\n",
    "self-contained RF-DETR checkpoint: it bundles the model weights together with the full training and model configs,\n",
    "including the class names.  You can share the file with a colleague and they can run inference with a single\n",
    "`RFDETRSegSmall.from_checkpoint(path)` call.  For full reproducibility, keep this checkpoint alongside the dataset\n",
    "version number printed in section 1."
   ]
  },
  {
   "cell_type": "code",
   "id": "684059b3",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "final_checkpoint = _save_final_checkpoint(model, trainer, train_config, variant.model_config, FINAL_CHECKPOINT_PATH)\n",
    "print(f\"saved_checkpoint={final_checkpoint}\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "dadba8e3",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "## 7 - Validate metrics\n",
    "\n",
    "This cell runs a clean post-training validation with no augmentation, using the best checkpoint loaded from\n",
    "`checkpoint_best_total.pth` written by `BestModelCallback`, and serialises results to JSON. Two metrics are most\n",
    "important to inspect. `bbox/map` is the standard COCO bounding-box mAP at IoU 0.50:0.95. `segm/map` (when logged)\n",
    "is the mask mAP — it measures how precisely the predicted binary masks overlap with the annotated polygons. A model\n",
    "can have high `bbox/map` but low `segm/map` if the detector locates objects well but the mask boundaries are coarse;\n",
    "this usually improves with more training epochs or by increasing `resolution`.\n",
    "\n",
    "**Note:** `checkpoint_best_total.pth` is written after the first completed validation epoch. If training was\n",
    "interrupted before any epoch finished, this file will not exist and the cell below will raise `FileNotFoundError`."
   ]
  },
  {
   "cell_type": "code",
   "id": "ea69ee92",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "# checkpoint_best_total.pth stores plain model weights under key \"model\", not a full\n",
    "# PTL checkpoint.  Passing it as ckpt_path to trainer.validate() triggers a KeyError\n",
    "# on \"validate_loop\" because PTL tries to restore loop state that isn't present.\n",
    "# Load the weights directly and run validate without ckpt_path instead.\n",
    "_ckpt = torch.load(OUTPUT_DIR / \"checkpoint_best_total.pth\", map_location=\"cpu\", weights_only=False)\n",
    "model.model.load_state_dict(_ckpt[\"model\"], strict=True)\n",
    "del _ckpt\n",
    "validation_results = trainer.validate(model, datamodule=datamodule)\n",
    "validation_metrics = {key: float(value) for key, value in validation_results[0].items()} if validation_results else {}\n",
    "VALIDATION_METRICS_JSON.write_text(json.dumps(validation_metrics, indent=2, sort_keys=True), encoding=\"utf-8\")\n",
    "print(f\"validation_metrics={validation_metrics}\")\n",
    "print(f\"validation_metrics_json={VALIDATION_METRICS_JSON}\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "e9b7a1cb",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "## 8 - Plot CSVLogger metrics\n",
    "\n",
    "Reading loss curves is one of the fastest ways to diagnose training problems.  Healthy segmentation runs show the\n",
    "combined detection + mask loss decreasing together for train and val.  A loss that rises or plateaus from epoch 1 is\n",
    "usually a sign that the learning rate is too high or the dataset is too small for the chosen `BATCH_SIZE`.  For the\n",
    "mAP curves, `segm mAP 50` (mask IoU ≥ 0.50) rises fastest and is the clearest signal of whether the model is\n",
    "learning to segment objects; `segm mAP 50:95` (averaged across stricter thresholds) follows more slowly.\n",
    "Set `PLOT_LOSS_LOG_SCALE = True` if the loss drops by an order of magnitude in the first few epochs and the later,\n",
    "more meaningful portion of the curve gets compressed into a flat line."
   ]
  },
  {
   "cell_type": "code",
   "id": "abab244f",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "print(f\"metrics_csv={METRICS_CSV}\")\n",
    "loss_figure = plot_loss_metrics(str(METRICS_CSV), loss_log_scale=PLOT_LOSS_LOG_SCALE)\n",
    "display(loss_figure)\n",
    "plt.close(loss_figure)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c93477c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "map_figure = plot_map_metrics(str(METRICS_CSV))\n",
    "display(map_figure)\n",
    "plt.close(map_figure)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e46c2424",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "## 9 - Load checkpoint/model\n",
    "\n",
    "`from_checkpoint` is the standard entry point for loading a saved RF-DETR model.  It reads both the weights and the\n",
    "stored config from the `.pth` file, so the reconstructed model has the correct number of classes and resolution\n",
    "without you having to pass them explicitly.  You can share this file with teammates and they can run inference\n",
    "immediately — the class names and architecture are all embedded alongside the weights."
   ]
  },
  {
   "cell_type": "code",
   "id": "58b00247",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": "loaded_model = RFDETRSegSmall.from_checkpoint(FINAL_CHECKPOINT_PATH)",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "d2252a7a",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "## 10 - Select inference images\n",
    "\n",
    "This cell implements a fallback chain — test split first, then validation, then train — so the notebook always finds\n",
    "images to run inference on even when the dataset has no dedicated test split. Using images from the test set gives\n",
    "you an unbiased view of model performance because those images were never seen during training or used to pick the\n",
    "best checkpoint. Replace `inference_image_paths` with a list of `Path` objects pointing to your own files to run\n",
    "inference on custom images."
   ]
  },
  {
   "cell_type": "code",
   "id": "f03e8dcc",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "_IMAGE_EXTS = {\".jpg\", \".jpeg\", \".png\"}\n",
    "\n",
    "\n",
    "def _find_images(directory: Path) -> list[Path]:\n",
    "    if not directory.is_dir():\n",
    "        return []\n",
    "    return sorted(p for p in directory.glob(\"**/*\") if p.is_file() and p.suffix.lower() in _IMAGE_EXTS)\n",
    "\n",
    "\n",
    "validation_image_paths = _find_images(DATASET_DIR / \"test\")\n",
    "if not validation_image_paths:\n",
    "    validation_image_paths = _find_images(DATASET_DIR / \"valid\")\n",
    "if not validation_image_paths:\n",
    "    validation_image_paths = _find_images(DATASET_DIR / \"train\")\n",
    "if not validation_image_paths:\n",
    "    raise FileNotFoundError(f\"No images ({', '.join(sorted(_IMAGE_EXTS))}) found under {DATASET_DIR}\")\n",
    "\n",
    "inference_image_paths = validation_image_paths[:INFERENCE_COUNT]\n",
    "print(f\"inference_images={[str(p) for p in inference_image_paths]}\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "74116bbe",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "## 11 - Visualize inference\n",
    "\n",
    "`INFERENCE_THRESHOLD` is a detection confidence threshold: only detections whose bounding-box score exceeds this\n",
    "value are shown.  Raise it to suppress false positives; lower it to surface low-confidence detections.  Each accepted\n",
    "detection includes a binary mask resized to the original image dimensions — `sv.MaskAnnotator` renders it as a\n",
    "translucent colour overlay so you can see both the mask shape and the image beneath it.\n",
    "\n",
    "If the mask boundaries look jagged, consider increasing `RESOLUTION` (in multiples of 24) and retraining; higher\n",
    "resolution gives the segmentation head more spatial detail to work with."
   ]
  },
  {
   "cell_type": "code",
   "id": "f3a68bfe",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "inference_grid_items: list[tuple[str, np.ndarray, sv.Detections]] = []\n",
    "for image_path in inference_image_paths:\n",
    "    with Image.open(image_path) as image:\n",
    "        image_rgb = image.convert(\"RGB\")\n",
    "    detections = loaded_model.predict(image_rgb, threshold=INFERENCE_THRESHOLD)\n",
    "    if not isinstance(detections, sv.Detections):\n",
    "        raise RuntimeError(f\"Expected RFDETRSegSmall.predict() to return sv.Detections, got {type(detections)!r}.\")\n",
    "    inference_grid_items.append((image_path.name, np.array(image_rgb), detections))\n",
    "\n",
    "if inference_grid_items:\n",
    "    figure = _segmentation_grid_figure(inference_grid_items, columns=INFERENCE_COLUMNS, class_names=CLASS_NAMES)\n",
    "    display(figure)\n",
    "    plt.close(figure)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "63027bc2",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "## 12 - Detection summary\n",
    "\n",
    "Print per-image detection counts and confidence scores as a quick sanity check.  If you see zero detections on most\n",
    "images, try lowering `INFERENCE_THRESHOLD` (towards 0.1) — the model may be calibrated to lower confidence scores\n",
    "on a custom dataset than the COCO-pretrained default.  If you see too many false positives, raise the threshold or\n",
    "add more negative examples to your training set."
   ]
  },
  {
   "cell_type": "code",
   "id": "d723ad36",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "summary_rows = []\n",
    "for image_path, (_, _, detections) in zip(inference_image_paths, inference_grid_items):\n",
    "    num_det = len(detections)\n",
    "    has_masks = detections.mask is not None and detections.mask.shape[0] > 0\n",
    "    avg_conf = (\n",
    "        float(np.mean(detections.confidence)) if detections.confidence is not None and num_det > 0 else float(\"nan\")\n",
    "    )\n",
    "    summary_rows.append(\n",
    "        {\n",
    "            \"image\": image_path.name,\n",
    "            \"detections\": num_det,\n",
    "            \"has_masks\": has_masks,\n",
    "            \"avg_confidence\": round(avg_conf, 3),\n",
    "        }\n",
    "    )\n",
    "\n",
    "summary_df = pd.DataFrame(summary_rows)\n",
    "display(summary_df)"
   ],
   "outputs": [],
   "execution_count": null
  }
 ],
 "metadata": {
  "jupytext": {
   "cell_metadata_filter": "title,-all",
   "main_language": "python",
   "notebook_metadata_filter": "-all"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}