example.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b7551eb7",
   "metadata": {},
   "source": [
    "# objdetect"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "207a03c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip3 install git+https://github.com/rpmcruz/objdetect.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4cddb90",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torchvision\n",
    "import torch\n",
    "import numpy as np\n",
    "import objdetect as od\n",
    "import albumentations as A\n",
    "from albumentations.pytorch import ToTensorV2\n",
    "from time import time\n",
    "import matplotlib.pyplot as plt\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "download = False\n",
    "data_path = '/data'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58716047",
   "metadata": {},
   "source": [
    "## Data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a413a9b",
   "metadata": {},
   "source": [
    "Let's use PASCAL VOC, which already comes with `torchvision`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62ba84db",
   "metadata": {},
   "outputs": [],
   "source": [
    "class VOC(torch.utils.data.Dataset):\n",
    "    def __init__(self, root, fold, transform=None, download=True):\n",
    "        super().__init__()\n",
    "        fold = 'test' if fold == 'val' else fold\n",
    "        self.ds = torchvision.datasets.VOCDetection(root, image_set=fold, download=download)\n",
    "        self.transform = transform\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.ds)\n",
    "\n",
    "    def __getitem__(self, i):\n",
    "        image, xml = self.ds[i]\n",
    "        image = np.array(image)\n",
    "        objs = xml['annotation']['object']\n",
    "        labels = [o['name'] for o in objs]\n",
    "        bboxes = [(\n",
    "            float(o['bndbox']['xmin']), float(o['bndbox']['ymin']),\n",
    "            float(o['bndbox']['xmax']), float(o['bndbox']['ymax']),\n",
    "            ) for o in objs]\n",
    "        d = {'image': image, 'bboxes': bboxes, 'labels': labels}\n",
    "        if self.transform:\n",
    "            d = self.transform(**d)\n",
    "        return d"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07984224",
   "metadata": {},
   "source": [
    "Let's detect only certain classes, such as animals."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66ddc663",
   "metadata": {},
   "outputs": [],
   "source": [
    "class FilterClass(torch.utils.data.Dataset):\n",
    "    def __init__(self, ds, whitelist):\n",
    "        super().__init__()\n",
    "        self.ds = ds\n",
    "        self.ix = [i for i in range(len(ds)) if any(label in whitelist for label in ds[i]['labels'])]\n",
    "        self.whitelist = whitelist\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.ix)\n",
    "\n",
    "    def __getitem__(self, i):\n",
    "        d = self.ds[self.ix[i]]\n",
    "        d['bboxes'] = [bbox for label, bbox in zip(d['labels'], d['bboxes']) if label in self.whitelist]\n",
    "        d['labels'] = [self.whitelist.index(label) for label in d['labels'] if label in self.whitelist]\n",
    "        return d"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bac4479c",
   "metadata": {},
   "source": [
    "Testing..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea63db58",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "animals = ['bird', 'cat', 'cow', 'dog', 'horse', 'sheep']\n",
    "voc = VOC(data_path, 'train', download=download)\n",
    "voc = FilterClass(voc, animals)\n",
    "data = voc[0]\n",
    "plt.imshow(data['image'])\n",
    "od.draw.bboxes(data['bboxes'], labels=data['labels'])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35c9a8a6",
   "metadata": {},
   "source": [
    "## Model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "219acdb9",
   "metadata": {},
   "source": [
    "We will implement the following model which is based on [FCOS](https://arxiv.org/abs/1904.01355). Although we are not going to support multi-scale grids or anchors here, we will separate the model into a `Grid` and `Model` class so that, if needed, you may more easily add grid multi-scale or anchors. (You may find implementations using multi-scale and anchors under the folder `implementations`.)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8b054ba",
   "metadata": {},
   "source": [
    "![Model diagram](model.svg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5265d264",
   "metadata": {},
   "outputs": [],
   "source": [
    "bboxes_loss = torchvision.ops.generalized_box_iou_loss\n",
    "centerness_loss = torch.nn.BCEWithLogitsLoss()\n",
    "labels_loss = torchvision.ops.sigmoid_focal_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2155a7ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Grid(torch.nn.Module):\n",
    "    def __init__(self, in_channels, nclasses, img_size):\n",
    "        super().__init__()\n",
    "        self.img_size = img_size\n",
    "        # like FCOS, we do not have a dedicated 'scores' prediction. it's just\n",
    "        # the argmax of the classes.\n",
    "        self.classes = torch.nn.Conv2d(in_channels, nclasses, 1)\n",
    "        self.bboxes = torch.nn.Conv2d(in_channels, 4, 1)\n",
    "        self.centerness = torch.nn.Conv2d(in_channels, 1, 1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        # like FCOS, the network is predicting bboxes in relative terms, we need\n",
    "        # to convert to absolute bboxes because the loss requires so.\n",
    "        bboxes = torch.exp(self.bboxes(x))\n",
    "        bboxes = od.transforms.rel_bboxes(bboxes, self.img_size)\n",
    "        return {'labels': self.classes(x), 'bboxes': bboxes,\n",
    "            'centerness': self.centerness(x)}\n",
    "\n",
    "    def post_process(self, preds, threshold=0.05):\n",
    "        scores, labels = torch.sigmoid(preds['labels']).max(1, keepdim=True)\n",
    "        bboxes = preds['bboxes']\n",
    "        centerness = torch.sigmoid(preds['centerness'])\n",
    "        mask = scores[:, 0] >= threshold\n",
    "        # like FCOS, centerness will help NMS choose the best bbox.\n",
    "        scores = scores * centerness\n",
    "        return {\n",
    "            'scores': od.grid.mask_select(mask, scores, True),\n",
    "            'bboxes': od.grid.mask_select(mask, bboxes, True),\n",
    "            'labels': od.grid.mask_select(mask, labels, True),\n",
    "        }\n",
    "\n",
    "    def compute_loss(self, preds, targets):\n",
    "        grid_size = preds['bboxes'].shape[2:]\n",
    "        mask, indices = od.grid.where(od.grid.slice_all_center, targets['bboxes'], grid_size, self.img_size)\n",
    "        # preds grid -> list\n",
    "        pred_bboxes = od.grid.mask_select(mask, preds['bboxes'])\n",
    "        pred_labels = od.grid.mask_select(mask, preds['labels'])\n",
    "        pred_centerness = od.grid.mask_select(mask, preds['centerness'])\n",
    "        # targets list -> list\n",
    "        target_bboxes = od.grid.indices_select(indices, targets['bboxes'])\n",
    "        target_labels = od.grid.indices_select(indices, targets['labels'])\n",
    "        # labels: must be one-hot since we use independent classifiers\n",
    "        target_labels = torch.nn.functional.one_hot(target_labels.long(),\n",
    "            preds['labels'].shape[1]).float()\n",
    "        # compute centerness: requires doing the transformation in grid-space\n",
    "        target_bboxes_grid = od.grid.to_grid(mask, indices, targets['bboxes'])\n",
    "        target_rel_bboxes = od.transforms.rel_bboxes(target_bboxes_grid, self.img_size)\n",
    "        target_centerness = od.transforms.centerness(target_rel_bboxes)\n",
    "        target_centerness = od.grid.mask_select(mask, target_centerness)\n",
    "        # compute losses\n",
    "        return bboxes_loss(pred_bboxes, target_bboxes).mean() + \\\n",
    "            labels_loss(pred_labels, target_labels).mean() + \\\n",
    "            centerness_loss(pred_centerness, target_centerness)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1e0b665",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Model(torch.nn.Module):\n",
    "    def __init__(self, nclasses, img_size):\n",
    "        super().__init__()\n",
    "        resnet = torchvision.models.resnet50(weights='DEFAULT')\n",
    "        self.backbone = torch.nn.Sequential(*list(resnet.children())[:-2])\n",
    "        self.grid = Grid(2048, nclasses, img_size)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.backbone(x)\n",
    "        return self.grid(x)\n",
    "\n",
    "    def post_process(self, x):\n",
    "        return self.grid.post_process(x)\n",
    "\n",
    "    def compute_loss(self, preds, targets):\n",
    "        return self.grid.compute_loss(preds, targets)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76ca5b23",
   "metadata": {},
   "source": [
    "## Data Augmentation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68e7fda2",
   "metadata": {},
   "outputs": [],
   "source": [
    "img_size = (256, 256)\n",
    "transform = A.Compose([\n",
    "    A.Resize(int(img_size[0]*1.1), int(img_size[1]*1.1)),\n",
    "    A.RandomCrop(*img_size),\n",
    "    A.HorizontalFlip(p=0.5),\n",
    "    A.RandomBrightnessContrast(p=1),\n",
    "    A.Normalize(),\n",
    "    ToTensorV2(),\n",
    "], bbox_params=A.BboxParams(format='pascal_voc', label_fields=['labels']))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b214cb47",
   "metadata": {},
   "source": [
    "## Training loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9eda199",
   "metadata": {},
   "outputs": [],
   "source": [
    "animals = ['bird', 'cat', 'cow', 'dog', 'horse', 'sheep']\n",
    "tr = VOC(data_path, 'train', transform, download)\n",
    "tr = FilterClass(tr, animals)\n",
    "tr = torch.utils.data.DataLoader(tr, 4, True, collate_fn=od.utils.collate_fn, num_workers=2, pin_memory=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0cf75b5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "K = len(animals)\n",
    "model = Model(K, img_size).to(device)\n",
    "opt = torch.optim.Adam(model.parameters(), 1e-4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d5bdfc2",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.train()\n",
    "epochs = 100\n",
    "for epoch in range(epochs):\n",
    "    tic = time()\n",
    "    avg_loss = 0\n",
    "    for images, targets in tr:\n",
    "        targets['bboxes'] = [bb.float() for bb in targets['bboxes']]\n",
    "        targets = {k: [v.to(device) for v in l] for k, l in targets.items()}\n",
    "        preds = model(images.to(device))\n",
    "        loss_value = model.compute_loss(preds, targets)\n",
    "        opt.zero_grad()\n",
    "        loss_value.backward()\n",
    "        opt.step()\n",
    "        avg_loss += float(loss_value) / len(tr)\n",
    "    toc = time()\n",
    "    print(f'Epoch {epoch+1}/{epochs} - {toc-tic:.0f}s - Avg loss: {avg_loss}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "495fae64",
   "metadata": {},
   "source": [
    "If you wish to evaluate the results, you may do so after the model is trained, or even inside the training loop..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8bb7b1e",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.eval()\n",
    "images, targets = next(iter(tr))\n",
    "preds_grid = model(images.to(device))\n",
    "preds = model.post_process(preds_grid)\n",
    "preds = od.post.NMS(preds)\n",
    "i = 0\n",
    "mean = torch.tensor([0.485, 0.456, 0.406])[None, None]\n",
    "std = torch.tensor([0.229, 0.224, 0.225])[None, None]\n",
    "plt.clf()\n",
    "plt.imshow(images[i].permute(1, 2, 0)*std+mean)\n",
    "od.draw.bboxes(preds['bboxes'][i].detach().cpu(), labels=[f'{int(l)} ({int(s*100)})' for l, s in zip(preds['labels'][i], preds['scores'][i])], color='cyan')\n",
    "od.draw.bboxes(targets['bboxes'][i], labels=[int(l) for l in targets['labels'][i]])\n",
    "plt.suptitle(f'Epoch {epoch+1} - Avg loss: {avg_loss}')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "72be8866",
   "metadata": {},
   "source": [
    "If you like, you can move the previous code block to inside the training loop itself. In such a case, we recommend either saving each image `plt.savefig()` or replace `plt.show()` with the following code to display it in a non-blocking fashion:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a744015e",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.show(block=False)\n",
    "plt.pause(0.1)"
   ]
  }
 ],
 "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.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}