{
 "cells": [
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "---\n",
    "execute:\n",
    "  enabled: false\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=object-detection-dataset-md-1 fw=all -->\n",
    "\n",
    "# The Object Detection Dataset\n",
    "\n",
    "There is no small dataset such as MNIST and Fashion-MNIST in the field of object detection.\n",
    "In order to quickly demonstrate object detection models,\n",
    "we collected and labeled a small dataset.\n",
    "First, we took photos of free bananas from our office\n",
    "and generated\n",
    "1000 banana images with different rotations and sizes.\n",
    "Then we placed each banana image\n",
    "at a random position on some background image.\n",
    "In the end, we labeled bounding boxes for those bananas on the images.\n",
    "\n",
    "\n",
    "## Downloading the Dataset\n",
    "\n",
    "The banana detection dataset with all the image and\n",
    "csv label files can be downloaded directly from the Internet."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "object-detection-dataset-downloading-the-dataset-1"
   },
   "source": [
    "%matplotlib inline\n",
    "from d2l import jax as d2l\n",
    "import jax\n",
    "from jax import numpy as jnp\n",
    "import optax\n",
    "import numpy as np\n",
    "import os\n",
    "import pandas as pd\n",
    "from PIL import Image"
   ],
   "execution_count": 1,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "object-detection-dataset-downloading-the-dataset-2"
   },
   "source": [
    "\n",
    "d2l.DATA_HUB['banana-detection'] = (\n",
    "    d2l.DATA_URL + 'banana-detection.zip',\n",
    "    '5de26c8fce5ccdea9f91267273464dc968d20d72')"
   ],
   "execution_count": 2,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=object-detection-dataset-md-2 fw=all -->\n",
    "\n",
    "## Reading the Dataset\n",
    "\n",
    "We are going to read the banana detection dataset in the `read_data_bananas`\n",
    "function below.\n",
    "The dataset includes a csv file for\n",
    "object class labels and\n",
    "ground-truth bounding box coordinates\n",
    "at the upper-left and lower-right corners."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "object-detection-dataset-reading-the-dataset-1"
   },
   "source": [
    "\n",
    "def read_data_bananas(is_train=True):\n",
    "    \"\"\"Read the banana detection dataset images and labels.\"\"\"\n",
    "    from PIL import Image\n",
    "    data_dir = d2l.download_extract('banana-detection')\n",
    "    csv_fname = os.path.join(data_dir, 'bananas_train' if is_train\n",
    "                             else 'bananas_val', 'label.csv')\n",
    "    csv_data = pd.read_csv(csv_fname)\n",
    "    csv_data = csv_data.set_index('img_name')\n",
    "    images, targets = [], []\n",
    "    for img_name, target in csv_data.iterrows():\n",
    "        img = Image.open(\n",
    "            os.path.join(data_dir, 'bananas_train' if is_train else\n",
    "                         'bananas_val', 'images', f'{img_name}'))\n",
    "        images.append(jnp.array(img).transpose(2, 0, 1))\n",
    "        # Here `target` contains (class, upper-left x, upper-left y,\n",
    "        # lower-right x, lower-right y), where all the images have the same\n",
    "        # banana class (index 0)\n",
    "        targets.append(list(target))\n",
    "    return images, jnp.expand_dims(jnp.array(targets), axis=1) / 256"
   ],
   "execution_count": 3,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=object-detection-dataset-md-3 fw=all -->\n",
    "\n",
    "By using the `read_data_bananas` function to read images and labels,\n",
    "the following `BananasDataset` class\n",
    "will allow us to create a customized `Dataset` instance\n",
    "for loading the banana detection dataset."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "object-detection-dataset-reading-the-dataset-2"
   },
   "source": [
    "\n",
    "class BananasDataset:\n",
    "    \"\"\"A customized dataset to load the banana detection dataset.\"\"\"\n",
    "    def __init__(self, is_train):\n",
    "        self.features, self.labels = read_data_bananas(is_train)\n",
    "        print('read ' + str(len(self.features)) + (f' training examples' if\n",
    "              is_train else f' validation examples'))\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return (self.features[idx].astype(jnp.float32), self.labels[idx])\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.features)"
   ],
   "execution_count": 4,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=object-detection-dataset-md-4 fw=all -->\n",
    "\n",
    "Finally, we define\n",
    "the `load_data_bananas` function to return two\n",
    "data iterator instances for both the training and test sets.\n",
    "For the test dataset,\n",
    "there is no need to read it in random order."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "object-detection-dataset-reading-the-dataset-3"
   },
   "source": [
    "\n",
    "def load_data_bananas(batch_size):\n",
    "    \"\"\"Load the banana detection dataset.\"\"\"\n",
    "    train_dataset = BananasDataset(is_train=True)\n",
    "    val_dataset = BananasDataset(is_train=False)\n",
    "    train_iter = d2l.ArrayDataLoader(\n",
    "        jnp.stack(train_dataset.features), train_dataset.labels,\n",
    "        batch_size=batch_size, shuffle=True)\n",
    "    val_iter = d2l.ArrayDataLoader(\n",
    "        jnp.stack(val_dataset.features), val_dataset.labels,\n",
    "        batch_size=batch_size)\n",
    "    return train_iter, val_iter"
   ],
   "execution_count": 5,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=object-detection-dataset-md-5 fw=all -->\n",
    "\n",
    "Let's read a minibatch and print the shapes of\n",
    "both images and labels in this minibatch.\n",
    "The shape of the image minibatch,\n",
    "(batch size, number of channels, height, width),\n",
    "looks familiar:\n",
    "it is the same as in our earlier image classification tasks.\n",
    "The shape of the label minibatch is\n",
    "(batch size, $m$, 5),\n",
    "where $m$ is the largest possible number of bounding boxes\n",
    "that any image has in the dataset.\n",
    "\n",
    "Although computation in minibatches is more efficient,\n",
    "it requires that all the image examples\n",
    "contain the same number of bounding boxes to form a minibatch via concatenation.\n",
    "In general,\n",
    "images may have a varying number of bounding boxes;\n",
    "thus,\n",
    "images with fewer than $m$ bounding boxes\n",
    "will be padded with illegal bounding boxes\n",
    "until $m$ is reached.\n",
    "Then\n",
    "the label of each bounding box is represented by an array of length 5.\n",
    "The first element in the array is the class of the object in the bounding box,\n",
    "where -1 indicates an illegal bounding box for padding.\n",
    "The remaining four elements of the array are\n",
    "the ($x$, $y$)-coordinate values\n",
    "of the upper-left corner and the lower-right corner\n",
    "of the bounding box (the range is between 0 and 1).\n",
    "For the banana dataset,\n",
    "since there is only one bounding box on each image,\n",
    "we have $m=1$."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "object-detection-dataset-reading-the-dataset-4"
   },
   "source": [
    "batch_size, edge_size = 32, 256\n",
    "train_iter, _ = load_data_bananas(batch_size)\n",
    "batch = next(iter(train_iter))\n",
    "batch[0].shape, batch[1].shape"
   ],
   "execution_count": 6,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "read 1000 training examples\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "read 100 validation examples\n"
     ]
    },
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "((32, 3, 256, 256), (32, 1, 5))"
     },
     "metadata": {},
     "execution_count": 6
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=object-detection-dataset-md-6 fw=all -->\n",
    "\n",
    "## Demonstration\n",
    "\n",
    "Let's demonstrate ten images with their labeled ground-truth bounding boxes.\n",
    "We can see that the rotations, sizes, and positions of bananas vary across all these images.\n",
    "Of course, this is just a simple artificial dataset.\n",
    "In practice, real-world datasets are usually much more complicated."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "object-detection-dataset-demonstration"
   },
   "source": [
    "imgs = jnp.transpose(batch[0][:10], (0, 2, 3, 1)) / 255\n",
    "axes = d2l.show_images(imgs, 2, 5, scale=2)\n",
    "for ax, label in zip(axes, np.array(batch[1][:10])):\n",
    "    d2l.show_bboxes(ax, [label[0][1:5] * edge_size], colors=['w'])"
   ],
   "execution_count": 7,
   "outputs": [
    {
     "output_type": "display_data",
     "data": {
      "image/png": "[encoded data, 512056 characters]"
     },
     "metadata": {}
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=object-detection-dataset-md-7 fw=all -->\n",
    "\n",
    "## Summary\n",
    "\n",
    "* The banana detection dataset we collected can be used to demonstrate object detection models.\n",
    "* The data loading for object detection is similar to that for image classification. However, in object detection the labels also contain information of ground-truth bounding boxes, which is missing in image classification.\n",
    "\n",
    "\n",
    "## Exercises\n",
    "\n",
    "1. Demonstrate other images with ground-truth bounding boxes in the banana detection dataset. How do they differ with respect to bounding boxes and objects?\n",
    "1. Say that we want to apply data augmentation, such as random cropping, to object detection. How can it be different from that in image classification? Hint: what if a cropped image only contains a small portion of an object?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=object-detection-dataset-md-8 fw=jax -->\n",
    "\n",
    "[Discussions](https://d2l.discourse.group/t/1608)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "d2l-jax",
   "display_name": "d2l (JAX)",
   "language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
