openimage_2_rec.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Convert dataset from Open Images to RecordIO format\n",
    "\n",
    "\n",
    "- Based on Open Images Dataset format (https://storage.googleapis.com/openimages/web/download.html)\n",
    "- Use `openimage_extractor.ipynb` to extract and download a subset of data from Open Images if needed.\n",
    "\n",
    "I am using a collection of existing resources from GluonCV and MXNet to perform some tasks. 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 Open Images Dataset 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",
    "import glob\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "import numpy as np\n",
    "import cv2\n",
    "\n",
    "from gluoncv.data import COCODetection, RecordFileDetection\n",
    "from gluoncv.utils import viz\n",
    "\n",
    "import importlib\n",
    "import utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "importlib.reload(utils)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Specify your own classes here\n",
    "CLASSES = ['monitor', 'vase', 'camera']\n",
    "COLORS = [[random.randint(0, 255) for _ in range(3)] for _ in range(len(CLASSES))]\n",
    "\n",
    "DATA_ROOT = 'data/openimage_like_sample'  # This would be path to folder 2017 if you are using a full COCO 2017 dataset \n",
    "DATA_SPLIT_NAME = 'test'\n",
    "IMAGE_EXT = '.png'\n",
    "\n",
    "lst_dir_path = os.path.join(DATA_ROOT, DATA_SPLIT_NAME)\n",
    "lst_file_path = lst_dir_path + '.lst'\n",
    "rec_file_path = lst_file_path.replace('.lst', '.rec')\n",
    "\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 OpenImageDetection:\n",
    "    CLASSES = CLASSES\n",
    "    def __init__(self, root, split, image_ext='.jpg'):\n",
    "        self.root = root\n",
    "        self.split = split\n",
    "        self._image_ext = image_ext\n",
    "        self._anno_ext = '.txt'\n",
    "        self._items = self._load_items()\n",
    "        self._index_map = dict(zip(self.classes, range(len(self.classes))))\n",
    "        \n",
    "    def __str__(self):\n",
    "        return self.__class__.__name__ + '(' + self.split + ')'\n",
    "\n",
    "    @property\n",
    "    def classes(self):\n",
    "        return type(self).CLASSES\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self._items)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        img_id, image_path, _ = self._items[idx]\n",
    "        labels = self._load_labels(idx)\n",
    "        img = cv2.cvtColor(cv2.imread(image_path), cv2.COLOR_BGR2RGB)\n",
    "        return img, labels\n",
    "\n",
    "    def _load_items(self):\n",
    "        ids = []\n",
    "        \n",
    "        # The same image may appear in many classes\n",
    "        processed = []\n",
    "        class_root = os.path.join(self.root, self.split)\n",
    "        images = glob.glob(os.path.join(class_root, f'*/*' + self._image_ext))\n",
    "        for image_path in images:\n",
    "            image_id = os.path.splitext(os.path.split(image_path)[1])[0]\n",
    "            if image_id not in processed:\n",
    "                labels = glob.glob(os.path.join(class_root, f'*/Label/{image_id}.txt'))\n",
    "                if len(labels) == 0:\n",
    "                    print(f'WARNING: no labels found for image {image_path}')\n",
    "                else:\n",
    "                    ids.append([image_id, image_path, labels])\n",
    "                    processed.append(image_id)\n",
    "        return ids\n",
    "\n",
    "    def _load_labels(self, idx):\n",
    "        \n",
    "        out_labels = []\n",
    "        for annot_path in self._items[idx][-1]:\n",
    "            with open(annot_path) as fin:\n",
    "                labels = [l.split(' ') for l in fin.read().strip().split('\\n')]\n",
    "\n",
    "            for cls, x1, y1, x2, y2 in labels:\n",
    "                if cls not in self.classes:\n",
    "                    continue\n",
    "                out_labels.append([float(x1), float(y1), float(x2), float(y2), self._index_map[cls]])\n",
    "            \n",
    "        return np.array(out_labels)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Load the included dataset and create an `.lst` file\n",
    "\n",
    "Also display the dataset images with bounding boxes, you may want to skip plotting those images for large datasets"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "dataset = OpenImageDetection(root=DATA_ROOT, split=DATA_SPLIT_NAME, image_ext='.png')\n",
    "print('Dataset length:', len(dataset))\n",
    "sample_image_path = None\n",
    "sample_image = None\n",
    "count = 0\n",
    "idx = 0\n",
    "\n",
    "image, labels = dataset[idx]\n",
    "h, w = image.shape[:2]\n",
    "\n",
    "image_path = '/'.join(dataset._items[idx][1].split('/')[-2:])\n",
    "bboxes, cids = labels[:, :4], labels[:, 4:5]\n",
    "lst_record = utils.build_lst_record(image_path, w, h,  bboxes, cids, idx)\n",
    "\n",
    "for bbox, cid in zip(bboxes, cids):\n",
    "    cid = int(cid)\n",
    "    utils.draw_bbox(image, bbox.astype(np.int), CLASSES[cid], COLORS[cid])\n",
    "    \n",
    "plt.imshow(image)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f'Creating LST file {lst_file_path}')\n",
    "dataset = OpenImageDetection(root=DATA_ROOT, split=DATA_SPLIT_NAME, image_ext='.png')\n",
    "print('Dataset length:', len(dataset))\n",
    "\n",
    "with open(lst_file_path, 'w') as lst_out:\n",
    "    for idx in range(len(dataset)):\n",
    "        image, labels = dataset[idx]\n",
    "        h, w = image.shape[:2]\n",
    "        image_path = '/'.join(dataset._items[idx][1].split('/')[-2:])\n",
    "        bboxes, cids = labels[:, :4], labels[:, 4:5]\n",
    "        lst_record = utils.build_lst_record(image_path, 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://raw.githubusercontent.com/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 $lst_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 $lst_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
}