train_mapillary_AWS_pilot_run.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Mask R-CNN - Train on Mapillary Dataset\n",
    "\n",
    "\n",
    "This notebook shows how to train Mask R-CNN on your own dataset. To keep things simple we use a synthetic dataset of shapes (squares, triangles, and circles) which enables fast training. You'd still need a GPU, though, because the network backbone is a Resnet101, which would be too slow to train on a CPU. On a GPU, you can start to get okay-ish results in a few minutes, and good results in less than an hour.\n",
    "\n",
    "The code of the *Shapes* dataset is included below. It generates images on the fly, so it doesn't require downloading any data. And it can generate images of any size, so we pick a small image size to train faster. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import coco\n",
    "import sys\n",
    "import random\n",
    "import math\n",
    "import re\n",
    "import time\n",
    "import numpy as np\n",
    "import cv2\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "import skimage.io\n",
    "import json\n",
    "from PIL import Image\n",
    "from config import Config\n",
    "import utils\n",
    "import model as modellib\n",
    "import visualize\n",
    "from model import log\n",
    "\n",
    "%matplotlib inline \n",
    "\n",
    "# Root directory of the project\n",
    "ROOT_DIR = os.getcwd()\n",
    "\n",
    "# Directory to save logs and trained model\n",
    "MODEL_DIR = os.path.join(ROOT_DIR, \"logs\")\n",
    "\n",
    "# Local path to trained weights file\n",
    "COCO_MODEL_PATH = os.path.join(ROOT_DIR, \"mask_rcnn_coco.h5\")\n",
    "\n",
    "# local path to mapillary dataset\n",
    "DATASET_DIR = os.path.join(ROOT_DIR, \"mapillary_dataset\")\n",
    "\n",
    "# Download COCO trained weights from Releases if needed\n",
    "if not os.path.exists(COCO_MODEL_PATH):\n",
    "    utils.download_trained_weights(COCO_MODEL_PATH)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MapillaryConfig(coco.CocoConfig):\n",
    "    \"\"\"Configuration for training on the mapillary dataset.\n",
    "    Derives from the base Config class and overrides values specific\n",
    "    to the mapillary dataset.\n",
    "    \"\"\"\n",
    "    # Give the configuration a recognizable name\n",
    "    NAME = \"train_4096\"\n",
    "\n",
    "    # Train on 1 GPU and 8 images per GPU. We can put multiple images on each\n",
    "    # GPU because the images are small. Batch size is 8 (GPUs * images/GPU).\n",
    "    GPU_COUNT = 1\n",
    "    IMAGES_PER_GPU = 8\n",
    "\n",
    "    # Number of classes (including background)\n",
    "    # this MUST be explicitly defined, or will run into index out of bound error\n",
    "    NUM_CLASSES = 1 + 11  # background + 11 objects\n",
    "\n",
    "    # Use small images for faster training. Set the limits of the small side\n",
    "    # the large side, and that determines the image shape.\n",
    "    IMAGE_MIN_DIM = 1024\n",
    "    IMAGE_MAX_DIM = 1024\n",
    "\n",
    "    # Use smaller anchors because our image and objects are small\n",
    "    RPN_ANCHOR_SCALES = (8, 16, 32, 64, 128)  # anchor side in pixels\n",
    "\n",
    "    # Reduce training ROIs per image because the images are small and have\n",
    "    # few objects. Aim to allow ROI sampling to pick 33% positive ROIs.\n",
    "    TRAIN_ROIS_PER_IMAGE = 32\n",
    "\n",
    "    # Use a small epoch since the data is simple\n",
    "    STEPS_PER_EPOCH = 1000\n",
    "\n",
    "    # use small validation steps since the epoch is small\n",
    "    VALIDATION_STEPS = 50\n",
    "    \n",
    "    LEARNING_RATE = 0.0001\n",
    "    WEIGHT_DECAY = 0.0005\n",
    "    \n",
    "config = MapillaryConfig()\n",
    "config.display()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MapillaryDataset(utils.Dataset):\n",
    "    \"\"\"Generates the shapes synthetic dataset. The dataset consists of simple\n",
    "    shapes (triangles, squares, circles) placed randomly on a blank surface.\n",
    "    The images are generated on the fly. No file access required.\n",
    "    \"\"\"\n",
    "    \n",
    "    DEBUG = False\n",
    "    CLASS_MAP = {}\n",
    "    CLASSES = [\"Bird\", \"Person\", \"Bicyclist\", \"Motorcyclist\", \"Bench\", \\\n",
    "     \"Car\", \"Person\", \"Fire Hydrant\", \"Traffic Light\", \"Bus\", \"Motorcycle\", \"Truck\"]\n",
    "    \n",
    "    # local path to image folder, choose 'dev', 'training', or 'testing'\n",
    "    SUBSET_DIR = \"\"\n",
    "\n",
    "    # local path to images inside development folder\n",
    "    IMG_DIR = \"\"\n",
    "\n",
    "    # local path to instance annotations inside development folder\n",
    "    INS_DIR = \"\"\n",
    "\n",
    "\n",
    "    def load_mapillary(self, dataset_dir, subset, class_ids=None,\n",
    "                  class_map=None):\n",
    "        \n",
    "        self.SUBSET_DIR = os.path.join(dataset_dir, subset)\n",
    "        self.IMG_DIR = os.path.join(self.SUBSET_DIR, 'images')\n",
    "        self.INS_DIR = os.path.join(self.SUBSET_DIR, 'instances')\n",
    "        \n",
    "        # load classes, start with id = 1 to account for background \"BG\"\n",
    "        class_id = 1\n",
    "        for label_id, label in enumerate(class_ids):\n",
    "            if label[\"instances\"] == True and label[\"readable\"] in self.CLASSES:\n",
    "                self.CLASS_MAP[label_id] = class_id\n",
    "                \n",
    "                if (self.DEBUG):\n",
    "                    print(\"{}: Class {} {} added\".format(label_id, class_id, label[\"readable\"]))\n",
    "                    \n",
    "                self.add_class(\"mapillary\", class_id, label[\"readable\"])\n",
    "                class_id = class_id + 1\n",
    "                \n",
    "        # add images \n",
    "        file_names = next(os.walk(self.IMG_DIR))[2]\n",
    "        for i in range(len(file_names)):\n",
    "            if file_names[i] != '.DS_Store':\n",
    "                image_path = os.path.join(self.IMG_DIR, file_names[i])\n",
    "                base_image = Image.open(image_path)\n",
    "                w, h = base_image.size\n",
    "\n",
    "                if (self.DEBUG):\n",
    "                    print(\"Image {} {} x {} added\".format(file_names[i], w, h))\n",
    "\n",
    "                self.add_image(\"mapillary\", image_id = i,\n",
    "                              path = file_names[i],\n",
    "                               width = w,\n",
    "                               height = h\n",
    "                              )\n",
    "\n",
    "    def load_image(self, image_id):\n",
    "        \"\"\"Generate an image from the specs of the given image ID.\n",
    "        This function loads the image from a file.\n",
    "        \"\"\"\n",
    "        info = self.image_info[image_id]\n",
    "        img_path = os.path.join(self.IMG_DIR, info[\"path\"])\n",
    "        image = Image.open(img_path)\n",
    "        image_array = np.array(image)\n",
    "        return image_array\n",
    "\n",
    "    def image_reference(self, image_id):\n",
    "        \"\"\"Return the local directory path of the image.\"\"\"\n",
    "        info = self.image_info[image_id]\n",
    "        img_path = os.path.join(self.IMG_DIR, info[\"path\"])\n",
    "        return img_path\n",
    "\n",
    "    def load_mask(self, image_id):\n",
    "        \"\"\"Generate instance masks for shapes of the given image ID.\n",
    "        \"\"\"\n",
    "        info = self.image_info[image_id]\n",
    "        instance_path = os.path.join(self.INS_DIR, info[\"path\"])\n",
    "        instance_image = Image.open(instance_path.rsplit(\".\", 1)[0] + \".png\")\n",
    "        \n",
    "        # convert labeled data to numpy arrays for better handling\n",
    "        instance_array = np.array(instance_image, dtype=np.uint16)\n",
    "\n",
    "        instances = np.unique(instance_array)\n",
    "        instaces_count = instances.shape[0]\n",
    "        \n",
    "        label_ids = instances // 256\n",
    "        label_id_count = np.unique(label_ids).shape[0]\n",
    "        \n",
    "        if (self.DEBUG):\n",
    "            print(\"There are {} instances, {} classes labelled instances in the image {}.\"\\\n",
    "                  .format(instaces_count, label_id_count, info[\"path\"]))\n",
    "            \n",
    "        mask = np.zeros([instance_array.shape[0], instance_array.shape[1], instaces_count], dtype=np.uint8)\n",
    "        mask_count = 0\n",
    "        loaded_class_ids = []\n",
    "        for instance in instances:\n",
    "            label_id = instance//256\n",
    "            if (label_id in self.CLASS_MAP):\n",
    "                m = np.zeros((instance_array.shape[0], instance_array.shape[1]), dtype=np.uint8)\n",
    "                m[instance_array == instance] = 1\n",
    "                m_size = np.count_nonzero(m == 1)\n",
    "                \n",
    "                # only load mask greater than threshold size, \n",
    "                # otherwise bounding box with area zero causes program to crash\n",
    "                if m_size > 4096:\n",
    "                    mask[:, :, mask_count] = m\n",
    "                    loaded_class_ids.append(self.CLASS_MAP[label_id])\n",
    "                    mask_count = mask_count + 1\n",
    "                    if (self.DEBUG):\n",
    "                        print(\"Non-zero: {}\".format(m_size))\n",
    "                        print(\"Mask {} created for instance {} of class {} {}\"\\\n",
    "                              .format(mask_count, instance, self.CLASS_MAP[label_id], \\\n",
    "                                      self.class_names[self.CLASS_MAP[label_id]]))\n",
    "        mask = mask[:, :, 0:mask_count]\n",
    "        return mask, np.array(loaded_class_ids)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# read in config file\n",
    "with open(os.path.join(DATASET_DIR, 'config.json')) as config_file:\n",
    "    class_config = json.load(config_file)\n",
    "# in this example we are only interested in the labels\n",
    "labels = class_config['labels']\n",
    "        \n",
    "# Training dataset\n",
    "dataset_train = MapillaryDataset()\n",
    "dataset_train.load_mapillary(DATASET_DIR, \"train_107\", class_ids = labels)\n",
    "dataset_train.prepare()\n",
    "\n",
    "# Validation dataset\n",
    "dataset_val = MapillaryDataset()\n",
    "dataset_val.load_mapillary(DATASET_DIR, \"dev_9\", class_ids = labels)\n",
    "dataset_val.prepare()\n",
    "\n",
    "print(\"mapping: \", class_config[\"mapping\"])\n",
    "print(\"version: \", class_config[\"version\"])\n",
    "print(\"folder_structure:\", class_config[\"folder_structure\"])\n",
    "print(\"There are {} classes in the config file\".format(len(labels)))\n",
    "print(\"There are {} classes in the model\".format(len(dataset_train.class_names)))\n",
    "for i in range(len(dataset_train.class_names)):\n",
    "    print(\"    Class {}: {}\".format(i, dataset_train.class_names[i]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create model in training mode\n",
    "model = modellib.MaskRCNN(mode=\"training\", config=config,\n",
    "                          model_dir=MODEL_DIR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Which weights to start with?\n",
    "init_with = \"coco\"  # imagenet, coco, or last\n",
    "\n",
    "if init_with == \"imagenet\":\n",
    "    model.load_weights(model.get_imagenet_weights(), by_name=True)\n",
    "    \n",
    "elif init_with == \"coco\":\n",
    "    # Load weights trained on MS COCO, but skip layers that\n",
    "    # are different due to the different number of classes\n",
    "    # See README for instructions to download the COCO weights\n",
    "    model.load_weights(COCO_MODEL_PATH, by_name=True,\n",
    "                       exclude=[\"mrcnn_class_logits\", \"mrcnn_bbox_fc\", \n",
    "                                \"mrcnn_bbox\", \"mrcnn_mask\"])\n",
    "elif init_with == \"last\":\n",
    "    # Load the last model you trained and continue training\n",
    "    model.load_weights(model.find_last()[1], by_name=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Train the head branches\n",
    "# Passing layers=\"heads\" freezes all layers except the head\n",
    "# layers. You can also pass a regular expression to select\n",
    "# which layers to train by name pattern.\n",
    "model.train(dataset_train, dataset_val, \n",
    "            learning_rate=config.LEARNING_RATE, \n",
    "            epochs=2, \n",
    "            layers='heads')"
   ]
  },
  {
   "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.6.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}