chore: import upstream snapshot with attribution

docs/cookbooks/NOTES.md

@@ -0,0 +1,58 @@
+# Notebooks
+
+Each `.ipynb` file here is rendered as a page under `/cookbooks/` in the docs site.
+
+Cards on the cookbooks landing page are driven by [`cards.yaml`](cards.yaml). The MkDocs hook
+`docs/hooks/cookbooks_cards.py` loads that file and exposes it to `docs/theme/notebooks.html`,
+which renders each entry as a card via a Jinja loop.
+
+## Converting a jupytext `.py` to `.ipynb`
+
+Cookbook source files live as jupytext percent-format `.py` scripts (e.g. `fine-tune_keypoints.py`) inside `docs/cookbooks/`. Each script requires at minimum a **docs render copy**; some also have a `notebooks/` copy for users who want to run it directly. Regenerate every existing copy after each edit:
+
+```bash
+# Docs render copy (served by mkdocs-jupyter at /cookbooks/) — always required
+jupytext --to notebook fine-tune_keypoints.py --output docs/cookbooks/fine-tune_keypoints.ipynb
+
+# Runnable copy in notebooks/ — only for notebooks explicitly placed there
+jupytext --to notebook fine-tune_keypoints.py --output notebooks/fine-tune_keypoints.ipynb
+```
+
+New notebooks default to the docs-only copy. Add a `notebooks/` copy only when there is an explicit need (e.g. a runnable starter notebook shipped with the repo). Omit `--execute` — notebooks are rendered statically by `mkdocs-jupyter` with `execute: false`, so pre-run outputs in the `.ipynb` are displayed as-is.
+
+If jupytext is not installed: `pip install jupytext` (or `uv add jupytext --dev`).
+
+## Adding a notebook
+
+1. Add the `.ipynb` file here, named after its content (e.g. `custom-augmentations.ipynb`, `onnx-export.ipynb`).
+2. Add a new entry to `docs/cookbooks/cards.yaml` under the `cards:` list:
+
+<!-- prettier-ignore -->
+
+```yaml
+  - href: content-slug/
+    name: Short Title
+    labels: [LABEL1, LABEL2]
+    version: vX.Y.0
+    author: GitHubUsername
+    description: One sentence describing what the notebook demonstrates.
+```
+
+Available labels (reuse these to keep tags standardised): `TRAINING`, `AUGMENTATION`, `EXPORT`, `TFLITE`, `PYTORCH LIGHTNING`, `INFERENCE`, `SEGMENTATION`, `DEPLOY`.
+Current tag colours are assigned dynamically by the docs UI, so they may change if cards or labels are added or reordered.
+
+## Removing a notebook
+
+1. Delete the `.ipynb` file.
+2. Remove the matching entry (the `- href: content-slug/` block) from `docs/cookbooks/cards.yaml`.
+
+## Current notebooks
+
+| File                                | Card title                                      | Version |
+| ----------------------------------- | ----------------------------------------------- | ------- |
+| `custom-augmentations.ipynb`        | Custom Augmentations and Live Training Progress | v1.5.0  |
+| `fine-tune_detection.ipynb`         | Fine-Tune RF-DETR Object Detection              | v1.8.0  |
+| `fine-tune_keypoints.ipynb`         | Fine-Tune RF-DETR Keypoint Detection            | v1.8.0  |
+| `fine-tune_segmentation.ipynb`      | Fine-Tune RF-DETR Instance Segmentation         | v1.8.2  |
+| `inference-latency-benchmark.ipynb` | Inference Latency Benchmark                     | v1.8.2  |
+| `pytorch-lightning.ipynb`           | Training with PyTorch Lightning                 | v1.6.0  |

docs/cookbooks/cards.yaml

@@ -0,0 +1,37 @@
+cards:
+  - href: fine-tune_detection/
+    name: "Fine-Tune RF-DETR Object Detection"
+    labels: [TRAINING, PYTORCH LIGHTNING, INFERENCE]
+    version: v1.8.2
+    author: Borda
+    description: "Fine-tune RF-DETR object detection on any COCO detection dataset from Roboflow Universe — covers PPE safety, traffic, sports, infrastructure, and marine-wildlife use cases."
+  - href: inference-latency-benchmark/
+    name: "Inference Latency Benchmark"
+    labels: [INFERENCE, ONNX, GPU]
+    version: v1.8.1
+    author: Borda
+    description: "Benchmark detection, segmentation, and keypoint models across FP32, FP16+JIT, and ONNX Runtime — measures latency and FPS on GPU using CUDA events."
+  - href: fine-tune_segmentation/
+    name: "Fine-Tune RF-DETR Instance Segmentation"
+    labels: [TRAINING, PYTORCH LIGHTNING, SEGMENTATION, INFERENCE]
+    version: v1.8.1
+    author: Borda
+    description: "Fine-tune RF-DETR instance segmentation on any COCO segmentation dataset from Roboflow Universe — covers training, mask loss metrics, and inference with mask overlays."
+  - href: fine-tune_keypoints/
+    name: "Fine-Tune RF-DETR Keypoint Detection"
+    labels: [TRAINING, PYTORCH LIGHTNING, INFERENCE]
+    version: v1.8.0
+    author: Borda
+    description: "Fine-tune RF-DETR keypoint detection on any COCO keypoint dataset from Roboflow Universe — covers training, metrics, and inference with uncertainty ellipses."
+  - href: custom-augmentations/
+    name: "Custom Augmentations and Live Training Progress"
+    labels: [TRAINING, AUGMENTATION]
+    version: v1.5.0
+    author: Borda
+    description: "Add custom Albumentations augmentations and monitor training metrics in real time."
+  - href: pytorch-lightning/
+    name: "Training with PyTorch Lightning"
+    labels: [TRAINING, PYTORCH LIGHTNING]
+    version: v1.6.0
+    author: Borda
+    description: "Train RF-DETR with PyTorch Lightning — callbacks, checkpointing, and the LightningModule API."

docs/cookbooks/custom-augmentations.ipynb

@@ -0,0 +1,417 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "0f72b166",
+   "metadata": {},
+   "source": [
+    "# Custom Augmentations and Live Training Progress\n",
+    "\n",
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/roboflow/rf-detr/blob/develop/docs/cookbooks/custom-augmentations.ipynb)\n",
+    "\n",
+    "**RF-DETR** is a real-time object detection model that combines the accuracy of\n",
+    "transformer-based detectors with inference speeds suitable for production.\n",
+    "1.5.0 brings two headline additions:\n",
+    "\n",
+    "1. **Custom training augmentations via Albumentations** — a flexible system that\n",
+    "   lets you control exactly how images are transformed during training, with bounding\n",
+    "   boxes and segmentation masks kept in sync automatically.  Four ready-made presets\n",
+    "   cover the most common scenarios out of the box.\n",
+    "2. **Live progress bars and structured epoch logs** — per-epoch rich / tqdm progress\n",
+    "   so you can monitor batch-level metrics without parsing raw log output.\n",
+    "\n",
+    "You will learn how to:\n",
+    "- Explore and compare the four built-in augmentation presets\n",
+    "- Visually inspect augmented samples *before* committing to a full training run\n",
+    "- Define a fully custom augmentation pipeline\n",
+    "- Train RF-DETR with your chosen augmentation config\n",
+    "- Run inference with the trained model"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "49ad10f9",
+   "metadata": {},
+   "source": [
+    "## 1. Install RF-DETR 1.5.0\n",
+    "\n",
+    "`rfdetr` includes the augmentation system and progress-bar support introduced in\n",
+    "this release.  `supervision` handles visualization."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5c1df817",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!pip install -q rfdetr==1.5.0"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "075cf0b4",
+   "metadata": {},
+   "source": [
+    "## 2. Check GPU availability\n",
+    "\n",
+    "RF-DETR trains on GPU when one is available and falls back to CPU otherwise.\n",
+    "The cell below detects your device and prints VRAM size — a useful sanity check\n",
+    "before choosing batch size."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d005aa12",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "\n",
+    "import torch\n",
+    "\n",
+    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
+    "print(f\"Using device: {device}\")\n",
+    "if device == \"cuda\":\n",
+    "    print(f\"GPU: {torch.cuda.get_device_name(0)}\")\n",
+    "    print(f\"VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB\")\n",
+    "\n",
+    "# Scale data-loader workers with available CPUs so the GPU is kept fed.\n",
+    "num_workers = max(os.cpu_count() or 0, 2)\n",
+    "print(f\"Data loader workers: {num_workers}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c926124e",
+   "metadata": {},
+   "source": [
+    "## 3. Download COCO 2017\n",
+    "\n",
+    "We use the official train/val splits — no manual splitting needed.\n",
+    "\n",
+    "| Split | Images | Size |\n",
+    "|---|---|---|\n",
+    "| `train2017` | ~118 000 | ~18 GB |\n",
+    "| `val2017` | 5 000 | ~1 GB |\n",
+    "| annotations | — | ~241 MB |"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a3cc746d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!wget -q --show-progress http://images.cocodataset.org/zips/train2017.zip -O train2017.zip\n",
+    "!wget -q --show-progress http://images.cocodataset.org/zips/val2017.zip -O val2017.zip\n",
+    "!wget -q --show-progress http://images.cocodataset.org/annotations/annotations_trainval2017.zip -O annotations.zip"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "985d1e18",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!unzip -q train2017.zip\n",
+    "!unzip -q val2017.zip\n",
+    "!unzip -q annotations.zip"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "27b7ae27",
+   "metadata": {},
+   "source": [
+    "### Set up the dataset directory structure\n",
+    "\n",
+    "`model.train()` expects:\n",
+    "```\n",
+    "dataset/\n",
+    "  train/  _annotations.coco.json  + images\n",
+    "  valid/  _annotations.coco.json  + images\n",
+    "```"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "da25896a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!mkdir -p coco_demo\n",
+    "!mv train2017 coco_demo/train\n",
+    "!mv val2017   coco_demo/valid\n",
+    "!cp annotations/instances_train2017.json coco_demo/train/_annotations.coco.json\n",
+    "!cp annotations/instances_val2017.json   coco_demo/valid/_annotations.coco.json"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "69343c86",
+   "metadata": {},
+   "source": [
+    "## 4. Explore the built-in augmentation presets\n",
+    "\n",
+    "RF-DETR ships four ready-made presets tuned for common use cases.  Each preset\n",
+    "is a plain Python dict mapping Albumentations transform names to their constructor\n",
+    "kwargs (including `p`, the probability of applying the transform).  You can\n",
+    "inspect, merge, or extend them like any other dict.\n",
+    "\n",
+    "| Preset | When to use |\n",
+    "|---|---|\n",
+    "| `AUG_CONSERVATIVE` | Small datasets (< 500 images) — gentle transforms to avoid overfitting |\n",
+    "| `AUG_AGGRESSIVE` | Large datasets (2 000+ images) — stronger augmentations for better generalisation |\n",
+    "| `AUG_AERIAL` | Satellite and overhead imagery — rotation-invariant transforms |\n",
+    "| `AUG_INDUSTRIAL` | Manufacturing and inspection data — handles structured textures |"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "409e006a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import rfdetr.datasets.aug_configs as aug_config\n",
+    "from rfdetr.datasets.aug_configs import AUG_AGGRESSIVE\n",
+    "\n",
+    "for name in (\"AUG_CONSERVATIVE\", \"AUG_AGGRESSIVE\", \"AUG_AERIAL\", \"AUG_INDUSTRIAL\"):\n",
+    "    preset = getattr(aug_config, name)\n",
+    "    print(f\"\\n{name}:\")\n",
+    "    for transform, params in preset.items():\n",
+    "        print(f\"  {transform}: {params}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d10837d5",
+   "metadata": {},
+   "source": [
+    "## 5. Inspect augmented samples before training\n",
+    "\n",
+    "Setting `save_dataset_grids=True` writes 3×3 JPEG grids of augmented training\n",
+    "and validation images to `output_dir` *before any weight updates occur*.  Use\n",
+    "this to visually sanity-check your pipeline in seconds — catching problems like\n",
+    "flipped labels or extreme colour shifts before committing to a full training run.\n",
+    "\n",
+    "Saved files:\n",
+    "```\n",
+    "output/\n",
+    "  train_batch0_grid.jpg  train_batch1_grid.jpg  train_batch2_grid.jpg\n",
+    "  val_batch0_grid.jpg    val_batch1_grid.jpg    val_batch2_grid.jpg\n",
+    "```"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d7797651",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "\n",
+    "from rfdetr import RFDETRNano  # smallest & fastest model — ideal for demos\n",
+    "\n",
+    "DATASET_DIR = \"coco_demo\"\n",
+    "OUTPUT_DIR = \"output\"\n",
+    "os.makedirs(OUTPUT_DIR, exist_ok=True)\n",
+    "\n",
+    "model = RFDETRNano()\n",
+    "model.train(\n",
+    "    dataset_dir=str(DATASET_DIR),\n",
+    "    epochs=1,  # one epoch is enough to generate the grids\n",
+    "    batch_size=12,\n",
+    "    aug_config=AUG_AGGRESSIVE,\n",
+    "    save_dataset_grids=True,\n",
+    "    output_dir=OUTPUT_DIR,\n",
+    "    device=device,\n",
+    "    num_workers=num_workers,\n",
+    "    run_test=False,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7b7cb782",
+   "metadata": {},
+   "source": [
+    "### Display the saved grids inline"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "28df5b59",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pathlib import Path\n",
+    "\n",
+    "%matplotlib inline\n",
+    "\n",
+    "import matplotlib.image as mpimg\n",
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "grids = sorted(Path(OUTPUT_DIR).glob(\"*_grid.jpg\"))\n",
+    "if grids:\n",
+    "    fig, axes = plt.subplots(len(grids), 1, figsize=(6, 6 * len(grids)))\n",
+    "    if len(grids) == 1:\n",
+    "        axes = [axes]\n",
+    "    for ax, grid_path in zip(axes, grids):\n",
+    "        ax.imshow(mpimg.imread(grid_path))\n",
+    "        ax.set_title(grid_path.name)\n",
+    "        ax.axis(\"off\")\n",
+    "    plt.tight_layout()\n",
+    "    plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d7f10179",
+   "metadata": {},
+   "source": [
+    "## 6. Choose your augmentation config\n",
+    "\n",
+    "Pick one of the four options below.  Option A (a built-in preset) is the fastest\n",
+    "way to get started.  Options B–D give increasing levels of control.  The selected\n",
+    "`AUG_CONFIG` is passed straight to `model.train()` in the next section."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f62b76de",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# --- Option A: built-in preset ---\n",
+    "AUG_CONFIG = AUG_AGGRESSIVE\n",
+    "\n",
+    "# --- Option B: extend a preset ---\n",
+    "# AUG_CONFIG = {**AUG_AGGRESSIVE, \"VerticalFlip\": {\"p\": 0.3}}\n",
+    "\n",
+    "# --- Option C: fully custom ---\n",
+    "# AUG_CONFIG = {\n",
+    "#     \"HorizontalFlip\": {\"p\": 0.5},\n",
+    "#     \"Rotate\": {\"limit\": 15, \"p\": 0.3},\n",
+    "#     \"RandomBrightnessContrast\": {\"brightness_limit\": 0.2, \"contrast_limit\": 0.2, \"p\": 0.4},\n",
+    "#     \"GaussianBlur\": {\"blur_limit\": 3, \"p\": 0.2},\n",
+    "# }\n",
+    "\n",
+    "# --- Option D: no augmentations ---\n",
+    "# AUG_CONFIG = {}\n",
+    "\n",
+    "print(\"Selected augmentation config:\")\n",
+    "for transform, params in AUG_CONFIG.items():\n",
+    "    print(f\"  {transform}: {params}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5a5ff7b9",
+   "metadata": {},
+   "source": [
+    "## 7. Train with augmentations\n",
+    "\n",
+    "Run a full training pass with the augmentation config chosen above.  The\n",
+    "`progress_bar=\"rich\"` setting activates per-epoch live progress bars — a batch\n",
+    "counter, loss, and learning rate are updated in real time.  Structured metric\n",
+    "tables are printed at the end of each epoch.\n",
+    "\n",
+    "`aug_config` accepts any dict of Albumentations transforms, including the\n",
+    "`AUG_CONFIG` defined above.  Passing `{}` disables all augmentations."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "79d29fd9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model = RFDETRNano()\n",
+    "model.train(\n",
+    "    dataset_dir=str(DATASET_DIR),\n",
+    "    epochs=2,\n",
+    "    batch_size=24,\n",
+    "    aug_config=AUG_CONFIG,\n",
+    "    output_dir=OUTPUT_DIR,\n",
+    "    device=device,\n",
+    "    num_workers=num_workers,\n",
+    "    progress_bar=\"rich\",\n",
+    "    run_test=False,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "beccaf40",
+   "metadata": {},
+   "source": [
+    "## 8. Run inference\n",
+    "\n",
+    "Load a sample image and run `model.predict()`.  The method returns a\n",
+    "`supervision.Detections` object, which makes it straightforward to draw\n",
+    "bounding boxes and labels with the `supervision` annotators."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d79c7612",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import requests\n",
+    "import supervision as sv\n",
+    "from PIL import Image\n",
+    "\n",
+    "from rfdetr.util.coco_classes import COCO_CLASSES\n",
+    "\n",
+    "image_url = \"https://media.roboflow.com/dog.jpg\"\n",
+    "image = Image.open(requests.get(image_url, stream=True).raw)\n",
+    "\n",
+    "detections = model.predict(image, threshold=0.5)\n",
+    "\n",
+    "labels = [COCO_CLASSES[class_id] for class_id in detections.class_id]\n",
+    "annotated = sv.BoxAnnotator().annotate(image.copy(), detections)\n",
+    "annotated = sv.LabelAnnotator().annotate(annotated, detections, labels)\n",
+    "sv.plot_image(annotated)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d69f8215",
+   "metadata": {},
+   "source": [
+    "## 9. Next steps\n",
+    "\n",
+    "You have seen the complete 1.5.0 augmentation workflow — from exploring presets\n",
+    "through live inspection to a full training run.  From here:\n",
+    "\n",
+    "- [Augmentation docs](https://rfdetr.roboflow.com/1.5.0/learn/train/augmentations/) — full transform reference and custom pipeline guide\n",
+    "- [Advanced training options](https://rfdetr.roboflow.com/1.5.0/learn/train/advanced/) — EMA, gradient accumulation, learning rate schedules\n",
+    "- [Logger integrations (ClearML, MLflow, W&B)](https://rfdetr.roboflow.com/1.5.0/learn/train/loggers/) — experiment tracking\n",
+    "- [Export your model](https://rfdetr.roboflow.com/1.5.0/learn/export/) — ONNX, TensorRT, CoreML\n",
+    "- [RF-DETR 1.6.0 notebook](https://colab.research.google.com/github/roboflow/rf-detr/blob/develop/notebooks/release-demo_1-6.ipynb) — PyTorch Lightning building blocks"
+   ]
+  }
+ ],
+ "metadata": {
+  "jupytext": {
+   "cell_metadata_filter": "-all",
+   "formats": "ipynb,py",
+   "main_language": "python"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

docs/cookbooks/fine-tune_detection.ipynb

@@ -0,0 +1,944 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "920e5414",
+   "metadata": {},
+   "source": [
+    "# RF-DETR object detection 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",
+    "## What is object detection?\n",
+    "\n",
+    "Object detection answers two questions simultaneously: *what* objects are in an image and *where* they are.\n",
+    "Each detected object gets a **bounding box** (four coordinates) and a **class label** (e.g. \"car\", \"person\").\n",
+    "This is the most common computer vision task — it powers traffic monitoring, safety compliance, retail analytics,\n",
+    "medical imaging, environmental surveys, and virtually every real-time vision application.\n",
+    "\n",
+    "## Why fine-tune instead of train from scratch?\n",
+    "\n",
+    "RF-DETR ships pre-trained on the COCO dataset (80 common object classes). Fine-tuning reuses the backbone's\n",
+    "rich visual representations — edges, textures, object parts — learned from millions of images.\n",
+    "Starting from these weights instead of random initialisation typically:\n",
+    "- **Reduces training time** from days to minutes or hours\n",
+    "- **Requires far fewer labelled images** (hundreds instead of tens of thousands)\n",
+    "- **Achieves higher accuracy** than training from scratch on small datasets\n",
+    "\n",
+    "The process is straightforward: swap out the final classification head for your number of classes,\n",
+    "then continue training with a small learning rate so the pre-trained features are preserved.\n",
+    "\n",
+    "## Dataset overview\n",
+    "\n",
+    "Five datasets covering different domains are pre-configured. Set `DATASET_KEY` to one of:\n",
+    "\n",
+    "| Key | Dataset | Domain | Classes |\n",
+    "|-----|---------|--------|---------|\n",
+    "| `\"hard_hat_ppes\"` | Hard Hat Worker Safety | Construction / PPE | hat, person, vest |\n",
+    "| `\"road_damage\"` | Pavement Distress Detection | Infrastructure Inspection | pot_hole, cracking, ravelling, ... |\n",
+    "| `\"traffic\"` | Traffic Detection | Autonomous Driving | car, bus, truck, motorbike, bike, traffic lights |\n",
+    "| `\"football\"` | Football Player Detection | Sports Analytics | player, ball, goalkeeper, referee |\n",
+    "| `\"brackish_underwater\"` | Brackish Underwater | Marine Wildlife | fish, crab, jellyfish, shrimp, starfish, small_fish |\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 bounding boxes."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b6eb7cd5",
+   "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\n",
+    "needed for the metrics table and optional plot styling. If you hit an `ImportError` after running this cell,\n",
+    "the most likely cause is a stale in-memory import — restart the Python runtime once and re-run from the top.\n",
+    "\n",
+    "**GPU recommended.** Fine-tuning 50 epochs on a ~1,500-image dataset takes roughly 10–25 minutes on a modern GPU\n",
+    "(RTX 3090, A10, T4). CPU-only is possible for a quick test but will be much slower."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e79facfd",
+   "metadata": {
+    "title": "[bash]"
+   },
+   "outputs": [],
+   "source": "!pip install -q \"rfdetr[train,visual]>=1.8.0\" roboflow pandas seaborn"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "62aff634",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\"\"\"Shared RF-DETR object detection fine-tuning demo for several Roboflow COCO detection 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 RFDETRSmall\n",
+    "from rfdetr.config import TrainConfig\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": "c9c78daa",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "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 (32 for RFDETRSmall: patch_size=16, num_windows=2).\n",
+    "# Higher resolution improves detection of small objects at the cost of GPU memory and training time.\n",
+    "# Practical range: 384–640. The default 512 suits most fine-tuning runs.\n",
+    "DATASETS: dict[str, dict[str, Any]] = {\n",
+    "    \"hard_hat_ppes\": {\n",
+    "        \"name\": \"Hard Hat Worker Safety\",\n",
+    "        \"workspace\": \"safety-first\",\n",
+    "        \"project\": \"hard-hat-worker-safety-equipments\",\n",
+    "        \"version\": 9,\n",
+    "        \"source_url\": \"https://universe.roboflow.com/safety-first/hard-hat-worker-safety-equipments\",\n",
+    "        \"output_name\": \"det_hard_hat_ppes\",\n",
+    "        # 512 — outdoor construction scenes; workers appear at medium scale so the\n",
+    "        # default resolution is sufficient to distinguish helmet vs no-helmet.\n",
+    "        \"resolution\": 512,\n",
+    "    },\n",
+    "    \"road_damage\": {\n",
+    "        \"name\": \"Pavement Distress Detection\",\n",
+    "        \"workspace\": \"pavement-distresses\",\n",
+    "        \"project\": \"distress-detection\",\n",
+    "        \"version\": 2,\n",
+    "        \"source_url\": \"https://universe.roboflow.com/pavement-distresses/distress-detection\",\n",
+    "        \"output_name\": \"det_road_damage\",\n",
+    "        # 544 — road-surface images captured from dashcams and drones; fine cracks\n",
+    "        # and potholes are small relative to the frame, so the extra resolution\n",
+    "        # helps the model distinguish closely spaced defect categories.\n",
+    "        \"resolution\": 544,\n",
+    "    },\n",
+    "    \"traffic\": {\n",
+    "        \"name\": \"Traffic Detection\",\n",
+    "        \"workspace\": \"redlightrunningdection\",\n",
+    "        \"project\": \"traffic-detection-sutq6\",\n",
+    "        \"version\": 36,\n",
+    "        \"source_url\": \"https://universe.roboflow.com/redlightrunningdection/traffic-detection-sutq6\",\n",
+    "        \"output_name\": \"det_traffic\",\n",
+    "        # 512 — street-level traffic scenes; vehicles span a wide size range\n",
+    "        # (large buses to distant motorbikes); 512 gives a good speed/accuracy trade-off.\n",
+    "        \"resolution\": 512,\n",
+    "    },\n",
+    "    \"football\": {\n",
+    "        \"name\": \"Football Player Detection\",\n",
+    "        \"workspace\": \"football-gozni\",\n",
+    "        \"project\": \"football-player-detection-bfswn\",\n",
+    "        \"version\": 1,\n",
+    "        \"source_url\": \"https://universe.roboflow.com/football-gozni/football-player-detection-bfswn\",\n",
+    "        \"output_name\": \"det_football\",\n",
+    "        # 544 — broadcast and stadium-angle shots; the ball is a tiny object relative\n",
+    "        # to the field; the extra resolution meaningfully reduces false negatives.\n",
+    "        \"resolution\": 544,\n",
+    "    },\n",
+    "    \"brackish_underwater\": {\n",
+    "        \"name\": \"Brackish Underwater\",\n",
+    "        \"workspace\": \"brad-dwyer\",\n",
+    "        \"project\": \"brackish-underwater\",\n",
+    "        \"version\": 2,\n",
+    "        \"source_url\": \"https://universe.roboflow.com/brad-dwyer/brackish-underwater\",\n",
+    "        \"output_name\": \"det_brackish_underwater\",\n",
+    "        # 512 — sonar and underwater camera frames; marine creatures appear at\n",
+    "        # medium-to-large scale; the default resolution captures enough detail.\n",
+    "        \"resolution\": 512,\n",
+    "    },\n",
+    "}\n",
+    "\n",
+    "DATASET_KEY = \"hard_hat_ppes\"\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": "cdce84bd",
+   "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",
+   "execution_count": null,
+   "id": "993a411e",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "outputs": [],
+   "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()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9c7f944e",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "source": [
+    "## 1 - Download dataset\n",
+    "\n",
+    "Roboflow exports object detection datasets in **COCO format**: a JSON file containing image metadata and per-image\n",
+    "annotations, each with a `bbox` field (x, y, width, height in pixels) and a `category_id` referencing the class.\n",
+    "The download cell fetches the exact dataset version listed in `DATASETS` and places it under\n",
+    "`datasets/<DATASET_KEY>/`. You need a Roboflow API key — get one for free at `app.roboflow.com/settings/api`\n",
+    "and set it as the `ROBOFLOW_API_KEY` environment variable (or as a Colab secret with the same name).\n",
+    "\n",
+    "The download is **idempotent**: if the target directory already exists Roboflow skips the network transfer,\n",
+    "so re-running this cell after a successful download is fast.\n",
+    "\n",
+    "> **Tip:** All five datasets in this notebook are open-access on Roboflow Universe. You can browse them in your\n",
+    "> browser to explore image samples and annotation statistics before downloading."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "20fea837",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "outputs": [],
+   "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}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6516f6ee",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "source": [
+    "## 2 - Infer class names\n",
+    "\n",
+    "The COCO annotation file defines a `categories` list mapping integer IDs to human-readable names. RF-DETR uses\n",
+    "**0-based class indices** internally, so we sort the categories by their original `id` field and read names in order.\n",
+    "The resulting `CLASS_NAMES` list maps each index back to a label, which is used during inference visualisation.\n",
+    "\n",
+    "If your own dataset has a single category, this list will have one element — that is expected and valid."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5ca32aa1",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "outputs": [],
+   "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}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3e39e02e",
+   "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`\n",
+    "call on any machine, without reconstructing the original config."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "89450ec9",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "outputs": [],
+   "source": [
+    "def _save_final_checkpoint(\n",
+    "    module: RFDETRModelModule,\n",
+    "    trainer: Any,\n",
+    "    train_config: TrainConfig,\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=\"RFDETRSmall\",\n",
+    "            model_config_dict=model_config.model_dump(),\n",
+    "        ),\n",
+    "        output_path,\n",
+    "    )\n",
+    "    return output_path"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e0a0ad62",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "source": [
+    "`_detection_grid_figure` arranges multiple annotated images into a fixed-column matplotlib grid.\n",
+    "Bounding boxes are drawn with `sv.BoxAnnotator` and labels with `sv.LabelAnnotator`, showing the class name and\n",
+    "confidence score for each detection. Using a grid keeps the visual summary compact and makes it easy to compare\n",
+    "predictions across images at a glance."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "efae67bb",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "outputs": [],
+   "source": [
+    "def _detection_grid_figure(\n",
+    "    items: list[tuple[str, np.ndarray, sv.Detections]],\n",
+    "    columns: int,\n",
+    "    class_names: list[str],\n",
+    ") -> Figure:\n",
+    "    \"\"\"Render bounding box detections in a fixed-column subplot grid.\"\"\"\n",
+    "    if columns <= 0:\n",
+    "        raise ValueError(f\"columns must be positive, got {columns}.\")\n",
+    "\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",
+    "        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"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "38359554",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "source": [
+    "## 3 - Configure training\n",
+    "\n",
+    "`TrainConfig` centralises every hyperparameter the training loop needs. The most important fields to understand\n",
+    "before customising them for your own dataset:\n",
+    "\n",
+    "**Learning rates** — `lr` and `lr_encoder` are both set to `1e-4` for fine-tuning. The `lr` controls the\n",
+    "detection head and decoder; `lr_encoder` controls the vision backbone (DINOv2). For fine-tuning, both are usually\n",
+    "equal. If you notice the backbone overfitting early — validation loss rising while training loss keeps falling —\n",
+    "halve `lr_encoder` to protect the pre-trained features while the head continues to adapt.\n",
+    "\n",
+    "**Batch size and gradient accumulation** — `batch_size=8` with `grad_accum_steps=2` gives an effective batch size\n",
+    "of 16 without requiring extra GPU memory. Effective batch size affects how smoothly gradients are estimated each\n",
+    "update; too small (< 8) leads to noisy updates; too large (> 64) may over-smooth and slow convergence.\n",
+    "\n",
+    "**EMA** — `use_ema=False` keeps this demo fast. For a production run you can enable EMA\n",
+    "(`use_ema=True`) to maintain an exponential moving average of the weights, which typically adds 0.5–1.0 mAP points\n",
+    "on small datasets by reducing the impact of noisy late-epoch updates.\n",
+    "\n",
+    "**Multi-scale training** — `multi_scale=False` and `expanded_scales=False` disable the multi-resolution\n",
+    "augmentation used during pre-training. Turning them off shortens epoch time by up to 40 % and is fine for\n",
+    "most fine-tuning runs on small custom datasets. Re-enable if your validation mAP plateaus early.\n",
+    "\n",
+    "**Notes dict** — stored verbatim in the checkpoint file alongside the weights, giving you a lightweight experiment\n",
+    "log that travels with the model file."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a8c0bbcf",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "outputs": [],
+   "source": [
+    "seed_all(SEED)\n",
+    "variant = RFDETRSmall(  # 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 = TrainConfig(\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\"detection 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)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c1eaa77a",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "source": [
+    "## 4 - Preview dataset inputs\n",
+    "\n",
+    "**Always look at your data before starting a long training run.** This cell renders a grid of annotated training\n",
+    "images drawn from the augmented pipeline — the same view the model will see during training.\n",
+    "\n",
+    "What to check here:\n",
+    "- **Box alignment** — do bounding boxes tightly surround the objects, or are they shifted or too large?\n",
+    "- **Class correctness** — are labels correct for each box? Mixed-up classes are one of the most common\n",
+    "  annotation errors and are hard to recover from after training.\n",
+    "- **Augmentation sanity** — horizontal flips, colour jitter, and random crops should look plausible;\n",
+    "  cropped boxes at image edges are normal.\n",
+    "- **Class imbalance** — if one class dominates the preview, the model will likely be biased toward it;\n",
+    "  rebalancing or oversampling rare classes can help.\n",
+    "\n",
+    "Catching label noise here costs a few seconds; catching it after 50 epochs of training costs much more."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a42a8135",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "outputs": [],
+   "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)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "58343884",
+   "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, bipartite\n",
+    "matching loss (classification + bounding-box L1 + GIoU), gradient accumulation, weight updates,\n",
+    "learning-rate scheduling, and periodic validation with COCO bounding-box mAP.\n",
+    "\n",
+    "After each epoch the trainer appends a row to `metrics.csv` and, if validation mAP improves,\n",
+    "saves a new best checkpoint via `BestModelCallback`.\n",
+    "\n",
+    "**Typical training times on a single GPU** (RTX 3090 or similar):\n",
+    "- PPE / hard hat dataset (~2,600 images): ~20–30 min for 50 epochs\n",
+    "- Traffic dataset (~1,400 images): ~15–20 min for 50 epochs\n",
+    "- Underwater dataset (~12,000 images): ~60–90 min for 50 epochs\n",
+    "\n",
+    "To **resume an interrupted run**, set `train_config.resume` to the path of the last PTL checkpoint\n",
+    "(e.g. `OUTPUT_DIR / \"last.ckpt\"`) before calling this cell again."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8c78f17f",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "outputs": [],
+   "source": [
+    "trainer.fit(model, datamodule=datamodule, ckpt_path=train_config.resume or None)\n",
+    "print(f\"output_dir={OUTPUT_DIR}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "92918976",
+   "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\n",
+    "are 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\n",
+    "configs, including the class names. You can share the file with a colleague and they can run inference with a\n",
+    "single `RFDETRSmall.from_checkpoint(path)` call, with no need to reconstruct the original config.\n",
+    "\n",
+    "For full reproducibility, keep this checkpoint alongside the dataset version number printed in section 1."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "47156e82",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "outputs": [],
+   "source": [
+    "final_checkpoint = _save_final_checkpoint(model, trainer, train_config, variant.model_config, FINAL_CHECKPOINT_PATH)\n",
+    "print(f\"saved_checkpoint={final_checkpoint}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9dbf2232",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "source": [
+    "## 7 - Validate metrics\n",
+    "\n",
+    "This cell runs a clean post-training validation: no augmentation, the best checkpoint loaded from\n",
+    "`checkpoint_best_total.pth` written by `BestModelCallback`, and results serialised to JSON for downstream\n",
+    "comparison or reporting.\n",
+    "\n",
+    "**Key metrics to inspect:**\n",
+    "\n",
+    "- **`bbox/map`** — the primary metric: standard COCO bounding-box mAP averaged across IoU thresholds 0.50–0.95\n",
+    "  in steps of 0.05. A value of 0.50 means the model correctly localises and classifies 50 % of objects at strict\n",
+    "  IoU thresholds. This is the most commonly reported number in object detection benchmarks.\n",
+    "- **`bbox/map_50`** — mAP at IoU ≥ 0.50 (a box is \"correct\" if it overlaps the ground-truth box by at least 50 %).\n",
+    "  Rises fastest and is the clearest early signal of whether the model is learning at all.\n",
+    "- **`bbox/map_75`** — mAP at the stricter IoU ≥ 0.75 threshold. Reflects how precisely the model localises\n",
+    "  objects, not just whether it finds them.\n",
+    "\n",
+    "A model with high `bbox/map_50` but low `bbox/map_75` finds objects reliably but draws imprecise boxes.\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",
+   "execution_count": null,
+   "id": "b35e7f96",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "outputs": [],
+   "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}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "cbb8112e",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "source": [
+    "## 8 - Plot CSVLogger metrics\n",
+    "\n",
+    "Reading loss and mAP curves is the fastest way to diagnose training problems.\n",
+    "\n",
+    "**Loss curves (first plot):**\n",
+    "- A healthy run shows both training loss and validation loss decreasing together.\n",
+    "- A **growing gap** between train and val loss is a classic overfitting signal — the model is memorising the\n",
+    "  training set. Add augmentation (`multi_scale=True`, or custom Albumentations), reduce learning rate, or collect\n",
+    "  more data.\n",
+    "- A **flat or rising loss from epoch 1** usually means the learning rate is too high. Try reducing `LR` to `5e-5`.\n",
+    "\n",
+    "**mAP curves (second plot):**\n",
+    "- `bbox_map_50` (solid line) rises fastest and is the clearest signal of learning.\n",
+    "- On small datasets (a few hundred images) you typically see rapid improvement in epochs 1–20 followed by\n",
+    "  a plateau around epoch 30–50.\n",
+    "- If mAP is **still rising at epoch 50**, extend `EPOCHS` to 75 or 100.\n",
+    "- If mAP **never rises above 0.05**, check class names, annotation quality, and whether\n",
+    "  `RESOLUTION` is appropriate for the objects' scale in your dataset.\n",
+    "\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",
+   "execution_count": null,
+   "id": "278892c2",
+   "metadata": {},
+   "outputs": [],
+   "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)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d0f5f468",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "outputs": [],
+   "source": [
+    "map_figure = plot_map_metrics(str(METRICS_CSV))\n",
+    "display(map_figure)\n",
+    "plt.close(map_figure)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "550cf8f7",
+   "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.\n",
+    "\n",
+    "This cell also confirms that the save–load round trip works before you proceed to inference, so any file corruption\n",
+    "or version mismatch surfaces here rather than silently producing wrong predictions later.\n",
+    "\n",
+    "To deploy this model in a production pipeline, you only need to ship the `.pth` file and call `from_checkpoint`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "dd0a0623",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "outputs": [],
+   "source": [
+    "loaded_model = RFDETRSmall.from_checkpoint(FINAL_CHECKPOINT_PATH)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "929edc54",
+   "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\n",
+    "finds images to run inference on even when the dataset has no dedicated test split.\n",
+    "\n",
+    "**Using test images gives an unbiased view of model performance** because those images were never seen during\n",
+    "training or used to select the best checkpoint. If your deployment images look different from the training set\n",
+    "(different lighting, camera angle, resolution), replace `inference_image_paths` with a list of `Path` objects\n",
+    "pointing to your own files to assess real-world performance before shipping the model."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "cc90e306",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "outputs": [],
+   "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",
+    "inference_image_paths = _find_images(DATASET_DIR / \"test\")\n",
+    "if not inference_image_paths:\n",
+    "    inference_image_paths = _find_images(DATASET_DIR / \"valid\")\n",
+    "if not inference_image_paths:\n",
+    "    inference_image_paths = _find_images(DATASET_DIR / \"train\")\n",
+    "if not inference_image_paths:\n",
+    "    raise FileNotFoundError(f\"No images ({', '.join(sorted(_IMAGE_EXTS))}) found under {DATASET_DIR}\")\n",
+    "\n",
+    "inference_image_paths = inference_image_paths[:INFERENCE_COUNT]\n",
+    "print(f\"inference_images={[str(p) for p in inference_image_paths]}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f49ac0a4",
+   "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\n",
+    "this value are shown. The default `0.3` is a reasonable starting point; adjust it based on what you see:\n",
+    "\n",
+    "- **Too many false positives** (boxes on background or wrong objects) → raise the threshold towards `0.5`–`0.7`.\n",
+    "- **Missing true objects** (objects in the image not detected) → lower the threshold towards `0.1`–`0.2`.\n",
+    "- **Zero detections on most images** → lower the threshold; the model may be well-calibrated to lower confidence\n",
+    "  scores on a custom dataset than the COCO-pretrained default.\n",
+    "\n",
+    "The confidence score reflects how certain the model is about each box. On a well-fine-tuned model you will\n",
+    "typically see scores above 0.7 for clear, unoccluded objects of the target classes."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4f7739a5",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "outputs": [],
+   "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 RFDETRSmall.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 = _detection_grid_figure(inference_grid_items, columns=INFERENCE_COLUMNS, class_names=CLASS_NAMES)\n",
+    "    display(figure)\n",
+    "    plt.close(figure)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "40582975",
+   "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. This table is useful for:\n",
+    "\n",
+    "- **Spotting images with zero detections** — if many images produce no detections at your chosen threshold,\n",
+    "  the model may be under-confident; try lowering `INFERENCE_THRESHOLD` first.\n",
+    "- **Checking average confidence** — a well-calibrated fine-tuned model typically reports average confidence\n",
+    "  between 0.5 and 0.9 on in-distribution images. Scores consistently below 0.3 suggest the model is uncertain\n",
+    "  and may benefit from more training epochs or additional labelled data.\n",
+    "- **Class distribution** — if certain classes never appear in inference results but were present in training data,\n",
+    "  they may be under-represented in the training set. Consider adding more examples for those classes."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "aa58d24e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "summary_rows = []\n",
+    "for image_path, (_, _, detections) in zip(inference_image_paths, inference_grid_items):\n",
+    "    num_det = len(detections)\n",
+    "    avg_conf = (\n",
+    "        float(np.mean(detections.confidence)) if detections.confidence is not None and num_det > 0 else float(\"nan\")\n",
+    "    )\n",
+    "    class_counts: dict[str, int] = {}\n",
+    "    if num_det > 0 and detections.class_id is not None:\n",
+    "        for cid in detections.class_id.tolist():\n",
+    "            label = CLASS_NAMES[cid] if cid < len(CLASS_NAMES) else str(cid)\n",
+    "            class_counts[label] = class_counts.get(label, 0) + 1\n",
+    "    summary_rows.append(\n",
+    "        {\n",
+    "            \"image\": image_path.name,\n",
+    "            \"detections\": num_det,\n",
+    "            \"avg_confidence\": round(avg_conf, 3),\n",
+    "            \"classes\": \", \".join(f\"{cls}×{cnt}\" for cls, cnt in sorted(class_counts.items())),\n",
+    "        }\n",
+    "    )\n",
+    "\n",
+    "summary_df = pd.DataFrame(summary_rows)\n",
+    "display(summary_df)"
+   ]
+  }
+ ],
+ "metadata": {
+  "jupytext": {
+   "cell_metadata_filter": "title,-all",
+   "main_language": "python",
+   "notebook_metadata_filter": "-all"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

docs/cookbooks/fine-tune_keypoints.ipynb

@@ -0,0 +1,878 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "7e465973",
+   "metadata": {},
+   "source": [
+    "# RF-DETR keypoint training demo on multiple 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 to predict structured keypoint skeletons — useful for human pose\n",
+    "estimation, sports field calibration, product inspection, and any task where *where* matters as much as *what*.\n",
+    "\n",
+    "**Datasets** — set `DATASET_KEY` to one of:\n",
+    "\n",
+    "| Key | Dataset |\n",
+    "|-----|---------|\n",
+    "| `\"dart\"` | Darts Detection |\n",
+    "| `\"human_pose\"` | Human Body Pose |\n",
+    "| `\"basketball_court\"` | Basketball Court Detection |\n",
+    "| `\"football_field\"` | Football Field Detection |\n",
+    "| `\"tennis_court\"` | Tennis Court Keypoint |\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\n",
+    "optional uncertainty ellipses."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6d361242",
+   "metadata": {},
+   "source": [
+    "## Setup\n",
+    "\n",
+    "Install `rfdetr` with the `train` and `visual` extras. The `train` extra pulls in the training loop dependencies (PyTorch Lightning, COCO evaluation tools, and data augmentation libraries), while `visual` adds the visualisation helpers used later in the notebook. The `roboflow` package handles dataset download; `pandas` and `seaborn` are needed for the metrics table and optional plot styling. If you hit an `ImportError` after running this cell, the most likely cause is a stale in-memory import — restart the Python runtime once and re-run from the top."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ca1b5461",
+   "metadata": {
+    "title": "[bash]"
+   },
+   "outputs": [],
+   "source": [
+    "!pip install -q \"rfdetr[train,visual]==1.8.2\" roboflow pandas seaborn"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d89771c9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\"\"\"Shared RF-DETR keypoint fine-tuning demo for several Roboflow COCO keypoint 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 RFDETRKeypointPreview\n",
+    "from rfdetr.config import KeypointTrainConfig\n",
+    "from rfdetr.datasets._keypoint_schema import infer_coco_keypoint_schema, infer_yolo_keypoint_schema\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.keypoints import _key_points_for_display, _keypoint_prediction_records\n",
+    "from rfdetr.visualize.training import plot_loss_metrics, plot_map_metrics"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0b7de4e2",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "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 the model patch_size (12 for RFDETRKeypointPreview).\n",
+    "# Larger resolutions improve fine-grained keypoint localisation at the cost of GPU memory\n",
+    "# and training time. Recommended choices per dataset are annotated below.\n",
+    "DATASETS: dict[str, dict[str, Any]] = {\n",
+    "    \"dart\": {\n",
+    "        \"name\": \"Darts Detection\",\n",
+    "        \"workspace\": \"dartdetection-vfgjd\",\n",
+    "        \"project\": \"darts_detection-yp7lt\",\n",
+    "        \"version\": 5,\n",
+    "        \"source_url\": \"https://universe.roboflow.com/dartdetection-vfgjd/darts_detection-yp7lt\",\n",
+    "        \"output_name\": \"keypoint_darts_detection_demo\",\n",
+    "        \"keypoint_flip_pairs\": [],\n",
+    "        # 576 — close-up shots where the dart tip occupies a large fraction of the frame;\n",
+    "        # default resolution provides sufficient detail for tip localisation.\n",
+    "        \"resolution\": 576,\n",
+    "    },\n",
+    "    \"human_pose\": {\n",
+    "        \"name\": \"Dataset Ridimensionato\",\n",
+    "        \"workspace\": \"poseestimation-wzidb\",\n",
+    "        \"project\": \"dataset-ridimensionato\",\n",
+    "        \"version\": 65,\n",
+    "        \"source_url\": \"https://universe.roboflow.com/poseestimation-wzidb/dataset-ridimensionato/dataset/65\",\n",
+    "        \"output_name\": \"keypoint_dataset_ridimensionato_demo\",\n",
+    "        \"format\": \"yolo\",\n",
+    "        \"keypoint_flip_pairs\": [],\n",
+    "        # 768 — body joints (wrists, ankles) require fine localisation across a range of\n",
+    "        # person scales; a larger canvas reduces positional error for distant people.\n",
+    "        \"resolution\": 768,\n",
+    "    },\n",
+    "    \"basketball_court\": {\n",
+    "        \"name\": \"Basketball Court Detection 2\",\n",
+    "        \"workspace\": \"roboflow-jvuqo\",\n",
+    "        \"project\": \"basketball-court-detection-2\",\n",
+    "        \"version\": 19,\n",
+    "        \"source_url\": \"https://universe.roboflow.com/roboflow-jvuqo/basketball-court-detection-2\",\n",
+    "        \"output_name\": \"keypoint_basketball_court_detection_demo\",\n",
+    "        \"keypoint_flip_pairs\": [],\n",
+    "        # 576 — court lines are coarse features that fill the frame; the default resolution\n",
+    "        # is sufficient to locate corners and centre-circle keypoints accurately.\n",
+    "        \"resolution\": 576,\n",
+    "    },\n",
+    "    \"football_field\": {\n",
+    "        \"name\": \"Football Field Detection\",\n",
+    "        \"workspace\": \"roboflow-jvuqo\",\n",
+    "        \"project\": \"football-field-detection-f07vi\",\n",
+    "        \"version\": 18,\n",
+    "        \"source_url\": \"https://universe.roboflow.com/roboflow-jvuqo/football-field-detection-f07vi\",\n",
+    "        \"output_name\": \"keypoint_football_field_detection_demo\",\n",
+    "        \"keypoint_flip_pairs\": [],\n",
+    "        # 768 — a full football pitch captured in a single frame compresses many line\n",
+    "        # intersections into a small pixel area; the higher resolution separates\n",
+    "        # closely spaced markings and reduces keypoint confusion near the centre circle.\n",
+    "        \"resolution\": 768,\n",
+    "    },\n",
+    "    \"tennis_court\": {\n",
+    "        \"name\": \"Tennis Court Keypoint\",\n",
+    "        \"workspace\": \"shj-pk1wt\",\n",
+    "        \"project\": \"tennis_court_keypoint\",\n",
+    "        \"version\": 6,\n",
+    "        \"source_url\": \"https://universe.roboflow.com/shj-pk1wt/tennis_court_keypoint\",\n",
+    "        \"output_name\": \"keypoint_tennis_court_demo\",\n",
+    "        \"format\": \"yolo\",\n",
+    "        \"keypoint_flip_pairs\": [],\n",
+    "        # 768 — a full court viewed from broadcast angles packs many line intersections\n",
+    "        # into a small pixel area; the higher resolution reduces confusion between\n",
+    "        # closely spaced baseline and service-box corners.\n",
+    "        \"resolution\": 768,\n",
+    "    },\n",
+    "}\n",
+    "\n",
+    "DATASET_KEY = \"dart\"\n",
+    "DATASET_INFO = DATASETS[DATASET_KEY]\n",
+    "DATASET_FORMAT = DATASET_INFO.get(\"format\", \"coco\")  # \"coco\" or \"yolo\"\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] | None = None\n",
+    "INFERENCE_COUNT = 6\n",
+    "INFERENCE_COLUMNS = 3\n",
+    "INFERENCE_THRESHOLD = 0.25\n",
+    "KEYPOINT_THRESHOLD = 0.1\n",
+    "PLOT_LOSS_LOG_SCALE = False\n",
+    "DRAW_UNCERTAINTY_ELLIPSES = True\n",
+    "ELLIPSE_SIGMA = 1.0\n",
+    "MAX_ELLIPSE_RADIUS = 36.0\n",
+    "\n",
+    "print(f\"dataset_key={DATASET_KEY}\")\n",
+    "print(f\"dataset={DATASET_INFO['name']}\")\n",
+    "print(f\"dataset_format={DATASET_FORMAT}\")\n",
+    "print(f\"source_url={DATASET_INFO['source_url']}\")\n",
+    "print(f\"resolution={RESOLUTION}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3d1ca469",
+   "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 notebook in Jupyter or Google Colab. If you are running this file as a plain Python script the function detects the absence of an IPython kernel and exits silently, so it is always safe to call."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b7920201",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "outputs": [],
+   "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()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "08de3f94",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "source": [
+    "## 1 - Download dataset\n",
+    "\n",
+    "Roboflow exports datasets in COCO keypoint format: a JSON file where each annotation contains a flat `keypoints` array of `[x, y, visibility, x, y, visibility, ...]` triplets, one per defined skeleton joint. The download cell fetches the exact dataset version listed in `DATASETS` and places it under `datasets/<DATASET_KEY>/`. You need a Roboflow API key — get one for free at app.roboflow.com/settings/api and set it as the `ROBOFLOW_API_KEY` environment variable (or as a Colab secret with the same name). The download is idempotent: if the target directory already exists Roboflow skips the network transfer, so re-running this cell after a successful download is fast."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b2109a1e",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "outputs": [],
+   "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",
+    "_rf_format = \"yolov8\" if DATASET_FORMAT == \"yolo\" else \"coco\"\n",
+    "dataset = (\n",
+    "    rf.workspace(str(DATASET_INFO[\"workspace\"]))\n",
+    "    .project(str(DATASET_INFO[\"project\"]))\n",
+    "    .version(int(DATASET_INFO[\"version\"]))\n",
+    "    .download(_rf_format, location=str(DATASETS_DIR / DATASET_KEY))\n",
+    ")\n",
+    "DATASET_DIR = Path(dataset.location)\n",
+    "print(f\"dataset_dir={DATASET_DIR}\")\n",
+    "print(f\"rf_format={_rf_format}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2f59cc9e",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "source": [
+    "## 2 - Infer keypoint schema\n",
+    "\n",
+    "`infer_coco_keypoint_schema` reads the training annotation JSON and extracts three things the model needs: the class names, the number of keypoints per class, and the OKS sigmas. OKS stands for Object Keypoint Similarity — it is the keypoint analogue of IoU and ranges from 0 to 1. Each sigma is a per-keypoint scale factor that controls how strictly localisation is penalised during evaluation: a smaller sigma (e.g. 0.025 for a wrist) means the model must predict that joint more precisely to score well, while a larger sigma (e.g. 0.107 for a hip) is more forgiving. `VALIDATE_KEYPOINT_METRICS` will be `False` when your dataset has categories with different keypoint counts, because the standard COCO OKS evaluator assumes a uniform skeleton across all instances; bounding-box mAP is always computed regardless."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3111b71a",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "outputs": [],
+   "source": [
+    "if DATASET_FORMAT == \"yolo\":\n",
+    "    schema = infer_yolo_keypoint_schema(DATASET_DIR / \"data.yaml\")\n",
+    "    _flip_idx = schema.flip_idx\n",
+    "    _pairs: list[int] = []\n",
+    "    _seen: set[int] = set()\n",
+    "    for _i, _j in enumerate(_flip_idx):\n",
+    "        if _i in _seen or _j in _seen or _i == _j:\n",
+    "            _seen.add(_i)\n",
+    "            continue\n",
+    "        if _j < len(_flip_idx) and _flip_idx[_j] == _i:\n",
+    "            _pairs.extend([_i, _j])\n",
+    "            _seen.update({_i, _j})\n",
+    "    KEYPOINT_FLIP_PAIRS = _pairs\n",
+    "else:\n",
+    "    TRAIN_ANNOTATIONS = DATASET_DIR / \"train\" / \"_annotations.coco.json\"\n",
+    "    schema = infer_coco_keypoint_schema(TRAIN_ANNOTATIONS)  # type: ignore[assignment]\n",
+    "    KEYPOINT_FLIP_PAIRS = list(DATASET_INFO.get(\"keypoint_flip_pairs\", []))\n",
+    "\n",
+    "CLASS_NAMES = schema.class_names\n",
+    "NUM_KEYPOINTS_PER_CLASS = schema.num_keypoints_per_class\n",
+    "NUM_CLASSES = len(CLASS_NAMES)\n",
+    "KEYPOINT_OKS_SIGMAS = schema.keypoint_oks_sigmas\n",
+    "ACTIVE_KEYPOINT_COUNTS = [count for count in NUM_KEYPOINTS_PER_CLASS if count > 0]\n",
+    "VALIDATE_KEYPOINT_METRICS = len(set(ACTIVE_KEYPOINT_COUNTS)) <= 1\n",
+    "\n",
+    "print(f\"class_names={CLASS_NAMES}\")\n",
+    "print(f\"num_keypoints_per_class={NUM_KEYPOINTS_PER_CLASS}\")\n",
+    "print(\"bbox_validation=True\")\n",
+    "print(f\"keypoint_oks_validation={VALIDATE_KEYPOINT_METRICS}\")\n",
+    "if VALIDATE_KEYPOINT_METRICS:\n",
+    "    print(f\"keypoint_oks_sigmas={len(KEYPOINT_OKS_SIGMAS)} values\")\n",
+    "else:\n",
+    "    print(\"keypoint_oks_validation=skipped because COCO OKS requires one keypoint count across categories\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4c9b5a6b",
+   "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 model config. This is separate from the PTL checkpoint because it can be loaded with a single `from_checkpoint` call on any machine, without reconstructing the original config."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e0bd03c1",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "outputs": [],
+   "source": [
+    "def _save_final_checkpoint(\n",
+    "    module: RFDETRModelModule,\n",
+    "    trainer: Any,\n",
+    "    train_config: KeypointTrainConfig,\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=\"RFDETRKeypointPreview\",\n",
+    "            model_config_dict=model_config.model_dump(),\n",
+    "        ),\n",
+    "        output_path,\n",
+    "    )\n",
+    "    return output_path"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f516409d",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "source": [
+    "`_annotated_keypoint_scene` renders keypoint dots and, when covariance data is available, uncertainty ellipses on top of an image. Drawing ellipses here rather than in the grid helper keeps the annotation logic in one place and makes it easy to toggle `DRAW_UNCERTAINTY_ELLIPSES` without touching the layout code."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "80f4317b",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "outputs": [],
+   "source": [
+    "def _annotated_keypoint_scene(image: np.ndarray, key_points: sv.KeyPoints) -> np.ndarray:\n",
+    "    \"\"\"Return an image with visible vertices and optional uncertainty ellipses.\"\"\"\n",
+    "    scene = image.copy()\n",
+    "    if DRAW_UNCERTAINTY_ELLIPSES and \"covariance\" in key_points.data:\n",
+    "        scene = sv.VertexEllipseAnnotator(\n",
+    "            sigma=ELLIPSE_SIGMA,\n",
+    "            color=sv.Color.ROBOFLOW,\n",
+    "            max_axis=MAX_ELLIPSE_RADIUS,\n",
+    "        ).annotate(scene=scene, key_points=key_points)\n",
+    "\n",
+    "    return sv.VertexAnnotator(radius=3).annotate(scene=scene, key_points=key_points)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "80f50cd8",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "source": [
+    "`_keypoint_grid_figure` arranges multiple annotated images into a fixed-column matplotlib grid. Using a grid rather than individual figures keeps the visual summary compact and makes it easy to compare predictions across images at a glance."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "51734ce1",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "outputs": [],
+   "source": [
+    "def _keypoint_grid_figure(items: list[tuple[str, np.ndarray, sv.KeyPoints]], columns: int) -> Figure:\n",
+    "    \"\"\"Render keypoint annotations in a fixed-column subplot grid.\"\"\"\n",
+    "    if columns <= 0:\n",
+    "        raise ValueError(f\"columns must be positive, got {columns}.\")\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, key_points) in zip(axes_array, items, strict=False):\n",
+    "        axis.imshow(_annotated_keypoint_scene(image, key_points))\n",
+    "        axis.set_title(title, fontsize=10)\n",
+    "        axis.axis(\"off\")\n",
+    "\n",
+    "    figure.tight_layout()\n",
+    "    return figure"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2e548c8c",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "source": [
+    "`_display_keypoint_records` prints per-image keypoint predictions as a pandas table. Fields that are constant across all keypoints in an image (such as the filename or detection score) are printed once as a header line rather than repeated in every row, keeping the output readable for images with many joints."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "757c151a",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "outputs": [],
+   "source": [
+    "def _display_keypoint_records(records: list[dict[str, Any]]) -> None:\n",
+    "    \"\"\"Display keypoint rows while printing per-image constant fields once.\"\"\"\n",
+    "    if not records:\n",
+    "        print(\"keypoint_rows=[]\")\n",
+    "        return\n",
+    "\n",
+    "    keypoint_columns = {\"detection_index\", \"keypoint_index\", \"x\", \"y\", \"keypoint_confidence\"}\n",
+    "    records_frame = pd.DataFrame(records)\n",
+    "    for image_name, image_frame in records_frame.groupby(\"image\", dropna=False, sort=False):\n",
+    "        constant_columns = [\n",
+    "            column\n",
+    "            for column in image_frame.columns\n",
+    "            if column not in keypoint_columns and image_frame[column].nunique(dropna=False) == 1\n",
+    "        ]\n",
+    "        if constant_columns:\n",
+    "            constants = image_frame.iloc[0][constant_columns]\n",
+    "            print(\", \".join(f\"{column}={constants[column]}\" for column in constant_columns))\n",
+    "        elif image_name is not None:\n",
+    "            print(f\"image={image_name}\")\n",
+    "\n",
+    "        display(image_frame.drop(columns=constant_columns).reset_index(drop=True))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "105aefc9",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "source": [
+    "## 3 - Configure training\n",
+    "\n",
+    "`KeypointTrainConfig` centralises every hyperparameter the training loop needs. A few fields are worth understanding before you customise them for your own dataset. `lr` and `lr_encoder` are usually set to the same value for fine-tuning; if you notice the backbone overfitting early you can halve `lr_encoder` to protect the pre-trained features while the detection head continues to adapt. `grad_accum_steps=2` means gradients are accumulated over two mini-batches before a weight update, which effectively doubles the batch size without requiring extra GPU memory — useful on smaller GPUs. `use_ema=False` keeps the demo fast; for a production run you can set this to `True` to maintain an exponential moving average of weights, which often improves final accuracy by a few tenths of a mAP point. `multi_scale=False` and `expanded_scales=False` disable the multi-resolution augmentation that RF-DETR uses during pre-training; turning them off shortens epoch time significantly and is fine for most fine-tuning runs. The `notes` dict is stored in the checkpoint alongside the weights, giving you a lightweight experiment log that travels with the model file."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fe060e29",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "outputs": [],
+   "source": [
+    "seed_all(SEED)\n",
+    "variant = RFDETRKeypointPreview(  # type: ignore[no-untyped-call]\n",
+    "    num_classes=NUM_CLASSES,\n",
+    "    num_keypoints_per_class=NUM_KEYPOINTS_PER_CLASS,\n",
+    "    resolution=RESOLUTION,\n",
+    ")\n",
+    "variant.model_config.model_name = type(variant).__name__\n",
+    "\n",
+    "train_config = KeypointTrainConfig(\n",
+    "    dataset_file=\"yolo\" if DATASET_FORMAT == \"yolo\" else \"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",
+    "    keypoint_flip_pairs=KEYPOINT_FLIP_PAIRS,\n",
+    "    keypoint_oks_sigmas=KEYPOINT_OKS_SIGMAS if VALIDATE_KEYPOINT_METRICS else None,\n",
+    "    notes={\n",
+    "        \"demo\": f\"keypoint-preview 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_keypoints_per_class\": NUM_KEYPOINTS_PER_CLASS,\n",
+    "        \"keypoint_flip_pairs\": KEYPOINT_FLIP_PAIRS,\n",
+    "        \"bbox_validation_enabled\": True,\n",
+    "        \"keypoint_oks_validation_enabled\": VALIDATE_KEYPOINT_METRICS,\n",
+    "        \"keypoint_oks_sigmas\": KEYPOINT_OKS_SIGMAS if VALIDATE_KEYPOINT_METRICS else None,\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)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4cdfa696",
+   "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 so you can confirm that keypoint skeletons are aligned with objects, that flips and colour jitter look reasonable, and that there are no systematic labelling errors such as swapped left/right joints or keypoints sitting outside their bounding box. Catching label noise here costs a few seconds; catching it after 50 epochs of training costs much more. If the preview shows empty images or missing annotations, the most common cause is a mismatch between the image filenames in the COCO JSON and the files on disk — check that the download completed fully."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "df126db6",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "outputs": [],
+   "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)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "76be07a3",
+   "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 computation, gradient accumulation, weight updates, learning-rate scheduling, and periodic validation. After each epoch the trainer appends a row to `metrics.csv` and, if validation mAP improves, saves a new best checkpoint via `BestModelCallback`. On a single consumer GPU (RTX 3090 or similar) 50 epochs on the dart dataset takes roughly 15–30 minutes depending on `BATCH_SIZE` and `NUM_WORKERS`. If you need to resume an interrupted run, set `train_config.resume` to the path of the last PTL checkpoint before calling this cell again."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "71d28853",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "outputs": [],
+   "source": [
+    "trainer.fit(model, datamodule=datamodule, ckpt_path=train_config.resume or None)\n",
+    "print(f\"output_dir={OUTPUT_DIR}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e1e72b80",
+   "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 not directly portable — they require the same class hierarchy to load. The `.pth` file written here is a self-contained RF-DETR checkpoint: it bundles the model weights together with the full training and model configs, including the keypoint schema and class names. This means you can share the file with a colleague and they can run inference with a single `RFDETRKeypointPreview.from_checkpoint(path)` call, with no need to reconstruct the original config. For full reproducibility, keep this checkpoint alongside the dataset version number printed in section 1."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5a052495",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "outputs": [],
+   "source": [
+    "final_checkpoint = _save_final_checkpoint(model, trainer, train_config, variant.model_config, FINAL_CHECKPOINT_PATH)\n",
+    "print(f\"saved_checkpoint_model={final_checkpoint}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "1482fa79",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "source": [
+    "## 7 - Validate metrics\n",
+    "\n",
+    "The validation pass that runs at the end of every training epoch uses the best weights seen so far and applies augmentation. This cell runs a clean post-training validation: no augmentation, the best checkpoint loaded from `checkpoint_best_total.pth` written by `BestModelCallback`, and results serialised to JSON for downstream comparison. Two metrics are most important to inspect here. `bbox/map` is the standard COCO bounding-box mAP at IoU 0.50:0.95 — it tells you how reliably the detector finds and bounds each object. `keypoint/oks` is the OKS-based mAP — it measures how precisely the model places each joint within the detected bounding boxes. A model can have a high `bbox/map` but a low `keypoint/oks` if it finds objects well but struggles to localise their joints; addressing that usually means more labelled data or tighter OKS sigmas.\n",
+    "\n",
+    "**Note:** `checkpoint_best_total.pth` is written after the first completed validation epoch. If training was interrupted before any epoch finished, this file will not exist and the cell below will raise `FileNotFoundError`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7aa26516",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "outputs": [],
+   "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",
+    "if not VALIDATE_KEYPOINT_METRICS:\n",
+    "    validation_metrics[\"keypoint_oks_skipped_mixed_keypoint_counts\"] = 1.0\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}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "0aa1a943",
+   "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 runs show both training loss and validation loss decreasing together; a large and growing gap between the two is a classic overfitting signal — the model is memorising the training set rather than generalising. For the mAP curves, on small datasets (a few hundred images) you typically see rapid improvement in the first 10–20 epochs followed by a plateau around epoch 30–50; if mAP is still climbing at epoch 50, consider extending `EPOCHS`. Set `PLOT_LOSS_LOG_SCALE = True` if the loss drops by an order of magnitude in the first few epochs and the later, more meaningful portion of the curve gets compressed into a flat line at the bottom of the plot."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "37589b33",
+   "metadata": {},
+   "outputs": [],
+   "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)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a5608045",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "outputs": [],
+   "source": [
+    "map_figure = plot_map_metrics(str(METRICS_CSV))\n",
+    "display(map_figure)\n",
+    "plt.close(map_figure)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "690a200a",
+   "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 stored config from the `.pth` file, so the reconstructed model has the correct number of classes and keypoints without you having to pass them explicitly. You can share this checkpoint file with teammates and they can run inference immediately — the keypoint schema, class names, and OKS sigmas are all embedded in the file alongside the weights. This cell also confirms that the save-load round trip works before you proceed to inference, so any corruption or version mismatch surfaces here rather than silently producing wrong predictions later."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c8d75b0f",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "outputs": [],
+   "source": [
+    "loaded_model = RFDETRKeypointPreview.from_checkpoint(FINAL_CHECKPOINT_PATH)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7bee8d28",
+   "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 images to run inference on even when the dataset has no dedicated test split. Using images from the test set gives you an unbiased view of model performance because those images were never seen during training or used to pick the best checkpoint. If you want to run inference on your own images, replace `inference_image_paths` with a list of `Path` objects pointing to your files before running the next cell."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a5182f47",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "outputs": [],
+   "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",
+    "    # Iterate through all subfolders to find images\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(path) for path in inference_image_paths]}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "1c7654a9",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "source": [
+    "## 11 - Visualize inference\n",
+    "\n",
+    "Two thresholds control what you see here. `INFERENCE_THRESHOLD` is a detection confidence threshold: only detections whose bounding-box score exceeds this value are shown. Raise it to suppress false positives; lower it to surface low-confidence detections. `KEYPOINT_THRESHOLD` operates within each accepted detection: individual keypoints whose confidence falls below this value are treated as not visible and omitted from the overlay. The uncertainty ellipses (drawn when `DRAW_UNCERTAINTY_ELLIPSES = True`) visualise the predicted covariance for each keypoint — a larger ellipse means the model is less certain about that joint's exact location."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e520b84c",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "outputs": [],
+   "source": [
+    "inference_grid_items = []\n",
+    "keypoint_rows = []\n",
+    "for image_path in inference_image_paths:\n",
+    "    with Image.open(image_path) as image:\n",
+    "        image = image.convert(\"RGB\")\n",
+    "        key_points_raw = loaded_model.predict(image, threshold=INFERENCE_THRESHOLD)\n",
+    "    if not isinstance(key_points_raw, sv.KeyPoints):\n",
+    "        raise RuntimeError(\n",
+    "            f\"Expected RFDETRKeypointPreview.predict() to return sv.KeyPoints, got {type(key_points_raw)!r}.\"\n",
+    "        )\n",
+    "\n",
+    "    key_points = _key_points_for_display(key_points_raw, keypoint_threshold=KEYPOINT_THRESHOLD)\n",
+    "    rows = _keypoint_prediction_records(key_points, image=image_path, keypoint_threshold=KEYPOINT_THRESHOLD)\n",
+    "    inference_grid_items.append((image_path.name, np.array(image), key_points))\n",
+    "    keypoint_rows.extend(rows)\n",
+    "\n",
+    "if inference_grid_items:\n",
+    "    figure = _keypoint_grid_figure(inference_grid_items, columns=INFERENCE_COLUMNS)\n",
+    "    display(figure)\n",
+    "    plt.close(figure)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "93df94eb",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "source": [
+    "## 12 - Keypoint prediction table\n",
+    "\n",
+    "The table below gives you the raw coordinates and per-keypoint confidence scores for every detection in the grid above. Fields that are constant across all keypoints in an image (filename, detection score) are printed once as a header line rather than repeated in every row. This format is easy to copy into a spreadsheet or feed into downstream analysis."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "de9b2707",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "_display_keypoint_records(keypoint_rows)"
+   ]
+  }
+ ],
+ "metadata": {
+  "jupytext": {
+   "cell_metadata_filter": "title,-all",
+   "main_language": "python",
+   "notebook_metadata_filter": "-all"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

docs/cookbooks/fine-tune_segmentation.ipynb

@@ -0,0 +1,846 @@
+{
+ "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
+}

docs/cookbooks/inference-latency-benchmark.ipynb

@@ -0,0 +1,375 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a5720dea",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# ------------------------------------------------------------------------\n",
+    "# RF-DETR\n",
+    "# Copyright (c) 2025 Roboflow. All Rights Reserved.\n",
+    "# Licensed under the Apache License, Version 2.0 [see LICENSE for details]\n",
+    "# ------------------------------------------------------------------------"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "66c62889",
+   "metadata": {},
+   "source": [
+    "# RF-DETR Inference Latency Benchmark\n",
+    "\n",
+    "Measures inference latency for three RF-DETR families across three configs:\n",
+    "\n",
+    "| Config | Description |\n",
+    "|--------|-------------|\n",
+    "| **FP32** | `predict()` — unoptimized baseline |\n",
+    "| **FP16+JIT** | `optimize_for_inference(dtype=torch.float16)` |\n",
+    "| **ONNX** | exported `.onnx` via `onnxruntime-gpu` |"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d1568b82",
+   "metadata": {},
+   "source": [
+    "## 1. Install\n",
+    "\n",
+    "We need `onnxruntime-gpu` built against CUDA 12 — the default PyPI wheel targets CUDA 11.8 and silently\n",
+    "falls back to CPU on modern GPUs.  The Microsoft CUDA-12 package index ships the correct build.\n",
+    "\n",
+    "> **Colab**: after running this cell, go to **Runtime → Restart session**, then run from the next cell."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6a116013",
+   "metadata": {
+    "title": "[bash]"
+   },
+   "outputs": [],
+   "source": [
+    "!pip uninstall -y onnxruntime onnxruntime-gpu\n",
+    "!pip install -q \"rfdetr[onnx]\" pillow pandas\n",
+    "!pip install -q onnxruntime-gpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "1ca0838b",
+   "metadata": {},
+   "source": [
+    "## 2. Config\n",
+    "\n",
+    "`WARMUP_RUNS` discards the first N inferences — GPU kernels are JIT-compiled on first use, so early timings\n",
+    "are outliers.  `MEASURE_RUNS` then collects the steady-state distribution.  20 + 100 is a reasonable balance\n",
+    "between statistical stability and total wall-clock time per model."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0e47b983",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from collections.abc import Callable\n",
+    "from pathlib import Path\n",
+    "from typing import Any, NamedTuple\n",
+    "\n",
+    "import numpy as np\n",
+    "import torch\n",
+    "from PIL import Image\n",
+    "\n",
+    "WARMUP_RUNS = 20\n",
+    "MEASURE_RUNS = 100\n",
+    "EXPORT_DIR = Path(\"benchmark_output\")\n",
+    "EXPORT_DIR.mkdir(exist_ok=True)\n",
+    "\n",
+    "if not torch.cuda.is_available():\n",
+    "    raise RuntimeError(\"This benchmark requires a CUDA GPU.\")\n",
+    "print(f\"GPU: {torch.cuda.get_device_name(0)}\")\n",
+    "\n",
+    "import onnxruntime as ort\n",
+    "\n",
+    "_ort_providers = ort.get_available_providers()\n",
+    "print(f\"ORT {ort.__version__}, providers: {_ort_providers}\")\n",
+    "if \"CUDAExecutionProvider\" not in _ort_providers:\n",
+    "    raise RuntimeError(\n",
+    "        f\"onnxruntime-gpu with CUDA support required. Available providers: {_ort_providers}. \"\n",
+    "        \"Fix: reinstall from the CUDA-12 index (see install cell) then restart runtime.\"\n",
+    "    )"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2ccdc87e",
+   "metadata": {},
+   "source": [
+    "## 3. Sample images\n",
+    "\n",
+    "Latency depends on resolution, not pixel content, so synthetic noise images are equivalent to real photos\n",
+    "for benchmarking purposes.  Using a fixed seed makes results reproducible across runs."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "06a16bb0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "rng = np.random.default_rng(42)\n",
+    "images: list[Image.Image] = [Image.fromarray(rng.integers(0, 256, (640, 640, 3), dtype=np.uint8)) for _ in range(10)]\n",
+    "print(f\"Generated {len(images)} synthetic 640×640 RGB images\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "db0d3d7c",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "source": [
+    "## 4. Latency helpers\n",
+    "\n",
+    "GPU kernels execute asynchronously — `time.perf_counter()` returns before the GPU finishes, giving\n",
+    "misleadingly low numbers.  CUDA events are inserted directly into the GPU command stream and timestamped\n",
+    "on the device, so `elapsed_time()` measures actual kernel execution.  `torch.cuda.synchronize()` after\n",
+    "each run flushes the stream and ensures the event fires before we read it."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7f928319",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class BenchmarkResult(NamedTuple):\n",
+    "    \"\"\"Single benchmark measurement.\"\"\"\n",
+    "\n",
+    "    label: str\n",
+    "    mean_ms: float\n",
+    "    std_ms: float\n",
+    "\n",
+    "    @property\n",
+    "    def fps(self) -> float:\n",
+    "        \"\"\"Frames per second.\"\"\"\n",
+    "        return 1000.0 / self.mean_ms\n",
+    "\n",
+    "\n",
+    "def measure_latency_gpu(\n",
+    "    fn: Callable[[], object],\n",
+    "    warmup: int = WARMUP_RUNS,\n",
+    "    runs: int = MEASURE_RUNS,\n",
+    ") -> tuple[float, float]:\n",
+    "    \"\"\"Return (mean_ms, std_ms) using CUDA events.\"\"\"\n",
+    "    for _ in range(warmup):\n",
+    "        fn()\n",
+    "    torch.cuda.synchronize()\n",
+    "    start = torch.cuda.Event(enable_timing=True)\n",
+    "    end = torch.cuda.Event(enable_timing=True)\n",
+    "    timings: list[float] = []\n",
+    "    for _ in range(runs):\n",
+    "        start.record()\n",
+    "        fn()\n",
+    "        end.record()\n",
+    "        torch.cuda.synchronize()\n",
+    "        timings.append(start.elapsed_time(end))\n",
+    "    arr = np.array(timings)\n",
+    "    return float(arr.mean()), float(arr.std())\n",
+    "\n",
+    "\n",
+    "_measure = measure_latency_gpu"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "8f7ee4f5",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "source": [
+    "## 5. Per-config benchmark functions\n",
+    "\n",
+    "Three inference paths are compared:\n",
+    "\n",
+    "- **FP32** — `predict()` as shipped.  Full 32-bit arithmetic on GPU.  Baseline.\n",
+    "- **FP16+JIT** — `optimize_for_inference(dtype=torch.float16)` fuses layers with `torch.jit.script` and\n",
+    "  halves the arithmetic precision.  Typically 1.5–2× faster than FP32 with negligible accuracy loss on\n",
+    "  modern tensor cores.\n",
+    "- **ONNX** — the model is exported to the Open Neural Network Exchange format and run through ONNX Runtime,\n",
+    "  bypassing PyTorch entirely.  ORT applies its own graph optimisations and can use the TensorRT execution\n",
+    "  provider for additional speedup.  Benchmarked separately on CPU and GPU to show the provider impact."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "533f037b",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "outputs": [],
+   "source": [
+    "from rfdetr.export._onnx.inference import _onnx_runtime\n",
+    "\n",
+    "\n",
+    "def _predict_fp32(model: Any, image: Image.Image) -> BenchmarkResult:\n",
+    "    \"\"\"Baseline FP32 predict() latency.\"\"\"\n",
+    "    mean, std = _measure(lambda: model.predict(image))\n",
+    "    return BenchmarkResult(\"predict() FP32\", mean, std)\n",
+    "\n",
+    "\n",
+    "def _predict_fp16(model: Any, image: Image.Image) -> BenchmarkResult:\n",
+    "    \"\"\"FP16+JIT latency — applies and removes optimize_for_inference.\"\"\"\n",
+    "    model.optimize_for_inference(dtype=torch.float16)\n",
+    "    mean, std = _measure(lambda: model.predict(image))\n",
+    "    model.remove_optimized_model()\n",
+    "    return BenchmarkResult(\"predict() FP16+JIT\", mean, std)\n",
+    "\n",
+    "\n",
+    "def _export_onnx(model: Any, export_dir: Path) -> Path:\n",
+    "    \"\"\"Export model to ONNX and return the path.\"\"\"\n",
+    "    return Path(model.export(output_dir=str(export_dir)))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2f6bc92a",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "source": [
+    "## 6. Model benchmark runner\n",
+    "\n",
+    "Each model is loaded fresh, exported once, and then each inference config is timed independently.\n",
+    "The ONNX path reuses the same exported file for both CPU and CUDA providers so export cost is not\n",
+    "counted in latency.  CUDA ONNX is skipped gracefully when no GPU is available; missing `onnxruntime-gpu`\n",
+    "raises immediately so misconfigured environments are caught early."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8afd179a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def run_model_benchmark(\n",
+    "    model_cls: type,\n",
+    "    model_name: str,\n",
+    "    images: list[Image.Image],\n",
+    ") -> list[BenchmarkResult]:\n",
+    "    \"\"\"Run FP32 / FP16 / ONNX benchmarks for one model and print results.\"\"\"\n",
+    "    print(f\"\\n{'=' * 62}\")\n",
+    "    print(f\"  {model_name}\")\n",
+    "    print(\"=\" * 62)\n",
+    "\n",
+    "    model: Any = model_cls()\n",
+    "    export_dir = EXPORT_DIR / model_name.split()[0]\n",
+    "    export_dir.mkdir(exist_ok=True)\n",
+    "    image = images[0]\n",
+    "\n",
+    "    fp32 = _predict_fp32(model, image)\n",
+    "    fp16 = _predict_fp16(model, image)\n",
+    "    results: list[BenchmarkResult] = [fp32, fp16]\n",
+    "\n",
+    "    onnx_path = _export_onnx(model, export_dir)\n",
+    "    for providers in ([\"CPUExecutionProvider\"], [\"CUDAExecutionProvider\", \"CPUExecutionProvider\"]):\n",
+    "        if providers[0] == \"CUDAExecutionProvider\" and not torch.cuda.is_available():\n",
+    "            print(\"  ⚠ ONNX (CUDA) skipped — no CUDA GPU\")\n",
+    "            continue\n",
+    "        mean_ms, std_ms, label = _onnx_runtime(onnx_path, image, providers, WARMUP_RUNS, MEASURE_RUNS)\n",
+    "        results.append(BenchmarkResult(f\"ONNX ({label})\", mean_ms, std_ms))\n",
+    "\n",
+    "    for r in results:\n",
+    "        print(f\"  {r.label:<30}  {r.mean_ms:6.2f} ms ± {r.std_ms:5.2f}   ({r.fps:6.1f} FPS)\")\n",
+    "\n",
+    "    onnx_results = [r for r in results if r.label.startswith(\"ONNX\")]\n",
+    "    speedups = [f\"FP16 {fp32.mean_ms / fp16.mean_ms:.1f}×\"]\n",
+    "    if onnx_results:\n",
+    "        speedups.append(f\"ONNX {fp32.mean_ms / onnx_results[0].mean_ms:.1f}×\")\n",
+    "    print(f\"  Speedup vs FP32: {' | '.join(speedups)}\")\n",
+    "\n",
+    "    return results"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e7ccd41c",
+   "metadata": {},
+   "source": [
+    "## 7. Benchmark loop — detection · segmentation · keypoint\n",
+    "\n",
+    "Three model families are benchmarked to show how task complexity affects latency.  Detection (`RFDETRMedium`)\n",
+    "outputs boxes only; segmentation (`RFDETRSegSmall`) additionally predicts per-object masks, which adds\n",
+    "decoder cost; keypoint (`RFDETRKeypointPreview`) predicts skeleton joints and is typically the lightest\n",
+    "of the three at smaller resolutions."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2ab7ffbc",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from rfdetr import RFDETRKeypointPreview, RFDETRMedium, RFDETRSegSmall\n",
+    "\n",
+    "MODELS: list[tuple[type, str]] = [\n",
+    "    (RFDETRMedium, \"RFDETRMedium — detection\"),\n",
+    "    (RFDETRSegSmall, \"RFDETRSegSmall — segmentation\"),\n",
+    "    (RFDETRKeypointPreview, \"RFDETRKeypointPreview — keypoint\"),\n",
+    "]\n",
+    "\n",
+    "all_results: dict[str, list[BenchmarkResult]] = {}\n",
+    "for _model_cls, _model_name in MODELS:\n",
+    "    all_results[_model_name] = run_model_benchmark(_model_cls, _model_name, images)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ebbba970",
+   "metadata": {},
+   "source": [
+    "## 8. Summary\n",
+    "\n",
+    "The table shows FPS (frames per second) for each model × config combination.  Higher is better.\n",
+    "Compare columns to see which model fits your latency budget; compare rows to choose the right\n",
+    "inference backend for your deployment target (Python server, edge CPU, or ONNX Runtime service)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5d8eddd9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import pandas as pd\n",
+    "\n",
+    "summary = {\n",
+    "    model_name.split()[0]: {r.label: round(r.fps, 1) for r in results} for model_name, results in all_results.items()\n",
+    "}\n",
+    "df = pd.DataFrame(summary)\n",
+    "df.index.name = \"Config \\\\ Model\"\n",
+    "print(df.to_string())\n",
+    "print(f\"\\nFPS — {MEASURE_RUNS} timed + {WARMUP_RUNS} warmup runs, batch 1, GPU: {torch.cuda.get_device_name(0)}.\")"
+   ]
+  }
+ ],
+ "metadata": {
+  "jupytext": {
+   "cell_metadata_filter": "title,-all",
+   "main_language": "python",
+   "notebook_metadata_filter": "-all"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

docs/cookbooks/pytorch-lightning.ipynb

@@ -0,0 +1,451 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "1777c5f9",
+   "metadata": {},
+   "source": [
+    "# Training with PyTorch Lightning\n",
+    "\n",
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/roboflow/rf-detr/blob/develop/docs/cookbooks/pytorch-lightning.ipynb)\n",
+    "\n",
+    "**RF-DETR** is a real-time object detection model that combines the accuracy of\n",
+    "transformer-based detectors with inference speeds suitable for production.\n",
+    "The training stack is built on **PyTorch Lightning** — giving you\n",
+    "composable building blocks you can adopt incrementally, without changing your\n",
+    "existing code.\n",
+    "\n",
+    "| Building block | Role |\n",
+    "|---|---|\n",
+    "| `RFDETRModelModule` | `LightningModule` — model, loss, optimizer, scheduler |\n",
+    "| `RFDETRDataModule` | `LightningDataModule` — datasets and dataloaders |\n",
+    "| `build_trainer()` | Factory that assembles a `Trainer` with all RF-DETR callbacks |\n",
+    "\n",
+    "**Key design principle:** start simple, then pick up building blocks without losing\n",
+    "your trained weights.\n",
+    "\n",
+    "- **Phase 1** — `model.train()` one-liner (`EPOCHS_PHASE_1` epochs)\n",
+    "- **Phase 2** — swap in the PTL components and continue for `EPOCHS_PHASE_2` more\n",
+    "  epochs from the same checkpoint, same output folder — no conversion required\n",
+    "- **End** — full training curve, single-image inference, and batch inference with\n",
+    "  `trainer.predict()`"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a750078f",
+   "metadata": {},
+   "source": [
+    "## 1. Install RF-DETR 1.6.0\n",
+    "\n",
+    "`rfdetr[train,loggers]` pulls in PyTorch Lightning, torchmetrics, and the full\n",
+    "callback stack. `roboflow` downloads the demo dataset."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a4688278",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!pip install -q rfdetr[train,loggers]==1.6.0 roboflow"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "77fab4e5",
+   "metadata": {},
+   "source": [
+    "## 2. Config\n",
+    "\n",
+    "All notebook-level knobs in one place. Adjust `EPOCHS_PHASE_*` and `BATCH_SIZE`\n",
+    "to match your hardware — every downstream cell reads from these variables.\n",
+    "\n",
+    "`num_workers` is set to `os.cpu_count()` inside a Jupyter/Colab kernel where\n",
+    "process forking is safe, and to `0` when running as a plain Python script.\n",
+    "On macOS and Windows, spawn-based multiprocessing would otherwise re-import\n",
+    "this module as `__main__` and retrigger training."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "174445e1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "from pathlib import Path\n",
+    "\n",
+    "DATASET_DIR = os.environ.get(\"DATASET_DIR\", \"\")\n",
+    "OUTPUT_DIR = \"output\"\n",
+    "EPOCHS_PHASE_1 = 20\n",
+    "EPOCHS_PHASE_2 = 10\n",
+    "BATCH_SIZE = 12\n",
+    "THRESHOLD = 0.3\n",
+    "\n",
+    "os.makedirs(OUTPUT_DIR, exist_ok=True)\n",
+    "\n",
+    "try:\n",
+    "    from IPython import get_ipython\n",
+    "\n",
+    "    _in_notebook = get_ipython() is not None\n",
+    "except Exception:\n",
+    "    _in_notebook = False\n",
+    "\n",
+    "# Outside a notebook kernel, macOS/Windows spawn-based multiprocessing will\n",
+    "# re-import this script as __main__, triggering training again.\n",
+    "# Use 0 workers in that case; inside a kernel the usual forking rules apply safely.\n",
+    "num_workers = (os.cpu_count() or 0) if _in_notebook else 0"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3a43c723",
+   "metadata": {},
+   "source": [
+    "## 3. Dataset\n",
+    "\n",
+    "[Aquarium Combined](https://universe.roboflow.com/brad-dwyer/aquarium-combined)\n",
+    "— 638 images across 7 classes (`fish`, `jellyfish`, `penguin`, `puffin`,\n",
+    "`shark`, `starfish`, `stingray`).  It is small enough to complete a demo run\n",
+    "in a few minutes yet diverse enough to produce meaningful detection results.\n",
+    "\n",
+    "Set `ROBOFLOW_API_KEY` as a Colab secret (Secrets panel, key icon) or as an\n",
+    "environment variable before running this cell.  The dataset is downloaded in\n",
+    "COCO format, which RF-DETR reads natively."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "34478ada",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import json\n",
+    "\n",
+    "from roboflow import Roboflow\n",
+    "\n",
+    "try:\n",
+    "    from google.colab import userdata  # type: ignore[import]\n",
+    "\n",
+    "    API_KEY = userdata.get(\"ROBOFLOW_API_KEY\")\n",
+    "except Exception:\n",
+    "    API_KEY = os.environ[\"ROBOFLOW_API_KEY\"]\n",
+    "\n",
+    "rf = Roboflow(api_key=API_KEY)\n",
+    "dataset = rf.workspace(\"brad-dwyer\").project(\"aquarium-combined\").version(1).download(\"coco\", location=\"datasets\")\n",
+    "DATASET_DIR = dataset.location\n",
+    "\n",
+    "with open(Path(DATASET_DIR) / \"train\" / \"_annotations.coco.json\") as f:\n",
+    "    _ann = json.load(f)\n",
+    "\n",
+    "CLASS_NAMES = [c[\"name\"] for c in sorted(_ann[\"categories\"], key=lambda c: c[\"id\"])]\n",
+    "NUM_CLASSES = len(CLASS_NAMES)\n",
+    "print(f\"Dataset : {DATASET_DIR}\")\n",
+    "print(f\"Classes : {NUM_CLASSES} — {CLASS_NAMES}\")\n",
+    "\n",
+    "with open(Path(DATASET_DIR) / \"valid\" / \"_annotations.coco.json\") as f:\n",
+    "    _val_ann = json.load(f)\n",
+    "\n",
+    "val_images_dir = Path(DATASET_DIR) / \"valid\"\n",
+    "val_image_files = [img[\"file_name\"] for img in _val_ann[\"images\"]]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "44bbcce0",
+   "metadata": {},
+   "source": [
+    "## 4. Phase 1 — `model.train()` one-liner\n",
+    "\n",
+    "The same high-level API that has been in RF-DETR since v1.0 — nothing here\n",
+    "changes from previous releases.  If you have existing training scripts, they\n",
+    "keep working without modification.\n",
+    "\n",
+    "`pretrain_weights=\"rf-detr-medium.pth\"` downloads COCO-pretrained backbone\n",
+    "weights automatically on first run and caches them locally.  `use_ema=True`\n",
+    "maintains an exponential moving average of the weights to stabilise validation\n",
+    "metrics.  `run_test=False` skips the final test-set evaluation to keep Phase 1\n",
+    "fast; Phase 2 turns it back on.\n",
+    "\n",
+    "After this cell completes, `OUTPUT_DIR/checkpoint_best_total.pth` holds the\n",
+    "best weights seen so far — the starting point for Phase 2."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "bf7b136e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from rfdetr import RFDETRMedium\n",
+    "\n",
+    "model = RFDETRMedium(num_classes=NUM_CLASSES, pretrain_weights=\"rf-detr-medium.pth\")\n",
+    "model.train(\n",
+    "    dataset_dir=DATASET_DIR,\n",
+    "    epochs=EPOCHS_PHASE_1,\n",
+    "    batch_size=BATCH_SIZE,\n",
+    "    grad_accum_steps=4,\n",
+    "    lr=1e-4,\n",
+    "    num_workers=num_workers,\n",
+    "    output_dir=OUTPUT_DIR,\n",
+    "    use_ema=True,\n",
+    "    run_test=False,\n",
+    "    progress_bar=\"rich\",\n",
+    "    tensorboard=True,\n",
+    "    seed=42,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "178dd044",
+   "metadata": {},
+   "source": [
+    "## 5. Phase 2 — PTL building blocks\n",
+    "\n",
+    "Pick up the three PTL components and call `trainer.fit()` pointing at the Phase 1\n",
+    "checkpoint.  No weight conversion is needed — `RFDETRModelModule.on_load_checkpoint`\n",
+    "detects the `.pth` format and remaps keys automatically.\n",
+    "\n",
+    "A lower learning rate (`5e-5`) is used because the model is already partially\n",
+    "converged.  `epochs=EPOCHS_PHASE_1 + EPOCHS_PHASE_2` sets the *absolute* epoch\n",
+    "ceiling; because the loaded checkpoint records the last completed epoch, PTL\n",
+    "runs exactly `EPOCHS_PHASE_2` additional epochs before stopping.  The same\n",
+    "`OUTPUT_DIR` is reused so checkpoints and metrics all land in one place."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d4526ab0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import pandas as pd\n",
+    "\n",
+    "from rfdetr import RFDETRDataModule, RFDETRModelModule, build_trainer\n",
+    "from rfdetr.config import RFDETRMediumConfig, TrainConfig\n",
+    "\n",
+    "# Read Phase 1 metrics before Phase 2 overwrites the CSV.\n",
+    "df1 = pd.read_csv(f\"{OUTPUT_DIR}/metrics.csv\")\n",
+    "\n",
+    "model_config = RFDETRMediumConfig(\n",
+    "    num_classes=NUM_CLASSES,\n",
+    "    pretrain_weights=\"rf-detr-medium.pth\",\n",
+    ")\n",
+    "\n",
+    "# epochs = EPOCHS_1 + EPOCHS_2 so PTL (which resumes the epoch counter from the\n",
+    "# checkpoint) runs exactly EPOCHS_2 additional epochs before reaching max_epochs.\n",
+    "train_config = TrainConfig(\n",
+    "    dataset_dir=DATASET_DIR,\n",
+    "    epochs=EPOCHS_PHASE_1 + EPOCHS_PHASE_2,\n",
+    "    batch_size=BATCH_SIZE,\n",
+    "    grad_accum_steps=4,\n",
+    "    lr=5e-5,\n",
+    "    num_workers=num_workers,\n",
+    "    output_dir=OUTPUT_DIR,\n",
+    "    use_ema=True,\n",
+    "    run_test=True,\n",
+    "    progress_bar=\"tqdm\",\n",
+    "    tensorboard=True,\n",
+    "    seed=42,\n",
+    ")\n",
+    "\n",
+    "module = RFDETRModelModule(model_config=model_config, train_config=train_config)\n",
+    "datamodule = RFDETRDataModule(model_config=model_config, train_config=train_config)\n",
+    "trainer = build_trainer(train_config, model_config)\n",
+    "\n",
+    "# Resume directly from the Phase 1 .pth — no conversion needed.\n",
+    "trainer.fit(module, datamodule, ckpt_path=f\"{OUTPUT_DIR}/checkpoint_best_total.pth\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "344147ee",
+   "metadata": {},
+   "source": [
+    "## 6. Training curve\n",
+    "\n",
+    "Phase 1 and Phase 2 each emit their own `metrics.csv` (Phase 2 overwrites Phase 1's\n",
+    "file when it starts).  We captured the Phase 1 copy before fitting so we can\n",
+    "concatenate both DataFrames and plot a single continuous curve across all epochs."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "751e8277",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%matplotlib inline\n",
+    "\n",
+    "import matplotlib.pyplot as plt\n",
+    "import numpy as np\n",
+    "import supervision as sv\n",
+    "from IPython.display import display\n",
+    "from PIL import Image\n",
+    "\n",
+    "from rfdetr.visualize.training import plot_metrics\n",
+    "\n",
+    "df2 = pd.read_csv(f\"{OUTPUT_DIR}/metrics.csv\")\n",
+    "\n",
+    "combined_csv = f\"{OUTPUT_DIR}/metrics_combined.csv\"\n",
+    "pd.concat([df1, df2], ignore_index=True).to_csv(combined_csv, index=False)\n",
+    "print(f\"Combined CSV: {combined_csv}  ({len(df1) + len(df2)} rows)\")\n",
+    "\n",
+    "fig = plot_metrics(combined_csv)\n",
+    "display(fig)\n",
+    "print(f\"Saved: {OUTPUT_DIR}/metrics_plot.png\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d3f29ab1",
+   "metadata": {},
+   "source": [
+    "## 7. Single-image inference — `model.predict()`\n",
+    "\n",
+    "`RFDETRMedium` can be instantiated directly from a checkpoint path for inference\n",
+    "— no training config needed.  `model.predict()` accepts a PIL `Image`, runs\n",
+    "preprocessing, the forward pass, and postprocessing internally, and returns a\n",
+    "`supervision.Detections` object that is ready to annotate and display."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "90e06abb",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%matplotlib inline\n",
+    "\n",
+    "model = RFDETRMedium(pretrain_weights=f\"{OUTPUT_DIR}/checkpoint_best_total.pth\", num_classes=NUM_CLASSES)\n",
+    "\n",
+    "image = Image.open(val_images_dir / val_image_files[0])\n",
+    "detections = model.predict(image, threshold=THRESHOLD)\n",
+    "\n",
+    "annotated = sv.BoxAnnotator().annotate(image.copy(), detections)\n",
+    "annotated = sv.LabelAnnotator().annotate(annotated, detections, labels=[CLASS_NAMES[c] for c in detections.class_id])\n",
+    "\n",
+    "plt.figure(figsize=(10, 7))\n",
+    "plt.imshow(np.array(annotated))\n",
+    "plt.axis(\"off\")\n",
+    "plt.show()\n",
+    "\n",
+    "print(f\"Detected {len(detections)} object(s)\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "32baa2d8",
+   "metadata": {},
+   "source": [
+    "## 8. Batch inference — `trainer.predict()`\n",
+    "\n",
+    "Instead of calling `model.predict()` one image at a time, `trainer.predict()`\n",
+    "streams the entire validation set through the model in batches and collects\n",
+    "all results — useful for dataset-level evaluation or offline export pipelines.\n",
+    "\n",
+    "**What happens under the hood:**\n",
+    "\n",
+    "1. PTL calls `datamodule.setup(\"predict\")` — this builds `_dataset_val` if it\n",
+    "   does not exist yet.  Because `trainer.fit()` already ran above, the dataset\n",
+    "   is already in memory and `setup` is a no-op.\n",
+    "2. PTL calls `datamodule.predict_dataloader()` — this returns the *validation*\n",
+    "   dataset wrapped in a `SequentialSampler` (no shuffle, no augmentation),\n",
+    "   identical to `val_dataloader`.\n",
+    "3. For each batch, `RFDETRModelModule.predict_step()` runs a forward pass under\n",
+    "   `torch.no_grad()` and returns a list of `{\"scores\", \"labels\", \"boxes\"}`\n",
+    "   dicts — one dict per image in the batch.\n",
+    "4. `trainer.predict()` collects all batch results into a\n",
+    "   `List[List[dict]]` (outer = batches, inner = images).\n",
+    "\n",
+    "Flatten, apply a confidence threshold, and wrap in `sv.Detections`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7d7543ff",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%matplotlib inline\n",
+    "\n",
+    "import itertools\n",
+    "\n",
+    "# Returns List[List[dict]] — outer: batches, inner: one dict per image\n",
+    "# Each dict has keys: \"scores\" (N,), \"labels\" (N,), \"boxes\" (N, 4) — all tensors\n",
+    "all_preds = trainer.predict(module, datamodule)\n",
+    "\n",
+    "# Flatten the batch dimension → one result dict per validation image\n",
+    "flat_preds = [img_result for batch in all_preds for img_result in batch]\n",
+    "print(f\"Ran predict on {len(flat_preds)} validation images\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ed942d37",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Build sv.Detections from raw tensors and visualise the first four images\n",
+    "annotated_images = []\n",
+    "for img_file, result in itertools.islice(zip(val_image_files, flat_preds), 4):\n",
+    "    keep = result[\"scores\"] > THRESHOLD\n",
+    "    detections = sv.Detections(\n",
+    "        xyxy=result[\"boxes\"][keep].cpu().float().numpy(),\n",
+    "        confidence=result[\"scores\"][keep].cpu().float().numpy(),\n",
+    "        class_id=result[\"labels\"][keep].cpu().long().numpy(),\n",
+    "    )\n",
+    "    image = Image.open(val_images_dir / img_file)\n",
+    "    annotated = sv.BoxAnnotator().annotate(image.copy(), detections)\n",
+    "    annotated = sv.LabelAnnotator().annotate(\n",
+    "        annotated, detections, labels=[CLASS_NAMES[c] for c in detections.class_id]\n",
+    "    )\n",
+    "    annotated_images.append(np.array(annotated))\n",
+    "    print(f\"  {img_file}: {len(detections)} detection(s)\")\n",
+    "\n",
+    "fig, axes = plt.subplots(2, 2, figsize=(14, 10))\n",
+    "for ax, img in zip(axes.flat, annotated_images):\n",
+    "    ax.imshow(img)\n",
+    "    ax.axis(\"off\")\n",
+    "plt.tight_layout()\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3387e060",
+   "metadata": {},
+   "source": [
+    "## 9. Next steps\n",
+    "\n",
+    "You have now seen the complete 1.6.0 stack — from a one-liner `model.train()`\n",
+    "through composable PTL components to batch inference.  From here:\n",
+    "\n",
+    "- [PyTorch Lightning training docs](https://rfdetr.roboflow.com/1.6.0/reference/training/) — custom callbacks, multi-GPU, mixed precision\n",
+    "- [Advanced training options](https://rfdetr.roboflow.com/1.6.0/learn/train/advanced/) — augmentations, EMA, learning rate schedules\n",
+    "- [Logger integrations (ClearML, MLflow, W&B)](https://rfdetr.roboflow.com/1.6.0/learn/train/loggers/) — experiment tracking\n",
+    "- [Export your model](https://rfdetr.roboflow.com/1.6.0/learn/export/) — ONNX, TensorRT, CoreML"
+   ]
+  }
+ ],
+ "metadata": {
+  "jupytext": {
+   "cell_metadata_filter": "-all",
+   "formats": "ipynb,py",
+   "main_language": "python"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}