workshop-11-03-2025-detect-huggingface-training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "941ee00e",
   "metadata": {},
   "source": [
    "# Workshop: Object Detection with Hugging Face, Ultralytics YOLOv8, and Lightning\n",
    "\n",
    "In this workshop we will explore three approaches for object detection using PyTorch:\n",
    "\n",
    "1. **Inference with Hugging Face DETR:** Using a pre-trained DETR model via a pipeline.\n",
    "2. **Inference & Training with Ultralytics YOLOv8:** Running inference and training a YOLO model on your own dataset (requires a YOLO-formatted dataset and a `data.yaml` configuration file).\n",
    "3. **Training Faster‑R‑CNN with PyTorch Lightning:** Wrapping a TorchVision Faster‑R‑CNN model in a LightningModule and training it on the PennFudanPed dataset, with data augmentation via TorchVision Transforms v2.\n",
    "\n",
    "Follow along for hands-on experience!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab08c400",
   "metadata": {},
   "source": [
    "## Installation\n",
    "\n",
    "Install the required packages by running the cell below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe35f62d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install -U transformers pillow matplotlib ultralytics timm lightning gdown"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d344a65d-01a4-482b-ac2e-51b2f08c31bf",
   "metadata": {},
   "source": [
    "## Organize Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "335f49d2-035a-41db-ba66-67a671b06622",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import gc\n",
    "import sys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40230fd5-4ac4-46ba-b5f5-73790f2a5e74",
   "metadata": {},
   "outputs": [],
   "source": [
    "import requests, zipfile, io\n",
    "from pathlib import Path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59ab2b49-c361-49e6-a36a-b6c53f295a2f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4e036a9-149c-4d49-868a-497877b6a020",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c42f45d-bce3-43c4-935e-824fca5032dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "import requests\n",
    "import glob\n",
    "from PIL import Image\n",
    "from IPython.display import display"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ec75c8b-ab7b-47f6-be38-37dbef200ea7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.utils.data import DataLoader, random_split\n",
    "import torch.optim as optim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b0efd8c-feb6-460e-a936-8fee8692a542",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torchvision\n",
    "# Use TorchVision Transforms v2 for data augmentation\n",
    "from torchvision.transforms import v2 as T2\n",
    "from torchvision.transforms import v2 as transforms\n",
    "from torchvision.utils import draw_bounding_boxes\n",
    "import torchvision.transforms.functional as F\n",
    "from torchvision import models, datasets, ops, utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df40e828-8536-486b-8dcb-b03a1d179eb6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd195d7c",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import lightning as pl\n",
    "import lightning as L"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c00093a-db41-4895-8fe6-5c5310288df6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ultralytics import YOLO"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b1bba29-0bd0-4324-b38e-58f4fc20b4b7",
   "metadata": {},
   "source": [
    "## Initialize Device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "efb8c247-f47a-437a-bd99-11e45b917b3a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def init_device():\n",
    "    # For the most part I'll try to import functions and classes near\n",
    "    # where they are used\n",
    "    # to make it clear where they come from.\n",
    "    if torch.backends.mps.is_available():\n",
    "        device = 'mps'\n",
    "    else:\n",
    "        device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "\n",
    "    print(f'Device: {device}')\n",
    "\n",
    "    return device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8cb4e29-6889-42b4-b626-4a5f59f79d1d",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = init_device()\n",
    "device"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f12ef5b-1126-4d7c-89d0-b2ba8d188a6e",
   "metadata": {},
   "source": [
    "## Initialize Paths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60002f74-84ef-4362-9929-9b5f1affe9cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "PATH = Path('data')\n",
    "DATA = PATH\n",
    "africanw = PATH / 'africanw' / 'african-wildlife.yaml'\n",
    "pennfped = PATH / 'PennFudanPed' / 'PennFudanPed'\n",
    "models_path = PATH / 'models'\n",
    "models_path.mkdir(exist_ok=True, parents=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8f854bc-67d4-4264-a5b9-3a6f6b7e4da7",
   "metadata": {},
   "outputs": [],
   "source": [
    "pennfped.parent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21b5d98e-832b-4dd3-b204-2836a3b7395a",
   "metadata": {},
   "outputs": [],
   "source": [
    "! ls {pennfped}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1eef5a2e-e3a7-41f3-98f4-b543fdf2a3b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "africanw"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "640d9d2c-f499-4b0a-af8b-2d7b1090a279",
   "metadata": {},
   "source": [
    "## Download Datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "559fe920-1b73-43d0-a370-c92db5bca01f",
   "metadata": {},
   "outputs": [],
   "source": [
    "africanw.parent.mkdir(exist_ok=True, parents=True)\n",
    "\n",
    "url = \"https://raw.githubusercontent.com/ultralytics/ultralytics/refs/heads/main/ultralytics/cfg/datasets/african-wildlife.yaml\"\n",
    "\n",
    "if africanw.exists():\n",
    "    print(f'File {africanw} exists')\n",
    "else:\n",
    "    response = requests.get(url)\n",
    "    \n",
    "    if response.status_code == 200:\n",
    "        with africanw.open(mode=\"w\") as f:\n",
    "            f.write(response.text)\n",
    "        print(\"File downloaded successfully as 'african-wildlife.yaml'!\")\n",
    "    else:\n",
    "        print(\"Failed to download file. Status code:\", response.status_code)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7de72ead-70c4-434e-bfb7-34a55183b8bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# URL for the PennFudanPed dataset zip file\n",
    "url = \"https://www.cis.upenn.edu/~jshi/ped_html/PennFudanPed.zip\"\n",
    "\n",
    "if pennfped.exists():\n",
    "    print('Data folder exists')\n",
    "else:\n",
    "    print(\"Downloading PennFudanPed dataset...\")\n",
    "    r = requests.get(url)\n",
    "    if r.status_code == 200:\n",
    "        z = zipfile.ZipFile(io.BytesIO(r.content))\n",
    "        z.extractall(pennfped.parent)\n",
    "        print(\"Downloaded and extracted PennFudanPed dataset to './PennFudanPed'\")\n",
    "    else:\n",
    "        print(\"Download failed with status code:\", r.status_code)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33552c78",
   "metadata": {},
   "source": [
    "## Part 1: Inference with Hugging Face DETR\n",
    "\n",
    "In this section we load a pre-trained DETR model via Hugging Face’s pipeline and run inference on a sample image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24761004",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download a sample image\n",
    "url = \"https://ultralytics.com/images/bus.jpg\"\n",
    "image = Image.open(requests.get(url, stream=True).raw)\n",
    "\n",
    "# Initialize the Hugging Face object detection pipeline (using DETR)\n",
    "detr_detector = pipeline(\"object-detection\", model=\"facebook/detr-resnet-50\")\n",
    "\n",
    "# Run inference\n",
    "results = detr_detector(image)\n",
    "\n",
    "print(\"DETR Inference Results:\")\n",
    "for r in results:\n",
    "    print(r)\n",
    "\n",
    "display(image)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66525c2f",
   "metadata": {},
   "source": [
    "## Part 2: Inference with Ultralytics YOLOv8\n",
    "\n",
    "Next, we use Ultralytics YOLOv8 to run inference on the same sample image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c62bfaed",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the pre-trained YOLOv8 nano model\n",
    "yolo_model = YOLO(\"yolov8n.pt\")\n",
    "\n",
    "# Run inference on the sample image\n",
    "results_yolo = yolo_model(\"https://ultralytics.com/images/bus.jpg\")\n",
    "\n",
    "# Print YOLOv8 results\n",
    "# print(results_yolo)\n",
    "\n",
    "# Plot the image with predictions\n",
    "plt.figure(figsize=(10, 10))\n",
    "plt.imshow(results_yolo[0].plot())\n",
    "plt.axis('off')\n",
    "plt.title('Ultralytics YOLOv8 Inference')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c43a970",
   "metadata": {},
   "source": [
    "## Part 3: Training with Ultralytics YOLOv8\n",
    "\n",
    "To train a YOLO model using Ultralytics, you need a dataset in YOLO format along with a YAML configuration file (e.g., `data/my_dataset/data.yaml`).\n",
    "\n",
    "For example, your `data.yaml` might look like:\n",
    "\n",
    "```yaml\n",
    "train: data/my_dataset/images/train\n",
    "val: data/my_dataset/images/val\n",
    "nc: 2\n",
    "names: ['class1', 'class2']\n",
    "```\n",
    "\n",
    "Make sure that the file exists at the specified path. Then run the cell below to start training for 5 epochs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "228b9176",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize the YOLOv8 nano model with pre-trained weights\n",
    "yolov8n = YOLO('yolov8n.pt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bda8e035-9d1c-411e-abfa-6d8c50029e8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "yolov8n = yolov8n.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96ac86e7-8d3d-4f6d-a5a5-d815dff3190c",
   "metadata": {},
   "outputs": [],
   "source": [
    "yolov8n.train(\n",
    "    data=africanw, \n",
    "    epochs=5, \n",
    "    imgsz=640,\n",
    "    device=device,\n",
    "    workers=8\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4594d023-de1b-4378-af05-b1a6fc748a6d",
   "metadata": {},
   "outputs": [],
   "source": [
    "yolov8n.save(models_path / 'yolov8n_afrwld.pt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "563d12f1-287a-4d9c-bf69-046dcff128a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "yolov8n.export(format='onnx')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ddf7139",
   "metadata": {},
   "source": [
    "## Part 4: Training Faster‑R‑CNN with PyTorch Lightning and TorchVision Transforms v2\n",
    "\n",
    "In this section we train a Faster‑R‑CNN model on the PennFudanPed dataset using PyTorch Lightning. We use a new data augmentation pipeline built with TorchVision Transforms v2. Make sure the PennFudanPed dataset is downloaded and extracted into a folder named `PennFudanPed`.\n",
    "\n",
    "The data augmentation pipeline includes random resized cropping, horizontal flipping, and color jitter. These augmentations help improve model robustness."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c226921-1c14-41f9-b571-bea6a48da3c3",
   "metadata": {},
   "source": [
    "#### Train Face DSetectors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ebc823ff-31ff-425d-8449-efb05c8c460d",
   "metadata": {},
   "outputs": [],
   "source": [
    "gc.collect()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b133155-6d2d-4522-b6b0-f90f8a46c785",
   "metadata": {},
   "source": [
    "#### Initialize Face Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0568aa32-787f-426e-87df-4a5d05dd35ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "class FacesData(L.LightningDataModule):\n",
    "    transform = transforms.Compose([\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n",
    "        transforms.Resize(size=(800,), max_size=1333),\n",
    "    ])\n",
    "\n",
    "    @staticmethod\n",
    "    def convert_inputs(imgs, annot, device, small_thr=0.001):\n",
    "        \"\"\"Conver dataset item to accepted target struture.\"\"\"\n",
    "        images, targets = [], []\n",
    "        for img, annot in zip(imgs, annot):\n",
    "            bbox = annot['bbox']\n",
    "            small = (bbox[:, 2] * bbox[:, 3]) <= (img.size[1] * img.size[0] * small_thr)\n",
    "            boxes = ops.box_convert(bbox[~small], in_fmt='xywh', out_fmt='xyxy')\n",
    "            output_dict = FacesData.transform({\"image\": img, \"boxes\": boxes})\n",
    "            images.append(output_dict['image'].to(device))\n",
    "            targets.append({\n",
    "                'boxes': output_dict['boxes'].to(device),\n",
    "                'labels': torch.ones(len(boxes), dtype=int, device=device)\n",
    "            })\n",
    "        return images, targets\n",
    "    \n",
    "    @staticmethod\n",
    "    def _collate_fn(batch):\n",
    "        \"\"\"Define a collate function to handle batches.\"\"\"\n",
    "        return tuple(zip(*batch))\n",
    "\n",
    "    def train_dataloader(self):# Step 4: Load the WIDERFace dataset using torchvision.datasets\n",
    "        train_dataset = datasets.WIDERFace(root=DATA, split='train', download=True)\n",
    "\n",
    "        # Step 5: Set up the DataLoader and train the model\n",
    "        return DataLoader(\n",
    "            train_dataset, batch_size=8, shuffle=True, num_workers=4, collate_fn=self._collate_fn\n",
    "        )\n",
    "\n",
    "    def val_dataloader(self):# Step 4: Load the WIDERFace dataset using torchvision.datasets\n",
    "        val_dataset = datasets.WIDERFace(root=DATA, split='val', download=True)\n",
    "\n",
    "        # Step 5: Set up the DataLoader and train the model\n",
    "        return DataLoader(\n",
    "            val_dataset, batch_size=8, shuffle=True, num_workers=4, collate_fn=self._collate_fn\n",
    "        )\n",
    "\n",
    "    def test_dataloader(self):# Step 4: Load the WIDERFace dataset using torchvision.datasets\n",
    "        test_dataset = datasets.WIDERFace(root=DATA, split='val', download=True)\n",
    "\n",
    "        # Step 5: Set up the DataLoader and train the model\n",
    "        return DataLoader(\n",
    "            test_dataset, batch_size=8, shuffle=True, num_workers=4, collate_fn=self._collate_fn\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c655d8fa-cddb-4987-8b29-778c67ba0663",
   "metadata": {},
   "source": [
    "#### Initialize the Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98dded5c-e710-4c3b-8e37-142ae3c5fb48",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use a pretrained Faster R-CNN model from torchvision and modify it\n",
    "class FaceDetectionModel(L.LightningModule):\n",
    "    def __init__(self):\n",
    "        super(FaceDetectionModel, self).__init__()\n",
    "        self.model = models.detection.fasterrcnn_mobilenet_v3_large_fpn(weights=\"DEFAULT\")\n",
    "\n",
    "    def forward(self, images, targets=None):\n",
    "        if targets is None:\n",
    "            return self.model(images)\n",
    "        return self.model(images, targets)\n",
    "\n",
    "    def training_step(self, batch, batch_idx):\n",
    "        imgs, annot = batch\n",
    "        images, targets = FacesData.convert_inputs(imgs, annot, device=self.device)\n",
    "        loss_dict = self.model(images, targets)\n",
    "        losses = sum(loss for loss in loss_dict.values())\n",
    "        self.log('train_loss', losses)\n",
    "        return losses\n",
    "\n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        imgs, annot = batch\n",
    "        images, targets = FacesData.convert_inputs(imgs, annot, device=self.device)\n",
    "        loss_dict = self.model(images, targets)\n",
    "        losses = sum(loss for loss in loss_dict.values())\n",
    "        self.log('val_loss', losses)\n",
    "        return losses\n",
    "\n",
    "    def test_step(self, batch, batch_idx):\n",
    "        imgs, annot = batch\n",
    "        images, targets = FacesData.convert_inputs(imgs, annot, device=self.device)\n",
    "        loss_dict = self.model(images, targets)\n",
    "        losses = sum(loss for loss in loss_dict.values())\n",
    "        self.log('test_loss', losses)\n",
    "        return losses\n",
    "\n",
    "    def configure_optimizers(self):\n",
    "        return optim.SGD(self.parameters(), lr=0.001, momentum=0.9, weight_decay=0.0005)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d30eede-474d-42db-98f5-234e6fe3f648",
   "metadata": {},
   "source": [
    "#### Run Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b94bc43-f9f9-44e7-9367-45ebcb1184bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "gc.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01e37d93-1f5d-44d4-ae20-44ea9e37395d",
   "metadata": {},
   "outputs": [],
   "source": [
    "data = FacesData()\n",
    "model = FaceDetectionModel()\n",
    "trainer = L.Trainer(\n",
    "    max_epochs=5, \n",
    "    precision='16-mixed', \n",
    "    log_every_n_steps=10\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c256e1c-c243-493f-bd3b-f34d7878e7f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.fit(model, data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5780bf7-e34b-4403-bd6f-b5fd75253174",
   "metadata": {},
   "outputs": [],
   "source": [
    "gc.collect()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86c81f4c-56aa-4d56-add0-ca45e93ad84e",
   "metadata": {},
   "source": [
    "## Test the Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53f075e2-2d5a-4ed5-acfe-85de1b443f09",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.rcParams[\"savefig.bbox\"] = \"tight\"\n",
    "sample_idx = 0\n",
    "print(f\"selected image sample: {sample_idx}\")\n",
    "\n",
    "def show(imgs):\n",
    "    if not isinstance(imgs, list):\n",
    "        imgs = [imgs]\n",
    "    fig, axs = plt.subplots(ncols=len(imgs), squeeze=False, figsize=(7 * len(imgs), 8))\n",
    "    for i, img in enumerate(imgs):\n",
    "        img = img.detach()\n",
    "        img = F.to_pil_image(img)\n",
    "        axs[0, i].imshow(np.asarray(img))\n",
    "        axs[0, i].set(xticklabels=[], yticklabels=[], xticks=[], yticks=[])\n",
    "    return fig\n",
    "\n",
    "# Step 1: Define the transform\n",
    "transform = transforms.Compose([transforms.ToTensor()])\n",
    "# define the transform\n",
    "normalize = transforms.Compose([\n",
    "    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n",
    "    transforms.Resize(size=(800,), max_size=1333),\n",
    "])\n",
    "\n",
    "# Step 2: Load the WIDERFace dataset using torchvision.datasets\n",
    "test_dataset = datasets.WIDERFace(root=DATA, split='test', download=True, transform=transform)\n",
    "img, target = train_dataset[sample_idx]\n",
    "img = F.convert_image_dtype(img, dtype=torch.uint8)\n",
    "boxes = ops.box_convert(target['bbox'], in_fmt='xywh', out_fmt='xyxy')\n",
    "\n",
    "# visualize the annotation\n",
    "annot = utils.draw_bounding_boxes(img, boxes, colors=\"red\", width=5)\n",
    "\n",
    "# Replace with path to your trained checkpoint 'lightning_logs/version_x/checkpoints/...'\n",
    "checkpoint_path = glob.glob(\"lightning_logs/version_6/checkpoints/*.ckpt\")[0]\n",
    "print(f\"loading model from checkpoint '{checkpoint_path}'\")\n",
    "\n",
    "# Load the model\n",
    "model = FaceDetectionModel.load_from_checkpoint(checkpoint_path).cpu()\n",
    "model.eval()\n",
    "\n",
    "# Get the model prediction\n",
    "img2, _ = train_dataset[sample_idx]\n",
    "with torch.no_grad():\n",
    "    output = model.model([normalize(img2)])\n",
    "print(f\"predistions: {output}\")\n",
    "boxes = output[0]['boxes'][output[0]['scores'] >= 0.15]\n",
    "# visualize the predictions\n",
    "preds = utils.draw_bounding_boxes(img, boxes, colors=\"pink\", width=5)\n",
    "\n",
    "# export figure\n",
    "fig = show([annot, preds])\n",
    "fig.savefig('figure.png')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4af298a-3163-4292-b347-6f8a934a61b0",
   "metadata": {},
   "source": [
    "#### Initialize Transforms for Data Augmentations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2dbac0a4-10ac-41a3-9010-b1736923a35d",
   "metadata": {},
   "outputs": [],
   "source": [
    "gc.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ca6a4e9-d34b-4f19-8ea3-70a32763a79f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_transform(train: bool):\n",
    "    if train:\n",
    "        return T2.Compose([\n",
    "            T2.RandomResizedCrop(size=(300, 300)),\n",
    "            T2.RandomHorizontalFlip(p=0.5),\n",
    "            T2.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.3, hue=0.1),\n",
    "            T2.ToTensor(),\n",
    "            T2.ConvertImageDtype(torch.float),\n",
    "            T2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
    "        ])\n",
    "    else:\n",
    "        return T2.Compose([\n",
    "            T2.Resize((300, 300)),\n",
    "            T2.ToTensor(),\n",
    "            T2.ConvertImageDtype(torch.float),\n",
    "            T2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
    "        ])\n",
    "\n",
    "# Minimal collate function for variable number of targets per image\n",
    "def collate_fn(batch):\n",
    "    return tuple(zip(*batch))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "942e1f48-ab47-49bb-8ed8-4a2fff59f1b7",
   "metadata": {},
   "source": [
    "## Train Model on PennFudanPed Dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07e33a1c-85dd-4da4-820c-2c9140aa1d80",
   "metadata": {},
   "source": [
    "#### Initialize Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6fb0e0e5-1a50-42a5-bb44-dac4b3a1b2ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the PennFudanPed Dataset (adapted from TorchVision tutorials)\n",
    "class PennFudanDataset(torch.utils.data.Dataset):\n",
    "    def __init__(self, root, transforms):\n",
    "        self.root = root\n",
    "        self.transforms = transforms\n",
    "        self.imgs = list(sorted(os.listdir(os.path.join(root, \"PNGImages\"))))\n",
    "        self.masks = list(sorted(os.listdir(os.path.join(root, \"PedMasks\"))))\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        img_path = os.path.join(self.root, \"PNGImages\", self.imgs[idx])\n",
    "        mask_path = os.path.join(self.root, \"PedMasks\", self.masks[idx])\n",
    "        img = Image.open(img_path).convert(\"RGB\")\n",
    "        mask = Image.open(mask_path)\n",
    "        mask = np.array(mask)\n",
    "\n",
    "        # Instances are encoded as different colors\n",
    "        obj_ids = np.unique(mask)[1:]\n",
    "        masks = mask == obj_ids[:, None, None]\n",
    "\n",
    "        boxes = []\n",
    "        for i in range(len(obj_ids)):\n",
    "            pos = np.where(masks[i])\n",
    "            xmin = np.min(pos[1])\n",
    "            xmax = np.max(pos[1])\n",
    "            ymin = np.min(pos[0])\n",
    "            ymax = np.max(pos[0])\n",
    "            boxes.append([xmin, ymin, xmax, ymax])\n",
    "        boxes = torch.as_tensor(boxes, dtype=torch.float32)\n",
    "        num_objs = len(obj_ids)\n",
    "        labels = torch.ones((num_objs,), dtype=torch.int64)  # one class: person\n",
    "\n",
    "        target = {}\n",
    "        target[\"boxes\"] = boxes\n",
    "        target[\"labels\"] = labels\n",
    "        target[\"image_id\"] = torch.tensor([idx])\n",
    "\n",
    "        if self.transforms is not None:\n",
    "            img = self.transforms(img)\n",
    "\n",
    "        return img, target\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.imgs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b28a3a9-5437-4727-b1a8-0f7485abdc98",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Prepare the dataset using the new transforms v2\n",
    "dataset_full = PennFudanDataset(pennfped.absolute(), get_transform(train=True))\n",
    "n = len(dataset_full)\n",
    "n_train = int(0.8 * n)\n",
    "n_val = n - n_train\n",
    "dataset_train, dataset_val = random_split(dataset_full, [n_train, n_val])\n",
    "\n",
    "workers = 0\n",
    "\n",
    "train_loader = DataLoader(\n",
    "    dataset_train, \n",
    "    batch_size=2, \n",
    "    shuffle=True, \n",
    "    num_workers=workers,\n",
    "    # persistent_workers=True,\n",
    "    collate_fn=collate_fn\n",
    ")\n",
    "val_loader = DataLoader(\n",
    "    dataset_val, \n",
    "    batch_size=4, \n",
    "    shuffle=False, \n",
    "    num_workers=workers,\n",
    "    # persistent_workers=True,\n",
    "    collate_fn=collate_fn\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a04a5e6-14c5-4c7a-8454-fe02df6c7f16",
   "metadata": {},
   "source": [
    "#### Initialize the Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5479fe8e-2dd5-42e1-90ce-a45f38742a1e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a PyTorch Lightning Module for Faster-RCNN\n",
    "class FasterRCNNLightning(pl.LightningModule):\n",
    "    def __init__(self, num_classes=2, lr=0.005):\n",
    "        super().__init__()\n",
    "        # Load pre-trained Faster-RCNN model\n",
    "        self.model = torchvision.models.detection.fasterrcnn_resnet50_fpn(\n",
    "            weights=torchvision.models.detection.faster_rcnn.FasterRCNN_ResNet50_FPN_Weights.DEFAULT\n",
    "        )\n",
    "        in_features = self.model.roi_heads.box_predictor.cls_score.in_features\n",
    "        self.model.roi_heads.box_predictor = torchvision.models.detection.faster_rcnn.FastRCNNPredictor(in_features, num_classes)\n",
    "        self.lr = lr\n",
    "\n",
    "    def forward(self, images, targets=None):\n",
    "        return self.model(images, targets)\n",
    "\n",
    "    def training_step(self, batch, batch_idx):\n",
    "        images, targets = batch\n",
    "        images = [img.to(self.device) for img in images]\n",
    "        targets = [{k: v.to(self.device) for k, v in t.items()} for t in targets]\n",
    "        loss_dict = self.model(images, targets)\n",
    "        loss = sum(loss for loss in loss_dict.values())\n",
    "        self.log(\"train_loss\", loss, prog_bar=True)\n",
    "        return loss\n",
    "\n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        images, targets = batch\n",
    "        images = [img.to(self.device) for img in images]\n",
    "        targets = [{k: v.to(self.device) for k, v in t.items()} for t in targets]\n",
    "        loss_dict = self.model(images, targets)\n",
    "        loss = sum(loss for loss in loss_dict.values())\n",
    "        self.log(\"val_loss\", loss, prog_bar=True)\n",
    "        return loss\n",
    "\n",
    "    def configure_optimizers(self):\n",
    "        optimizer = torch.optim.SGD(self.model.parameters(), lr=self.lr, momentum=0.9, weight_decay=0.0005)\n",
    "        scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)\n",
    "        return [optimizer], [scheduler]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "709752c0-7a04-43bb-8b05-598120103ce7",
   "metadata": {},
   "source": [
    "#### Train the Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "732a9b02-0a72-4fb9-a4ad-a41bfe0d6917",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Instantiate the Lightning module\n",
    "model_lightning = FasterRCNNLightning(num_classes=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "292a9515-5705-43b3-bdcd-6c39f898f57b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize a PyTorch Lightning Trainer\n",
    "trainer = pl.Trainer(\n",
    "    max_epochs=2, \n",
    "    accelerator='auto', \n",
    "    devices=1\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce125c34-2f2d-4012-87f0-51ca5321ef00",
   "metadata": {},
   "outputs": [],
   "source": [
    "gc.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3fcdac54-8373-4c6f-9108-dc3f34aee763",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train the Faster-RCNN model\n",
    "trainer.fit(\n",
    "    model_lightning, \n",
    "    train_loader, \n",
    "    val_loader\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25bbdfec-e94b-42a9-aa3d-21c70b3f5599",
   "metadata": {},
   "source": [
    "## Visualize the Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9b864b5-5d0b-4cf1-bd18-8af671673431",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch = next(iter(val_loader))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ada6ad4b-c72b-4007-8114-29d0f26a3a53",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Ensure the underlying TorchVision model is in eval mode\n",
    "model_lightning.model.eval()\n",
    "\n",
    "# Define mean and std used during training (for un-normalization)\n",
    "mean = torch.tensor([0.485, 0.456, 0.406]).view(3,1,1)\n",
    "std = torch.tensor([0.229, 0.224, 0.225]).view(3,1,1)\n",
    "\n",
    "def unnormalize(img):\n",
    "    \"\"\"Reverse the normalization on an image tensor.\"\"\"\n",
    "    return img * std + mean\n",
    "\n",
    "# Get a batch from the validation DataLoader (val_loader from training section)\n",
    "batch = next(iter(val_loader))\n",
    "images, targets = batch\n",
    "# Move images to device (assumed same device as model)\n",
    "images = [img.to(model_lightning.device) for img in images]\n",
    "\n",
    "# Run inference (without gradients)\n",
    "with torch.no_grad():\n",
    "    outputs = model_lightning.model(images)\n",
    "\n",
    "# Loop over each image in the batch and plot predictions\n",
    "for i, img in enumerate(images):\n",
    "    # Unnormalize the image for visualization\n",
    "    img_unnorm = unnormalize(img).clamp(0, 1)\n",
    "    # Convert tensor to uint8 for drawing\n",
    "    img_uint8 = (img_unnorm * 255).type(torch.uint8)\n",
    "    \n",
    "    # Get predictions for the image and filter by confidence threshold (e.g., 0.5)\n",
    "    boxes = outputs[i][\"boxes\"].detach().cpu()\n",
    "    scores = outputs[i][\"scores\"].detach().cpu()\n",
    "    keep = scores >= 0.5\n",
    "    boxes = boxes[keep]\n",
    "    \n",
    "    # Draw boxes on the image\n",
    "    drawn_img = draw_bounding_boxes(img_uint8, boxes, colors=\"red\", width=2)\n",
    "    \n",
    "    # Convert to PIL image and display\n",
    "    plt.figure(figsize=(8, 8))\n",
    "    plt.imshow(F.to_pil_image(drawn_img))\n",
    "    plt.title(f\"Validation Image {i} Predictions\")\n",
    "    plt.axis(\"off\")\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aeee1d94",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "\n",
    "In this notebook we demonstrated:\n",
    "\n",
    "- **Inference with Hugging Face DETR:** Running inference on a sample image using a DETR model via Transformers.\n",
    "- **Inference & Training with Ultralytics YOLOv8:** Running inference on a sample image and training a YOLO model using Ultralytics (ensure your dataset YAML file exists at the specified path).\n",
    "- **Training Faster‑R‑CNN with PyTorch Lightning:** Wrapping TorchVision’s Faster‑R‑CNN in a LightningModule, using TorchVision Transforms v2 for data augmentation on the PennFudanPed dataset, and training the model.\n",
    "\n",
    "Feel free to experiment further with hyperparameters, dataset splits, and alternative models. Happy detecting and training!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27f03374-aded-47ae-94d3-cef4333c313b",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Ensure the underlying TorchVision model is in eval mode\n",
    "model_lightning.model.eval()\n",
    "\n",
    "# Define mean and std used during training (for un-normalization)\n",
    "mean = torch.tensor([0.485, 0.456, 0.406]).view(3,1,1)\n",
    "std = torch.tensor([0.229, 0.224, 0.225]).view(3,1,1)\n",
    "\n",
    "def unnormalize(img):\n",
    "    \"\"\"Reverse the normalization on an image tensor.\"\"\"\n",
    "    return img * std + mean\n",
    "\n",
    "# Get a batch from the validation DataLoader (val_loader from training section)\n",
    "batch = next(iter(val_loader))\n",
    "images, targets = batch\n",
    "# Move images to device (assumed same device as model)\n",
    "images = [img.to(model_lightning.device) for img in images]\n",
    "\n",
    "# Run inference (without gradients)\n",
    "with torch.no_grad():\n",
    "    outputs = model_lightning.model(images)\n",
    "\n",
    "# Loop over each image in the batch and plot predictions\n",
    "for i, img in enumerate(images):\n",
    "    # Unnormalize the image for visualization\n",
    "    img_unnorm = unnormalize(img).clamp(0, 1)\n",
    "    # Convert tensor to uint8 for drawing\n",
    "    img_uint8 = (img_unnorm * 255).type(torch.uint8)\n",
    "    \n",
    "    # Get predictions for the image and filter by confidence threshold (e.g., 0.5)\n",
    "    boxes = outputs[i][\"boxes\"].detach().cpu()\n",
    "    scores = outputs[i][\"scores\"].detach().cpu()\n",
    "    keep = scores >= 0.5\n",
    "    boxes = boxes[keep]\n",
    "    \n",
    "    # Draw boxes on the image\n",
    "    drawn_img = draw_bounding_boxes(img_uint8, boxes, colors=\"red\", width=2)\n",
    "    \n",
    "    # Convert to PIL image and display\n",
    "    plt.figure(figsize=(8, 8))\n",
    "    plt.imshow(F.to_pil_image(drawn_img))\n",
    "    plt.title(f\"Validation Image {i} Predictions\")\n",
    "    plt.axis(\"off\")\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60e046cb-0cef-4c83-b3c9-87847c13defd",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}