openimage_extractor.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "chronic-oxford",
   "metadata": {},
   "source": [
    "## Open Images Dataset Extractor\n",
    "\n",
    "***Extract and download a subset of images for specific classes from Open Images Dataset together with bounding box labels***\n",
    "\n",
    "- Using Open Images Dataset version: `V6` - https://storage.googleapis.com/openimages/web/download.html\n",
    "- If running on SageMaker Notebook Instance use `conda_python3` kernel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "positive-sequence",
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget https://raw.githubusercontent.com/openimages/dataset/master/downloader.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "three-interval",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "import random\n",
    "import importlib\n",
    "\n",
    "from matplotlib import pyplot as plt\n",
    "import cv2\n",
    "\n",
    "import utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "vietnamese-injury",
   "metadata": {},
   "outputs": [],
   "source": [
    "importlib.reload(utils)"
   ]
  },
  {
   "cell_type": "raw",
   "id": "premier-pollution",
   "metadata": {},
   "source": [
    "!python downloader.py -h"
   ]
  },
  {
   "cell_type": "raw",
   "id": "inappropriate-discrimination",
   "metadata": {},
   "source": [
    "!wget https://storage.googleapis.com/openimages/v5/class-descriptions-boxable.csv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "urban-aging",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('class-descriptions-boxable.csv') as fin:\n",
    "    all_classes = fin.read().strip().split('\\n')\n",
    "print(f'Loaded {len(all_classes)} class names')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "comparative-alarm",
   "metadata": {},
   "outputs": [],
   "source": [
    "EXTRACT_CLASSES = ['Man', 'Woman']\n",
    "COLOURS = {c:[random.randint(0, 255) for _ in range(3)] for c in EXTRACT_CLASSES}\n",
    "DATA_ROOT = 'data'\n",
    "IMAGE_EXT = '.jpg'\n",
    "\n",
    "TRAIN_SAMPLE_COUNT = 10000\n",
    "VALIDATION_SAMPLE_COUNT = int(TRAIN_SAMPLE_COUNT / 10)\n",
    "\n",
    "TRAIN_SPLIT = ('train', 'https://storage.googleapis.com/openimages/v6/oidv6-train-annotations-bbox.csv', \n",
    "               TRAIN_SAMPLE_COUNT)\n",
    "VALIDATION_SPLIT = ('validation', 'https://storage.googleapis.com/openimages/v5/validation-annotations-bbox.csv', \n",
    "                    VALIDATION_SAMPLE_COUNT)\n",
    "\n",
    "CLASS_LOOKUP = {cl[0]: cl[1] for cl in [c.split(',') for c in all_classes] if cl[1] in EXTRACT_CLASSES}\n",
    "\n",
    "print(f'Selected classes: {EXTRACT_CLASSES}')"
   ]
  },
  {
   "cell_type": "raw",
   "id": "noticed-relay",
   "metadata": {},
   "source": [
    "!rm -rf $DATA_ROOT/validation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "finished-bankruptcy",
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract(split, url, count_per_class, download_images=True, class_lookup=CLASS_LOOKUP, data_root=DATA_ROOT):\n",
    "    print(f'--> Extract split [{split}]')\n",
    "    annot = os.path.split(url)[-1]\n",
    "    if not os.path.exists(annot):\n",
    "        print('\\t- download annotations file')\n",
    "        !wget $url\n",
    "\n",
    "    for cid in class_lookup.keys():\n",
    "        cls = class_lookup[cid]\n",
    "        print(f'- Processing class {cls}')\n",
    "        extracted = {}\n",
    "        to_download = []\n",
    "        remaining_count = count_per_class\n",
    "\n",
    "        print('\\t- extracting')\n",
    "        with open(annot) as fin:\n",
    "            rec = fin.readline()\n",
    "            while rec != '':\n",
    "                rec = rec.strip().split(',')\n",
    "                if len(rec) > 2 and cid == rec[2]:\n",
    "                    img_id = rec[0]\n",
    "\n",
    "                    # Get all the bounding boxes, even if we got the required image count\n",
    "                    if remaining_count > 0:\n",
    "                        if img_id not in extracted:\n",
    "                            extracted[img_id] = []\n",
    "                            to_download.append(img_id)\n",
    "                            remaining_count -= 1\n",
    "                        add = True\n",
    "                    else:\n",
    "                        add = img_id in extracted\n",
    "                    if add:\n",
    "                        # Source box is in (x1, x2, y1, y2), move around\n",
    "                        extracted[img_id].append(' '.join([rec[4], rec[6], rec[5], rec[7]]))  \n",
    "\n",
    "                rec = fin.readline()\n",
    "\n",
    "        split_path = os.path.join(data_root, split)  \n",
    "        class_path = os.path.join(split_path, cls)\n",
    "        label_path = os.path.join(class_path, 'NormalisedLabel')\n",
    "        os.makedirs(label_path, exist_ok=True)\n",
    "\n",
    "        image_list_path = os.path.join(split_path, f'{cls}_images.txt')\n",
    "        with open(image_list_path, 'w') as fout:\n",
    "            fout.write(split + '/' + f'\\n{split}/'.join(to_download))\n",
    "\n",
    "        for img_id in extracted.keys():\n",
    "            with open(os.path.join(label_path, f'{img_id}.txt'), 'w') as fout:\n",
    "                boxes = [f'{cls} {b}' for b in extracted[img_id]]\n",
    "                fout.write('\\n'.join(boxes))\n",
    "\n",
    "        if download_images:\n",
    "            print('\\t- download images')\n",
    "            !python downloader.py $image_list_path --download_folder=$class_path --num_processes=5\n",
    "\n",
    "        denorm_label_path = os.path.join(class_path, 'Label')\n",
    "        print(f'\\t- denormalise into {denorm_label_path}', end='') \n",
    "        if not os.path.exists(denorm_label_path):\n",
    "            os.makedirs(denorm_label_path, exist_ok=True)\n",
    "\n",
    "        count = 0\n",
    "        for img in os.listdir(class_path):\n",
    "            if img.endswith(IMAGE_EXT):\n",
    "                \n",
    "                img_id = os.path.splitext(img)[0]   \n",
    "                labels = extracted[img_id]\n",
    "                \n",
    "                denorm_labels = []\n",
    "                image = cv2.imread(os.path.join(class_path, img))\n",
    "                for box in labels:\n",
    "                    box = box.split(' ')\n",
    "                    denorm_labels.append(' '.join([cls, *[str(b) for b in utils.denormalise_box(box, image)]]))\n",
    "\n",
    "                with open(os.path.join(denorm_label_path, img.replace(IMAGE_EXT, '.txt')), 'w') as fout:\n",
    "                    fout.write('\\n'.join(denorm_labels))     \n",
    "                \n",
    "                count += 1\n",
    "        print(f' - {count} images')\n",
    "    print('Done')       "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "searching-forwarding",
   "metadata": {},
   "outputs": [],
   "source": [
    "CURRENT_SPLIT = TRAIN_SPLIT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "north-elephant",
   "metadata": {},
   "outputs": [],
   "source": [
    "extract(*CURRENT_SPLIT, download_images=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "provincial-roman",
   "metadata": {},
   "source": [
    "##### Show some samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "coastal-heater",
   "metadata": {},
   "outputs": [],
   "source": [
    "def show_samples(split, classes=EXTRACT_CLASSES, max_per_class=3):\n",
    "    show_count = 3\n",
    "    _, ax = plt.subplots(ncols=len(classes), nrows=max_per_class, figsize=(20, 20))\n",
    "\n",
    "    for cls_i, cls in enumerate(classes):\n",
    "        print(f'--> Processing class {cls}')\n",
    "        count = 0\n",
    "        images_path = os.path.join(DATA_ROOT, split, cls)\n",
    "        labels_path = os.path.join(images_path, 'Label')\n",
    "        color = COLOURS[cls]\n",
    "\n",
    "        for image_name in os.listdir(images_path):\n",
    "            if image_name.endswith('.jpg'):\n",
    "\n",
    "                print(f'--> Processing {image_name}')\n",
    "                image = cv2.imread(os.path.join(images_path, image_name))\n",
    "                with open(os.path.join(labels_path, image_name.replace('.jpg', '.txt'))) as fin:\n",
    "                    labels = fin.read().strip().split('\\n')\n",
    "                lbl = [l.split(' ') for l in labels]    \n",
    "\n",
    "                for c, x1, y1, x2, y2 in lbl:\n",
    "                    x1, y1, x2, y2 = int(float(x1)), int(float(y1)), int(float(x2)), int(float(y2))\n",
    "                    utils.draw_box(image, [x1, y1, x2, y2], c, color)\n",
    "\n",
    "                ax[count, cls_i].imshow(image)\n",
    "\n",
    "                count += 1\n",
    "                if count >= max_per_class:\n",
    "                    break\n",
    "    plt.show()            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ordinary-handling",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_samples(CURRENT_SPLIT[0], max_per_class=10)"
   ]
  }
 ],
 "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": 5
}