create-custom-coco-dataset.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import glob\n",
    "\n",
    "from src.create_annotations import *\n",
    "\n",
    "# Label ids of the dataset\n",
    "category_ids = {\n",
    "    \"outlier\": 0,\n",
    "    \"window\": 1,\n",
    "    \"wall\": 2,\n",
    "    \"balcony\": 3,\n",
    "    \"door\": 4,\n",
    "    \"roof\": 5,\n",
    "    \"sky\": 6,\n",
    "    \"shop\": 7,\n",
    "    \"chimney\": 8\n",
    "}\n",
    "\n",
    "# Define which colors match which categories in the images\n",
    "category_colors = {\n",
    "    \"(0, 0, 0)\": 0, # Outlier\n",
    "    \"(255, 0, 0)\": 1, # Window\n",
    "    \"(255, 255, 0)\": 2, # Wall\n",
    "    \"(128, 0, 255)\": 3, # Balcony\n",
    "    \"(255, 128, 0)\": 4, # Door\n",
    "    \"(0, 0, 255)\": 5, # Roof\n",
    "    \"(128, 255, 255)\": 6, # Sky\n",
    "    \"(0, 255, 0)\": 7, # Shop\n",
    "    \"(128, 128, 128)\": 8 # Chimney\n",
    "}\n",
    "\n",
    "# Define the ids that are a multiplolygon. In our case: wall, roof and sky\n",
    "multipolygon_ids = [2, 5, 6]\n",
    "\n",
    "# Get \"images\" and \"annotations\" info \n",
    "def images_annotations_info(maskpath):\n",
    "    # This id will be automatically increased as we go\n",
    "    annotation_id = 0\n",
    "    image_id = 0\n",
    "    annotations = []\n",
    "    images = []\n",
    "    \n",
    "    for mask_image in glob.glob(maskpath + \"*.png\"):\n",
    "        # The mask image is *.png but the original image is *.jpg.\n",
    "        # We make a reference to the original file in the COCO JSON file\n",
    "        original_file_name = os.path.basename(mask_image).split(\".\")[0] + \".jpg\"\n",
    "\n",
    "        # Open the image and (to be sure) we convert it to RGB\n",
    "        mask_image_open = Image.open(mask_image).convert(\"RGB\")\n",
    "        w, h = mask_image_open.size\n",
    "        \n",
    "        # \"images\" info \n",
    "        image = create_image_annotation(original_file_name, w, h, image_id)\n",
    "        images.append(image)\n",
    "\n",
    "        sub_masks = create_sub_masks(mask_image_open, w, h)\n",
    "        for color, sub_mask in sub_masks.items():\n",
    "            category_id = category_colors[color]\n",
    "\n",
    "            # \"annotations\" info\n",
    "            polygons, segmentations = create_sub_mask_annotation(sub_mask)\n",
    "\n",
    "            # Check if we have classes that are a multipolygon\n",
    "            if category_id in multipolygon_ids:\n",
    "                # Combine the polygons to calculate the bounding box and area\n",
    "                multi_poly = MultiPolygon(polygons)\n",
    "                                \n",
    "                annotation = create_annotation_format(multi_poly, segmentations, image_id, category_id, annotation_id)\n",
    "\n",
    "                annotations.append(annotation)\n",
    "                annotation_id += 1\n",
    "            else:\n",
    "                for i in range(len(polygons)):\n",
    "                    # Cleaner to recalculate this variable\n",
    "                    segmentation = [np.array(polygons[i].exterior.coords).ravel().tolist()]\n",
    "                    \n",
    "                    annotation = create_annotation_format(polygons[i], segmentation, image_id, category_id, annotation_id)\n",
    "                    \n",
    "                    annotations.append(annotation)\n",
    "                    annotation_id += 1\n",
    "        image_id += 1\n",
    "    return images, annotations, annotation_id\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    # Get the standard COCO JSON format\n",
    "    coco_format = get_coco_json_format()\n",
    "    \n",
    "    for keyword in [\"train\", \"val\"]:\n",
    "        mask_path = \"dataset/{}_mask/\".format(keyword)\n",
    "        \n",
    "        # Create category section\n",
    "        coco_format[\"categories\"] = create_category_annotation(category_ids)\n",
    "    \n",
    "        # Create images and annotations sections\n",
    "        coco_format[\"images\"], coco_format[\"annotations\"], annotation_cnt = images_annotations_info(mask_path)\n",
    "\n",
    "        with open(\"output/{}.json\".format(keyword),\"w\") as outfile:\n",
    "            json.dump(coco_format, outfile)\n",
    "        \n",
    "        print(\"Created %d annotations for images in folder: %s\" % (annotation_cnt, mask_path))\n"
   ]
  }
 ],
 "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}