diff --git a/notebooks/train-rf-detr-with-augmentation.ipynb b/notebooks/train-rf-detr-with-augmentation.ipynb new file mode 100644 index 00000000..dbe97f00 --- /dev/null +++ b/notebooks/train-rf-detr-with-augmentation.ipynb @@ -0,0 +1,404 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "1c44f3cf", + "metadata": {}, + "source": [ + "# RF-DETR 1.5 — `AUG_AERIAL`: Augmentations for Overhead Imagery\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/notebooks/aerial-augmentation-demo.ipynb)\n", + "\n", + "Aerial images break one of detection's core assumptions: objects always\n", + "appear upright. From a UAV or satellite, a car travelling north looks\n", + "nothing like the same car travelling east. Without rotation augmentation,\n", + "a detector trained on mostly north-facing objects will miss the others.\n", + "\n", + "RF-DETR 1.5 ships an **`AUG_AERIAL` preset** built for this problem.\n", + "It applies 90° discrete rotations together with horizontal / vertical\n", + "flips — giving the model all 8 canonical orientations of every training\n", + "image for free.\n", + "\n", + "**What you will learn:**\n", + "- Why rotation augmentation is essential for top-down imagery\n", + "- How to visually inspect augmented batches *before* training\n", + "- How to measure the mAP gain from `AUG_AERIAL` with a side-by-side comparison" + ] + }, + { + "cell_type": "markdown", + "id": "9f252bbd", + "metadata": {}, + "source": [ + "## 1. Install RF-DETR 1.5" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0ede61d3", + "metadata": {}, + "outputs": [], + "source": [ + "!pip install -q rfdetr>=1.5.0 roboflow" + ] + }, + { + "cell_type": "markdown", + "id": "ae1760ef", + "metadata": {}, + "source": [ + "## 2. Check GPU\n", + "\n", + "RF-DETR trains on GPU when available and falls back to CPU.\n", + "The cell below prints VRAM — a handy sanity check before choosing batch size." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "162fab6b", + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "\n", + "import torch\n", + "\n", + "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n", + "print(f\"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", + "num_workers = min(os.cpu_count() or 2, 8)\n", + "print(f\"Data-loader workers: {num_workers}\")\n", + "\n", + "EPOCHS = 50" + ] + }, + { + "cell_type": "markdown", + "id": "d20e825b", + "metadata": {}, + "source": [ + "## 3. Download an aerial dataset\n", + "\n", + "We use [**Aerial Cows**](https://universe.roboflow.com/roboflow-100/aerial-cows)\n", + "from Roboflow 100 — 1,084 drone images of cattle filmed from directly above.\n", + "Cows face every compass direction, so rotation augmentation has a real job to do.\n", + "\n", + "The dataset is CC BY 4.0. You need a **free Roboflow API key**:\n", + "add it as a Colab secret named `ROBOFLOW_API_KEY` (or set the env var locally).\n", + "\n", + "> To swap in a different aerial dataset, change `WORKSPACE`, `PROJECT`, `VERSION`\n", + "> below. Browse https://universe.roboflow.com/browse/aerial for alternatives." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "81145168", + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "\n", + "from rfdetr import RFDETRSmall\n", + "from roboflow import Roboflow\n", + "\n", + "WORKSPACE = \"roboflow-100\"\n", + "PROJECT = \"aerial-cows\"\n", + "VERSION = 2\n", + "\n", + "try:\n", + " from google.colab import userdata # type: ignore[import]\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(WORKSPACE).project(PROJECT).version(VERSION).download(\"coco\")\n", + "DATASET_DIR = dataset.location\n", + "print(f\"Dataset root: {DATASET_DIR}\")" + ] + }, + { + "cell_type": "markdown", + "id": "8feaf494", + "metadata": {}, + "source": [ + "## 4. Explore the `AUG_AERIAL` preset\n", + "\n", + "`AUG_AERIAL` is a plain Python dict — inspect, extend, or override it freely.\n", + "\n", + "| Transform | Why it matters for overhead imagery |\n", + "|---|---|\n", + "| `HorizontalFlip` | Objects can face left or right |\n", + "| `VerticalFlip` | Objects can face toward or away from the sensor |\n", + "| `Rotate(limit=(90, 90))` | Combined with the flips, covers all 8 cardinal orientations |\n", + "| `RandomBrightnessContrast` | Handles lighting changes across time of day, altitude, and cloud cover |\n", + "\n", + "The first three transforms together form the dihedral group D₄ of the\n", + "square — meaning the model sees every 90° rotation and reflection of\n", + "every training image." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5d1ff00d", + "metadata": {}, + "outputs": [], + "source": [ + "from rfdetr.datasets.aug_config import AUG_AERIAL\n", + "\n", + "print(\"AUG_AERIAL preset:\")\n", + "for transform, params in AUG_AERIAL.items():\n", + " print(f\" {transform}: {params}\")" + ] + }, + { + "cell_type": "markdown", + "id": "ca0bac8f", + "metadata": {}, + "source": [ + "## 5. Train: baseline vs `AUG_AERIAL`\n", + "\n", + "We run two training runs on the same dataset and compare validation mAP@50:95.\n", + "The baseline uses no augmentations; the second run uses `AUG_AERIAL`.\n", + "\n", + "> **Tip:** On a T4 GPU with a ~500-image aerial dataset, each 50-epoch run\n", + "> takes roughly 5–10 minutes." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0dedf007", + "metadata": {}, + "outputs": [], + "source": [ + "# ── Baseline: no augmentations ──────────────────────────────────────────────\n", + "OUTPUT_BASE = \"output_baseline\"\n", + "os.makedirs(OUTPUT_BASE, exist_ok=True)\n", + "\n", + "model_base = RFDETRSmall()\n", + "model_base.train(\n", + " dataset_dir=DATASET_DIR,\n", + " epochs=EPOCHS,\n", + " batch_size=8,\n", + " aug_config={}, # no augmentations\n", + " output_dir=OUTPUT_BASE,\n", + " device=device,\n", + " num_workers=num_workers,\n", + " run_test=False,\n", + " progress_bar=True,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "18b4076d", + "metadata": {}, + "outputs": [], + "source": [ + "# ── With AUG_AERIAL ──────────────────────────────────────────────────────────\n", + "OUTPUT_AUG = \"output_aerial\"\n", + "os.makedirs(OUTPUT_AUG, exist_ok=True)\n", + "\n", + "model_aug = RFDETRSmall()\n", + "model_aug.train(\n", + " dataset_dir=DATASET_DIR,\n", + " epochs=EPOCHS,\n", + " batch_size=8,\n", + " aug_config=AUG_AERIAL,\n", + " save_dataset_grids=True, # writes 3×3 grids before training starts\n", + " output_dir=OUTPUT_AUG,\n", + " device=device,\n", + " num_workers=num_workers,\n", + " run_test=False,\n", + " progress_bar=True,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "a22592d4", + "metadata": {}, + "source": [ + "### Augmented training batches\n", + "\n", + "RF-DETR writes the grids before the first weight update — you can see exactly\n", + "what the model will be trained on." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0f9504a8", + "metadata": {}, + "outputs": [], + "source": [ + "from pathlib import Path\n", + "\n", + "import matplotlib.image as mpimg\n", + "import matplotlib.pyplot as plt\n", + "\n", + "grids = sorted(Path(OUTPUT_AUG).glob(\"train_batch*_grid.jpg\"))[:3]\n", + "if grids:\n", + " fig, axes = plt.subplots(1, len(grids), figsize=(18, 6))\n", + " if len(grids) == 1:\n", + " axes = [axes]\n", + " for ax, g in zip(axes, grids):\n", + " ax.imshow(mpimg.imread(g))\n", + " ax.set_title(g.stem)\n", + " ax.axis(\"off\")\n", + " plt.suptitle(\"Augmented training batches — AUG_AERIAL\", fontsize=13)\n", + " plt.tight_layout()\n", + " plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "9bae1c07", + "metadata": {}, + "source": [ + "## 6. Compare results\n", + "\n", + "RF-DETR writes a `log.txt` (one JSON object per epoch) to `output_dir`.\n", + "We parse `test_coco_eval_bbox[0]` — the standard COCO mAP@50:95 — from\n", + "both runs and plot them side by side." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "746731a0", + "metadata": {}, + "outputs": [], + "source": [ + "import json\n", + "\n", + "import matplotlib.pyplot as plt\n", + "\n", + "\n", + "def read_map_history(output_dir: str) -> list[float]:\n", + " \"\"\"Return per-epoch mAP@50:95 from RF-DETR's log.txt.\"\"\"\n", + " maps = []\n", + " log_path = Path(output_dir) / \"log.txt\"\n", + " with open(log_path) as f:\n", + " for line in f:\n", + " entry = json.loads(line.strip())\n", + " bbox = entry.get(\"test_coco_eval_bbox\") or entry.get(\"ema_test_coco_eval_bbox\")\n", + " if bbox:\n", + " maps.append(float(bbox[0]))\n", + " return maps\n", + "\n", + "\n", + "base_maps = read_map_history(OUTPUT_BASE)\n", + "aug_maps = read_map_history(OUTPUT_AUG)\n", + "\n", + "best_base = max(base_maps) if base_maps else float(\"nan\")\n", + "best_aug = max(aug_maps) if aug_maps else float(\"nan\")\n", + "delta = best_aug - best_base\n", + "\n", + "# ── Side-by-side: training curve | best-mAP bar chart ──────────────────────\n", + "fig, (ax_curve, ax_bar) = plt.subplots(1, 2, figsize=(14, 4))\n", + "\n", + "# Left: mAP over epochs\n", + "ax_curve.plot(base_maps, linewidth=2, label=\"No augmentation\")\n", + "ax_curve.plot(aug_maps, linewidth=2, label=\"AUG_AERIAL\")\n", + "ax_curve.set_xlabel(\"Epoch\")\n", + "ax_curve.set_ylabel(\"mAP@50:95 (validation)\")\n", + "ax_curve.set_title(\"mAP over training\")\n", + "ax_curve.legend()\n", + "ax_curve.grid(alpha=0.3)\n", + "\n", + "# Right: best mAP bar chart\n", + "labels = [\"No augmentation\", \"AUG_AERIAL\"]\n", + "values = [best_base, best_aug]\n", + "bars = ax_bar.bar(labels, values, width=0.4)\n", + "bars[1].set_color(\"C1\")\n", + "ax_bar.set_ylabel(\"Best mAP@50:95\")\n", + "ax_bar.set_title(f\"Best result ({delta:+.4f})\")\n", + "ax_bar.set_ylim(0, max(values) * 1.2)\n", + "for bar, val in zip(bars, values):\n", + " ax_bar.text(bar.get_x() + bar.get_width() / 2, val + 0.005, f\"{val:.4f}\", ha=\"center\")\n", + "\n", + "plt.tight_layout()\n", + "plt.show()\n", + "\n", + "print(f\"Best mAP@50:95 — no augmentation : {best_base:.4f}\")\n", + "print(f\"Best mAP@50:95 — AUG_AERIAL : {best_aug:.4f} ({delta:+.4f})\")" + ] + }, + { + "cell_type": "markdown", + "id": "0372c2bd", + "metadata": {}, + "source": [ + "## 7. Run inference on a validation image" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fa9aa2b7", + "metadata": {}, + "outputs": [], + "source": [ + "import json\n", + "from pathlib import Path\n", + "\n", + "import supervision as sv\n", + "from PIL import Image\n", + "\n", + "# Load the first image from the validation split\n", + "val_ann_path = Path(DATASET_DIR) / \"valid\" / \"_annotations.coco.json\"\n", + "with open(val_ann_path) as f:\n", + " ann_data = json.load(f)\n", + "\n", + "first_filename = ann_data[\"images\"][0][\"file_name\"]\n", + "image = Image.open(Path(DATASET_DIR) / \"valid\" / first_filename)\n", + "\n", + "detections = model_aug.predict(image, threshold=0.3)\n", + "annotated = sv.BoxAnnotator().annotate(image.copy(), detections)\n", + "sv.plot_image(annotated)\n", + "print(f\"Detections: {len(detections)}\")" + ] + }, + { + "cell_type": "markdown", + "id": "07089325", + "metadata": {}, + "source": [ + "## Next steps\n", + "\n", + "**Extend the preset** — add transforms that match your specific conditions:\n", + "```python\n", + "custom_aerial = {\n", + " **AUG_AERIAL,\n", + " \"GaussianBlur\": {\"blur_limit\": 3, \"p\": 0.2}, # UAV motion blur\n", + " \"RandomScale\": {\"scale_limit\": 0.3, \"p\": 0.4}, # altitude variation\n", + "}\n", + "```\n", + "\n", + "- [Augmentation docs](https://rfdetr.roboflow.com/develop/learn/train/augmentations/)\n", + "- [Advanced training options](https://rfdetr.roboflow.com/develop/learn/train/advanced/)\n", + "- [Logger integrations (ClearML, MLflow, W&B)](https://rfdetr.roboflow.com/develop/learn/train/loggers/)\n", + "- [Export your model](https://rfdetr.roboflow.com/develop/learn/export/)" + ] + } + ], + "metadata": { + "jupytext": { + "cell_metadata_filter": "-all", + "main_language": "python", + "notebook_metadata_filter": "-all" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +}