{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {"collapsed": true},
   "outputs": [],
   "source": [
    "The purpose of this notebook is to provide a brief introduction to the training data used for mask rcnn, as well as demonstrate tools for subsetting the data for use in testing and evaluation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os \n",
    "import sys\n",
    "sys.path.append(os.getcwd())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The annotations are stored in json format. Pycocotools reads and parses these annotations.\n",
    "\n",
    "COCO is used by pycocotools, while the COCOSubsetter is written specifically for this notebook for the purpose of creating smaller data subsets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils.coco import COCOSubsetter\n",
    "from pycocotools.coco import COCO\n",
    "\n",
    "data_root=\"/data\" # set your data root\n",
    "coco_subsetter = COCOSubsetter(data_root)\n",
    "coco = COCO(os.path.join(data_root, \"annotations\", \"instances_train2017.json\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Start by getting a random image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "image_id = np.random.choice(list(coco_subsetter.images.keys()))\n",
    "image_id"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can load the image itself by passing the image id to the images attribute in the subsetter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "plt.imshow(plt.imread(coco_subsetter.images[image_id]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The pycocotools lets us apply the ground truth masks to the image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(plt.imread(coco_subsetter.images[image_id]))\n",
    "plt.axis('off')\n",
    "annIds = coco.getAnnIds(imgIds=[image_id])\n",
    "anns = coco.loadAnns(annIds)\n",
    "coco.showAnns(anns)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Pycocotools does not have a utility to display ground truth boxes, but this can be done using the patches module in matplotlib."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.patches as patches\n",
    "fig,ax = plt.subplots(1)\n",
    "ax.imshow(plt.imread(coco_subsetter.images[image_id]))\n",
    "for anno in anns:\n",
    "    rect = patches.Rectangle((anno['bbox'][0],anno['bbox'][1]),\n",
    "                             anno['bbox'][2],anno['bbox'][3],\n",
    "                             linewidth=1,edgecolor='r',facecolor='none')\n",
    "    ax.add_patch(rect)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Subsetting\n",
    "\n",
    "The steps below show how to create a small subset of the data. This is useful for debugging.\n",
    "\n",
    "Start by getting a random subset of 25 images. You can also use pycocotools to specify image ids based on category, or create your own custom list."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# select a subset of images\n",
    "images_sample = coco_subsetter.random_subset(25)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Place the 25 images and their annotations in a new directory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "os.makedirs(\"./data\", exist_ok=True)\n",
    "sample_dir = f'./data/sample_{time.time()}'\n",
    "coco_subsetter.create_subset(list(images_sample.keys()), sample_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We also might want to create a dataset that repeats a few images. Again, this can be useful for debugging.\n",
    "\n",
    "First, read back in the image and annotations for the small subset we just created, then run the duplicate dataset function with the number of duplicates, and the new directory to write to."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "coco_duplicator = COCOSubsetter(sample_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dup_dir = f'./data/dup_{time.time()}'\n",
    "coco_duplicator.duplicate_dataset(10, dup_dir)"
   ]
  }
 ],
 "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.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}