versions/2022/datasets/python/dataloader_demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### PyTorch Dataset and Dataloader Demo\n",
    "\n",
    "We illustrate how to build a custom dataset and dataloader for object detection. \n",
    "\n",
    "We will use our collected and labeled images for object detection. Over 1,000 `640x480` RGB images were collected using an off-the-shelf USB camera (A4TECH PK-635G).\n",
    "\n",
    "The images were labeled using [VIA](https://www.robots.ox.ac.uk/~vgg/software/via/) and the image filenames and labels are stored in a CSV file. \n",
    "\n",
    "Before continuiing, please download the dataset from [here](https://bit.ly/adl2-ssd). Extract the dataset on the same directory as this file. The directory structure is something like this.\n",
    "\n",
    "```\n",
    "--> datasets --> python --> config.py\n",
    "                        --> dataloader_demo.ipynb\n",
    "                        --> drinks/\n",
    "                        --> label_utils.py\n",
    "                        --> sample_labels.png\n",
    "                        ...\n",
    "```\n",
    "\n",
    "**Note**: Before running this demo, please make sure that you have `wandb.ai` account. See our discussion on [`wandb.ai`](https://github.com/roatienza/Deep-Learning-Experiments/blob/master/versions/2022/tools/python/wandb_demo.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Sample image annotation\n",
    "\n",
    "A sample image annotation is shown below. There are only 3 categories: `Water`, `Soda`, and `Juice`. By default the backgroud is the first category. The bounding boxes and classs names as shown. Each bounding box is defined by 4 numbers. The numbers define 2 corners of the bounding box: xmin, xmax, ymin, and ymax in pixel coordinates.\n",
    "\n",
    "<img src=\"sample_labels.png\" width=\"640\" height=\"480\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Import** the required modules.\n",
    "\n",
    "[`label_utils`](label_utils.py) is a helper module for loading the CSV file and converting a label to class name. Basically, `0` is `background`, `1` is `Water`, `2` is `Soda` and `3` is `Juice`. It also contains helper functions to build the label dictionary from the CSV file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "import wandb\n",
    "import label_utils\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision import transforms\n",
    "from PIL import Image"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Login to and initialize** `wandb`. You will need to use your `wandb` API key to run this demo.\n",
    "\n",
    "We will use the following dataset and dataloader configuration. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Failed to detect the name of this notebook, you can set it manually with the WANDB_NOTEBOOK_NAME environment variable to enable code saving.\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mrowel\u001b[0m. Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mrowel\u001b[0m (\u001b[33mupeee\u001b[0m). Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n",
      "2022-09-29 09:41:42.036174: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcudart.so.11.0\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "Tracking run with wandb version 0.12.16"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Run data is saved locally in <code>/home/rowel/github/roatienza/Deep-Learning-Experiments/versions/2022/datasets/python/wandb/run-20220929_094135-2p2mvf8f</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Syncing run <strong><a href=\"https://wandb.ai/upeee/dataloader-project/runs/2p2mvf8f\" target=\"_blank\">smooth-waterfall-12</a></strong> to <a href=\"https://wandb.ai/upeee/dataloader-project\" target=\"_blank\">Weights & Biases</a> (<a href=\"https://wandb.me/run\" target=\"_blank\">docs</a>)<br/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "wandb.login()\n",
    "config = {\n",
    "    \"num_workers\": 4,\n",
    "    \"pin_memory\": True,\n",
    "    \"batch_size\": 32,\n",
    "    \"dataset\": \"drinks\",\n",
    "    \"train_split\": \"drinks/labels_train.csv\",\n",
    "    \"test_split\": \"drinks/labels_test.csv\",}\n",
    "run = wandb.init(project=\"dataloader-project\", entity=\"upeee\", config=config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Dataset and Dataloader for Custom Object Detection\n",
    "\n",
    "The dataset CSV file is a list of image filenames and their labels. The image filenames and their labels are stored in a CSV file using the following format.\n",
    "\n",
    "```\n",
    "frame,xmin,xmax,ymin,ymax,class_id\n",
    "0001000.jpg,310,445,104,443,1\n",
    "0000999.jpg,194,354,96,478,1\n",
    "0000998.jpg,105,383,134,244,1\n",
    "0000997.jpg,157,493,89,194,1\n",
    "0000996.jpg,51,435,207,347,1\n",
    "...\n",
    "```\n",
    "\n",
    "A label represents the coordinates of the object bounding box.\n",
    "\n",
    "We will build a dictionary of `path_to_image` to `label` mapping. The `label` is a tensor of the form `xmin,xmax,ymin,ymax,class_id`. There can be multiple labels for an image since there can be multiple objects in an image.\n",
    "\n",
    "The `ImageDataset` class is a custom dataset class that loads the images and labels using the dictionary. The `ImageDataset` class is a subclass of the abstract class `torch.utils.data.Dataset` that supports `__len__()` and `__getitem__()` methods. This is also known as **map-style** method. A dataset can also be **iterable-style** that supports the `__iter__()` method.\n",
    "\n",
    "Our train and test dataloaders use the `wandb` configuration. \n",
    "\n",
    "We also create a custom `collate_fn` function to handle the labels per image. `collate_fn` pads all labels in a mini-batch to the same size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train split len: 996\n",
      "Test split len: 51\n"
     ]
    }
   ],
   "source": [
    "test_dict, test_classes = label_utils.build_label_dictionary(\n",
    "    config['test_split'])\n",
    "train_dict, train_classes = label_utils.build_label_dictionary(\n",
    "    config['train_split'])\n",
    "\n",
    "\n",
    "class ImageDataset(torch.utils.data.Dataset):\n",
    "    def __init__(self, dictionary, transform=None):\n",
    "        self.dictionary = dictionary\n",
    "        self.transform = transform\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.dictionary)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        # retrieve the image filename\n",
    "        key = list(self.dictionary.keys())[idx]\n",
    "        # retrieve all bounding boxes\n",
    "        boxes = self.dictionary[key]\n",
    "        # open the file as a PIL image\n",
    "        img = Image.open(key)\n",
    "        # apply the necessary transforms\n",
    "        # transforms like crop, resize, normalize, etc\n",
    "        if self.transform:\n",
    "            img = self.transform(img)\n",
    "        \n",
    "        # return a list of images and corresponding labels\n",
    "        return img, boxes\n",
    "\n",
    "\n",
    "train_split = ImageDataset(train_dict, transforms.ToTensor())\n",
    "test_split = ImageDataset(test_dict, transforms.ToTensor())\n",
    "\n",
    "# This is approx 95/5 split\n",
    "print(\"Train split len:\", len(train_split))\n",
    "print(\"Test split len:\", len(test_split))\n",
    "\n",
    "# We do not have a validation split\n",
    "\n",
    "def collate_fn(batch):\n",
    "    maxlen = max([len(x[1]) for x in batch])\n",
    "    images = []\n",
    "    boxes = []\n",
    "    for i in range(len(batch)):\n",
    "        img, box = batch[i]\n",
    "        images.append(img)\n",
    "        # pad with zeros if less than maxlen\n",
    "        if len(box) < maxlen:\n",
    "            box = np.concatenate(\n",
    "                (box, np.zeros((maxlen-len(box), box.shape[-1]))), axis=0)\n",
    "\n",
    "        box = torch.from_numpy(box)\n",
    "        boxes.append(box)\n",
    "\n",
    "    return torch.stack(images, 0), torch.stack(boxes, 0)\n",
    "\n",
    "\n",
    "train_loader = DataLoader(train_split,\n",
    "                          batch_size=config['batch_size'],\n",
    "                          shuffle=True,\n",
    "                          num_workers=config['num_workers'],\n",
    "                          pin_memory=config['pin_memory'],\n",
    "                          collate_fn=collate_fn)\n",
    "\n",
    "test_loader = DataLoader(test_split,\n",
    "                         batch_size=config['batch_size'],\n",
    "                         shuffle=False,\n",
    "                         num_workers=config['num_workers'],\n",
    "                         pin_memory=config['pin_memory'],\n",
    "                         collate_fn=collate_fn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualizing sample data from train split\n",
    "\n",
    "We visualize sample images from the train split by creating a `wandb` table with one column to visualize an image and the objects using bounding boxes. The annotation is stored in a list of dictionaries named `dict`. One dictionary per image using `position`, `class_id`, `domain` and `box_caption` as keys. Please check the [`wandb` media documentation](https://docs.wandb.ai/guides/track/log/media) for more details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/rowel/anaconda3/lib/python3.7/site-packages/IPython/core/display.py:724: UserWarning: Consider using IPython.display.IFrame instead\n",
      "  warnings.warn(\"Consider using IPython.display.IFrame instead\")\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<iframe src=\"https://wandb.ai/upeee/dataloader-project/runs/2p2mvf8f?jupyter=true\" style=\"border:none;width:100%;height:1000px;\"></iframe>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Waiting for W&B process to finish... <strong style=\"color:green\">(success).</strong>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "51931a7014354a34a94e584aa6b65bfd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(children=(Label(value='14.967 MB of 14.967 MB uploaded (0.000 MB deduped)\\r'), FloatProgress(value=0.9999…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Synced <strong style=\"color:#cdcd00\">smooth-waterfall-12</strong>: <a href=\"https://wandb.ai/upeee/dataloader-project/runs/2p2mvf8f\" target=\"_blank\">https://wandb.ai/upeee/dataloader-project/runs/2p2mvf8f</a><br/>Synced 6 W&B file(s), 1 media file(s), 34 artifact file(s) and 1 other file(s)"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Find logs at: <code>./wandb/run-20220929_094135-2p2mvf8f/logs</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# sample one mini-batch\n",
    "images, boxes = next(iter(train_loader))\n",
    "# map of label to class name\n",
    "class_labels = {i: label_utils.index2class(i) for i in train_classes}\n",
    "\n",
    "run.display(height=1000)\n",
    "table = wandb.Table(columns=['Image'])\n",
    "\n",
    "# we use wandb to visualize the objects and bounding boxes\n",
    "for image, box in zip(images, boxes):\n",
    "    dict = []\n",
    "    for i in range(box.shape[0]):\n",
    "        if box[i, -1] == 0:\n",
    "            continue\n",
    "        dict_item = {}\n",
    "        dict_item[\"position\"] = {\n",
    "            \"minX\": box[i, 0].item(),\n",
    "            \"maxX\": box[i, 1].item(),\n",
    "            \"minY\": box[i, 2].item(),\n",
    "            \"maxY\": box[i, 3].item(),\n",
    "        }\n",
    "        dict_item[\"domain\"] = \"pixel\"\n",
    "        dict_item[\"class_id\"] = (int)(box[i, 4].item())\n",
    "        dict_item[\"box_caption\"] = label_utils.index2class(\n",
    "            dict_item[\"class_id\"])\n",
    "        dict.append(dict_item)\n",
    "\n",
    "    img = wandb.Image(image, boxes={\n",
    "        \"ground_truth\": {\n",
    "            \"box_data\": dict,\n",
    "            \"class_labels\": class_labels\n",
    "        }\n",
    "    })\n",
    "    table.add_data(img)\n",
    "\n",
    "wandb.log({\"train_loader\": table})\n",
    "wandb.finish()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "52772734322c44a04e342c358be70f2ff4d97da358cb3cd38ceb0f6066598be5"
  },
  "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}