{
 "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
}