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