notebooks/roboflow-100-pytorch.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b471e7ce",
   "metadata": {},
   "source": [
    "# RoboFlow 100 - Visualize images with PyTorch\n",
    "\n",
    "Let's see how to parse and visualize images inside RF100 with PyTorch. Let's start by importing a bunch of stuff"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "3ead1e30",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "from ipywidgets import interact\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import torchvision.transforms as T\n",
    "from typing import Callable, Tuple, Optional, List, Dict\n",
    "import torch\n",
    "from PIL import Image\n",
    "from dataclasses import dataclass"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1113bef0",
   "metadata": {},
   "source": [
    "## Downloading RF100\n",
    "\n",
    "We will assume you are `RF100` saved on your disk, we provided a detailed guide on the official [README](https://github.com/roboflow-ai/roboflow-100-benchmark/blob/main/README.md) with instruction for dowloading it.\n",
    "\n",
    "**We assume the dataset was downloaded in yolov5/7 format** using the `-f yolov5` flag"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "03f09607",
   "metadata": {},
   "source": [
    "## Dataset\n",
    "\n",
    "We need to know where `RF100` is stored, feel free to change `ROOT`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "cd85ea27",
   "metadata": {},
   "outputs": [],
   "source": [
    "ROOT = Path('../rf100/') # <- change me :)\n",
    "datasets = sorted(list(ROOT.iterdir()), key=lambda x: x.stem)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2d6ba27",
   "metadata": {},
   "source": [
    "Next, we need the classic [PyTorch `Dataset`](https://pytorch.org/tutorials/beginner/basics/data_tutorial.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "28f824fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ImageDataset(Dataset):\n",
    "    def __init__(self, root: Path, split: str = \"train\", transform: Callable = None):\n",
    "        super().__init__()\n",
    "        # we have three splits, \"train\", \"valid\" and \"test\"\n",
    "        self.src = root / split\n",
    "        # images and labels are linked by the same name\n",
    "        self.names = list(map(lambda x: x.stem, (self.src / \"labels\").glob(\"*.txt\")))\n",
    "        self.transform = transform\n",
    "\n",
    "    def get_image(self, image_path: Path) -> Image.Image:\n",
    "        \"\"\"\n",
    "        This function opens the image and returns it\n",
    "        \"\"\"\n",
    "        image = Image.open(image_path).convert(\"RGB\")\n",
    "        return image\n",
    "\n",
    "    def get_labels(self, labels_path: Path) -> Optional[torch.Tensor]:\n",
    "        \"\"\"\n",
    "        This function reads the label txt file in yolo format, each line of the file looks like\n",
    "        \n",
    "        <label_id> <x_center> <y_center> <width> <height>\n",
    "        \n",
    "        The coordinates are with respect to the image's width and height, so between 0 - 1\n",
    "        \n",
    "        We parse the labels with the following steps:\n",
    "            1) read line by line\n",
    "            2) for each line, get the label id, x_center, y_center, width and height\n",
    "            3) convert to a tensor\n",
    "            4) add to the previous one by stacking them vertically\n",
    "            \n",
    "        The return tensor has shape `batch_size, 5`\n",
    "        \"\"\"\n",
    "        labels = None\n",
    "        with labels_path.open(\"r\") as f:\n",
    "            for line in f.readlines():\n",
    "                parsed_line = [float(e) for e in line.strip().split(\" \")]\n",
    "                if len(parsed_line) != 5:\n",
    "                    continue\n",
    "                c, x, y, w, h = [float(e) for e in line.strip().split(\" \")]\n",
    "                label = torch.as_tensor([[c, x, y, w, h]])\n",
    "                if labels is None:\n",
    "                    labels = label\n",
    "                else:\n",
    "                    labels = torch.vstack([labels, label])\n",
    "        return labels\n",
    "\n",
    "    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:\n",
    "        \"\"\"\n",
    "        Parse image and associated labels are return them as a tensor.\n",
    "        \"\"\"\n",
    "        image_path = self.src / \"images\" / f\"{self.names[idx]}.jpg\"\n",
    "        labels_path = self.src / \"labels\" / f\"{self.names[idx]}.txt\"\n",
    "        image = self.get_image(image_path)\n",
    "        labels = self.get_labels(labels_path)\n",
    "        bboxes = None\n",
    "        if labels is not None:\n",
    "            bboxes = labels[..., 1:]\n",
    "            labels = labels[..., 0].to(torch.int)\n",
    "            image, bboxes = self.transform(image, bboxes)\n",
    "        else:\n",
    "            # if we don't have labels, let's use this values\n",
    "            labels = torch.tensor([-1.0])\n",
    "            bboxes = torch.as_tensor([[0.0, 0.0, 0.0, 0.0]])\n",
    "        image, bboxes = self.transform(image, bboxes)\n",
    "        image = T.functional.to_tensor(image)\n",
    "        return {\"image\": image, \"bboxes\": bboxes, \"labels\": labels}\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.names)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "05f393fd",
   "metadata": {},
   "source": [
    "### Transformations\n",
    "\n",
    "Since we will probably need to resize the image, we can create a custom transformation that works on `image` and `bboxes` as well! [Latest version of `torchvision`](https://pytorch.org/blog/extending-torchvisions-transforms-to-object-detection-segmentation-and-video-tasks/) supports both images and bboxes, feel free to use it instead! "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "a8e7f8cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class Resize:\n",
    "    size: Tuple[int, int]\n",
    "    \n",
    "    def __call__(self, image: Image.Image, bboxes: Optional[torch.Tensor]) -> Tuple[Image.Image, torch.Tensor]:\n",
    "        # in PIL they are inverted LoL\n",
    "        w, h = image.size\n",
    "        if bboxes is not None:\n",
    "            bboxes *= torch.as_tensor([w, h, w, h])\n",
    "        image = T.functional.resize(image, self.size)\n",
    "        new_w, new_h = image.size\n",
    "        if bboxes is not None:\n",
    "            # map to new sizes\n",
    "            bboxes /= torch.as_tensor([w / new_w , h / new_h, w / new_w, h / new_h])\n",
    "            # to 0 - 1\n",
    "            bboxes /= torch.as_tensor([new_w, new_h, new_w, new_h])\n",
    "        return image, bboxes"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8f52bb4",
   "metadata": {},
   "source": [
    "## Visualisation\n",
    "Let's create `AnnotatedImage` to help us draw the bboxes on the image, since `torchvision.utils.draw_bounding_boxes` uses `xyxy` format, we will use `torchvision.ops.boxes.box_convert` to convert yolo style annotation (`x_center`, `y_center`, `width`, `height`) to `xyxy`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "9d1d6e72",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision.utils import draw_bounding_boxes\n",
    "from torchvision.ops.boxes import box_convert\n",
    "\n",
    "class AnnotatedImage:\n",
    "    def __init__(self, image: torch.Tensor, bboxes: torch.Tensor, labels: torch.Tensor):\n",
    "        self.image = image\n",
    "        self.bboxes = bboxes\n",
    "        self.labels = labels\n",
    "        \n",
    "    def draw(self):\n",
    "        c, h, w = self.image.shape\n",
    "        bboxes = self.bboxes * torch.as_tensor([w, h, w, h])\n",
    "        return draw_bounding_boxes(\n",
    "            (self.image * 255).to(torch.uint8), \n",
    "            bboxes, \n",
    "            width=3,\n",
    "            colors=[\"yellow\"] * len(self.labels),\n",
    "            labels=[str(i.item()) for i in self.labels])\n",
    "    \n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d1e1736",
   "metadata": {},
   "source": [
    "Finally, we can use jupyter `ipywidgets` module to visualize the images and the labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "579f0244",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c7e5f633c7d84f8b9b3e3019f3fd1cb6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(Dropdown(description='dataset_path', options=('4-fold-defect', 'abdomen-mri', 'acl-x-ray…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from torchvision.ops.boxes import box_convert\n",
    "\n",
    "from ipywidgets import interactive,Dropdown, IntSlider\n",
    "\n",
    "dataset_path_dropdown = Dropdown(options=map(lambda x: x.stem, datasets))\n",
    "dataset_path_split = Dropdown(options=[\"train\", \"valid\", \"test\"])\n",
    "image_idx_slider = IntSlider(min=0, max=1)\n",
    "\n",
    "@interact(dataset_path=dataset_path_dropdown, split=dataset_path_split, image_idx=image_idx_slider)\n",
    "def visualize(dataset_path: Path, split: str = \"train\", image_idx: int = 0):\n",
    "    ds = ImageDataset(ROOT / dataset_path, split=split, transform=Resize((640, 640)))\n",
    "    image_idx_slider.max = len(ds)\n",
    "    # let's be sure we are within range\n",
    "    image_idx = min(image_idx_slider.value, len(ds) - 1)\n",
    "    data = ds[image_idx]\n",
    "    bboxes = box_convert(data[\"bboxes\"], in_fmt=\"cxcywh\", out_fmt=\"xyxy\")\n",
    "    img = AnnotatedImage(data[\"image\"], bboxes, data[\"labels\"]).draw()\n",
    "    return Image.fromarray(img.permute(1,2,0).numpy())\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "162a9360",
   "metadata": {},
   "source": [
    "Et voilà 🥳"
   ]
  }
 ],
 "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.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}