Mask_R_CNN_for_multi_class_custom_dataset.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Cd6zo0dZHLlm"
      },
      "source": [
        "**Segmentation study with Mask R-CNN for multi-class custom dataset**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "V1w-LKZ4HDOE"
      },
      "outputs": [],
      "source": [
        "import os \n",
        "os.getcwd() # Checking directory"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QJmxlApnHRuV"
      },
      "outputs": [],
      "source": [
        "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 json\n",
        "import datetime\n",
        "import skimage.draw\n",
        "\n",
        "# Root directory of the project\n",
        "ROOT_DIR = os.path.abspath(\"\")\n",
        "\n",
        "# Import Mask RCNN\n",
        "sys.path.append(ROOT_DIR)  # To find local version of the library\n",
        "\n",
        "from mrcnn.config import Config\n",
        "from mrcnn import utils\n",
        "import mrcnn.model as modellib\n",
        "from mrcnn import visualize\n",
        "from mrcnn.model import log\n",
        "\n",
        "%matplotlib inline \n",
        "\n",
        "# Directory to save logs and trained model\n",
        "MODEL_DIR = os.path.join(ROOT_DIR, \"logfiles\")\n",
        "\n",
        "# Local path to trained weights file\n",
        "COCO_MODEL_PATH = os.path.join(ROOT_DIR, \"mask_rcnn_coco.h5\")\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": "markdown",
      "metadata": {
        "id": "RbnLvfo5Hdhm"
      },
      "source": [
        "**Set the configuration parameters of each object**\n",
        "## Configurations"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Yd4eprU9HZIc"
      },
      "outputs": [],
      "source": [
        "class classnameConfig(Config):\n",
        "    \"\"\"Configuration for training on the custom dataset.\n",
        "    Derives from the base Config class and overrides some values.\n",
        "    \"\"\"\n",
        "    # Give the configuration a recognizable name\n",
        "    NAME = \"configclassname\"\n",
        "\n",
        "    # Specify how many graphics cards will be used in the background while the neural networks are running for the model. (The default is usually 1 or 2)\n",
        "    GPU_COUNT = 1 \n",
        "\n",
        "    # We use a GPU with 12GB memory, which can fit two images.\n",
        "    # Adjust down if you use a smaller GPU.\n",
        "    IMAGES_PER_GPU = 2 \n",
        "\n",
        "    # Number of classes (including background)\n",
        "    NUM_CLASSES = 1 + 5  # Background + Number of Classes. For instance, if there are 5 classes, specify 1+5.\n",
        "\n",
        "    # Number of training steps per epoch\n",
        "    STEPS_PER_EPOCH = 100\n",
        "\n",
        "    # Skip detections with < 90% confidence\n",
        "    DETECTION_MIN_CONFIDENCE = 0.9\n",
        "    \n",
        "    DETECTION_NMS_THRESHOLD = 0.3\n",
        "\n",
        "config = classnameConfig()\n",
        "config.display()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4nQ5K3fYIIrs"
      },
      "source": [
        "**After installing TensorFlow GPU, check it.**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gEFF9gHDIExt"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "tf.test.gpu_device_name()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bbDIP-SUIRNc"
      },
      "source": [
        "### Notebook Preferences"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iYEniLpXIPle"
      },
      "outputs": [],
      "source": [
        "def get_ax(rows=1, cols=1, size=16):\n",
        "    \"\"\"Return a Matplotlib Axes array to be used in\n",
        "    all visualizations in the notebook. Provide a\n",
        "    central point to control graph sizes.\n",
        "    \n",
        "    Change the default size attribute to control the size\n",
        "    of rendered images\n",
        "    \"\"\"\n",
        "    _, ax = plt.subplots(rows, cols, figsize=(size*cols, size*rows))\n",
        "    return ax"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DvPO-BUFIWUd"
      },
      "source": [
        "## Dataset\n",
        "### Create a synthetic dataset\n",
        "\n",
        "Extend the Dataset class and add a method to load the shapes dataset, load_shapes(), and override the following methods:\n",
        "\n",
        "*  load_image()\n",
        "*  load_mask()\n",
        "*  image_reference()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PQCWA6I9IUQ9"
      },
      "outputs": [],
      "source": [
        "class customDataset(utils.Dataset):\n",
        "\n",
        "    def load_class(self, dataset_dir, subset):\n",
        "        \"\"\"Load a subset of the custom dataset.\n",
        "        dataset_dir: Root directory of the dataset.\n",
        "        subset: Subset to load: train or val\n",
        "        \"\"\"\n",
        "        # Add classes. (for instance, 5 classes)\n",
        "        # configclassname specifies your superclass in json annotation file inside the \"Categories\" key.\n",
        "        self.add_class(\"configclassname\", 1, \"class1\")\n",
        "        self.add_class(\"configclassname\", 2, \"class2\")\n",
        "        self.add_class(\"configclassname\", 3, \"class3\")\n",
        "        self.add_class(\"configclassname\", 4, \"class4\")\n",
        "        self.add_class(\"configclassname\", 5, \"class5\")\n",
        "\n",
        "        # Train or validation dataset?\n",
        "        assert subset in [\"train\", \"val\"]  # or test\n",
        "        dataset_dir = os.path.join(dataset_dir, subset)\n",
        "\n",
        "        # Load annotations\n",
        "        # VGG Image Annotator (up to version 1.6) saves each image in the form:\n",
        "        # { 'filename': '28503151_5b5b7ec140_b.jpg',\n",
        "        #   'regions': {\n",
        "        #       '0': {\n",
        "        #           'region_attributes': {},\n",
        "        #           'shape_attributes': {\n",
        "        #               'all_points_x': [...],\n",
        "        #               'all_points_y': [...],\n",
        "        #               'name': 'polygon'}},\n",
        "        #       ... more regions ...\n",
        "        #   },\n",
        "        #   'size': 100202\n",
        "        # }\n",
        "        # We mostly care about the x and y coordinates of each region\n",
        "        # Note: In VIA 2.0, regions was changed from a dict to a list.\n",
        "        \n",
        "        annotations1 = json.load(open(os.path.join(dataset_dir, \"dataset.json\")))\n",
        "        \n",
        "        # \"Annotation\" key elements in json file\n",
        "        annotations = list(annotations1.values())[2]  # don't need the dict keys\n",
        "\n",
        "        # \"Images\" key elements in json file\n",
        "        images= list(annotations1.values())[1]  # don't need the dict keys\n",
        "\n",
        "        # The VIA tool saves images in the JSON even if they don't have any\n",
        "        # annotations. Skip unannotated images.\n",
        "        \n",
        "\n",
        "        # Add images\n",
        "        for img in images:\n",
        "            annotations_img = [a for a in annotations if a['segmentation'] and a[\"image_id\"]==img[\"id\"]]\n",
        "            class_ = (a['category_id'] for a in annotations if a['category_id'] and a[\"image_id\"]==img[\"id\"])\n",
        "            # Get the x, y coordinaets of points of the polygons that make up\n",
        "            # the outline of each object instance. These are stores in the\n",
        "            # shape_attributes (see json format above)\n",
        "            # The if condition is needed to support VIA versions 1.x and 2.x.\n",
        "            polygons = list(annotations_img)\n",
        "            \n",
        "            # load_mask() needs the image size to convert polygons to masks.\n",
        "            # Unfortunately, VIA doesn't include it in JSON, so we must read\n",
        "            # the image. This is only managable since the dataset is tiny.\n",
        "            #fname=dict(list(filter(lambda person: person['id'] == img[\"image_id\"], images))[0])[\"file_name\"]\n",
        "            #print(img[\"file_name\"]);\n",
        "            image_path = os.path.join(dataset_dir, img[\"file_name\"])\n",
        "            image = skimage.io.imread(image_path)\n",
        "            height, width = image.shape[:2]\n",
        "\n",
        "            self.add_image(\n",
        "                \"configclassname\", #custom super class name\n",
        "                image_id=img['id'],  # use file name as a unique image id\n",
        "                path=image_path,\n",
        "                width=width, \n",
        "                height=height,\n",
        "                polygons=polygons,\n",
        "                class_ids=np.array(list(class_)))\n",
        "\n",
        "\n",
        "    def load_mask(self, image_id):\n",
        "        \"\"\"Generate instance masks for an image.\n",
        "       Returns:\n",
        "        masks: A bool array of shape [height, width, instance count] with\n",
        "            one mask per instance.\n",
        "        class_ids: a 1D array of class IDs of the instance masks.\n",
        "        \"\"\"\n",
        "        # If not a balloon dataset image, delegate to parent class.\n",
        "        image_info = self.image_info[image_id]\n",
        "        if image_info[\"source\"] != \"configclassname\":\n",
        "            return super(self.__class__, self).load_mask(image_id)\n",
        "        class_ids = image_info['class_ids']\n",
        "        \n",
        "        # Convert polygons to a bitmap mask of shape\n",
        "        # [height, width, instance_count]\n",
        "        \n",
        "        mask = np.zeros([image_info[\"height\"], image_info[\"width\"], len(image_info[\"polygons\"])],dtype=np.uint8)\n",
        "        for i, p in enumerate(image_info[\"polygons\"]):\n",
        "            # Get indexes of pixels inside the polygon and set them to 1\n",
        "            list=p['segmentation'][0] #polygons\n",
        "            cat=p['category_id'] #categories\n",
        "            rr, cc = skimage.draw.polygon(list[1::2], list[0::2])\n",
        "            mask[rr, cc, i] = 1\n",
        "              \n",
        "        # Return mask, and array of class IDs of each instance. Since we have\n",
        "        # one class ID only, we return an array of 1s\n",
        "        \n",
        "        #return mask, np.ones([mask.shape[-1]], dtype=np.int32)\n",
        "        return mask, class_ids\n",
        "        \n",
        "    def image_reference(self, image_id):\n",
        "        \"\"\"Return the path of the image.\"\"\"\n",
        "        info = self.image_info[image_id]\n",
        "        if info[\"source\"] == \"configclassname\":\n",
        "            return info[\"path\"]\n",
        "        else:\n",
        "            super(self.__class__, self).image_reference(image_id)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "s5g1Q2wLI7nX"
      },
      "outputs": [],
      "source": [
        "# Get the dataset from the releases page\n",
        "# https://github.com/matterport/Mask_RCNN/releases\n",
        "\n",
        "dentiAssist_DIR = os.path.join(ROOT_DIR,\"dataset\")\n",
        "\n",
        "# Training dataset.\n",
        "dataset_train = customDataset()\n",
        "dataset_train.load_class(dentiAssist_DIR, \"train\")\n",
        "dataset_train.prepare()\n",
        "\n",
        "# Validation dataset\n",
        "dataset_val = customDataset()\n",
        "dataset_val.load_class(dentiAssist_DIR, \"val\")  # or test\n",
        "dataset_val.prepare()\n",
        "\n",
        "print(\"Train Image Count: {}\".format(len(dataset_train.image_ids)))\n",
        "print(\"Train Class Count: {}\".format(dataset_train.num_classes))\n",
        "for i, info in enumerate(dataset_train.class_info):\n",
        "    print(\"{:3}. {:50}\".format(i, info['name']))\n",
        "\n",
        "print(\"Val Image Count: {}\".format(len(dataset_val.image_ids)))\n",
        "print(\"Val Class Count: {}\".format(dataset_val.num_classes))\n",
        "for i, info in enumerate(dataset_val.class_info):\n",
        "    print(\"{:3}. {:50}\".format(i, info['name']))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oQ6SttkDJbGt"
      },
      "source": [
        "**Visualize of masks**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aELjPsFWJYU2"
      },
      "outputs": [],
      "source": [
        "# Load and display 5 random samples\n",
        "image_ids = np.random.choice(dataset_train.image_ids, 5)\n",
        "for image_id in image_ids:\n",
        "    image = dataset_train.load_image(image_id)\n",
        "    mask, class_ids = dataset_train.load_mask(image_id)\n",
        "    visualize.display_top_masks(image, mask, class_ids, dataset_train.class_names,2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JakG4_21JjAF"
      },
      "source": [
        "## Create model in training mode"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2AJwKV0dJg91"
      },
      "outputs": [],
      "source": [
        "model = modellib.MaskRCNN(mode=\"training\", config=config,\n",
        "                          model_dir=MODEL_DIR)\n",
        "                          \n",
        "# 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",
        "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(), by_name=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7tzeM8tqJpAF"
      },
      "source": [
        "## Training\n",
        "\n",
        "Train in two stages:\n",
        "\n",
        "\n",
        "1.   Only the heads. Here we're freezing all the backbone layers and training only the randomly initialized layers (i.e. the ones that we didn't use pre-trained weights from MS COCO). To train only the head layers, pass layers='heads' to the train() function.\n",
        "\n",
        "2.   Fine-tune all layers. For this simple example it's not necessary, but we're including it to show the process. Simply pass layers=\"all to train all layers.\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qCMLlmwqJm9d"
      },
      "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=100, \n",
        "            layers='heads')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rZpxvbJcJy6l"
      },
      "outputs": [],
      "source": [
        "# Fine tune all layers\n",
        "# Passing layers=\"all\" trains all layers. You can also \n",
        "# pass a regular expression to select which layers to\n",
        "# train by name pattern.\n",
        "model.train(dataset_train, dataset_val, \n",
        "            learning_rate=config.LEARNING_RATE / 10,\n",
        "            epochs=100, \n",
        "            layers=\"all\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QMXn9T6SJ2cF"
      },
      "outputs": [],
      "source": [
        "# Save weights\n",
        "# Typically not needed because callbacks save after every epoch\n",
        "# Uncomment to save manually\n",
        "model_path = os.path.join(MODEL_DIR, \"model.h5\")\n",
        "model.keras_model.save_weights(model_path)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Oy5V7y1WJ5dV"
      },
      "source": [
        "## Detection"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ATqbcXRJJ6qW"
      },
      "outputs": [],
      "source": [
        "class InferenceConfig(classnameConfig):\n",
        "    GPU_COUNT = 1\n",
        "    IMAGES_PER_GPU = 1\n",
        "\n",
        "inference_config = InferenceConfig()\n",
        "\n",
        "# Recreate the model in inference mode\n",
        "model = modellib.MaskRCNN(mode=\"inference\", \n",
        "                          config=inference_config,\n",
        "                          model_dir=MODEL_DIR)\n",
        "\n",
        "# Get path to saved weights\n",
        "# Either set a specific path or find last trained weights\n",
        "# model_path = os.path.join(ROOT_DIR, \".h5 file name here\")\n",
        "model_path = model.find_last()\n",
        "\n",
        "# Load trained weights\n",
        "print(\"Loading weights from \", model_path)\n",
        "model.load_weights(model_path, by_name=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "s7p_O6UuKEUl"
      },
      "source": [
        "## Test on a random image\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "So-VbtEuJ4Nd"
      },
      "outputs": [],
      "source": [
        "image_id=random.choice(dataset_val.image_ids)\n",
        "original_image, image_meta, gt_class_id, gt_bbox, gt_mask =\\\n",
        "    modellib.load_image_gt(dataset_val, inference_config, \n",
        "                           image_id, use_mini_mask=False)\n",
        "\n",
        "log(\"original_image\", original_image)\n",
        "log(\"image_meta\", image_meta)\n",
        "log(\"gt_class_id\", gt_class_id)\n",
        "log(\"gt_bbox\", gt_bbox)\n",
        "log(\"gt_mask\", gt_mask)\n",
        "\n",
        "visualize.display_instances(original_image, gt_bbox, gt_mask, gt_class_id, \n",
        "                            dataset_train.class_names, figsize=(16, 16))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6Xt88Vm-KIQG"
      },
      "outputs": [],
      "source": [
        "results = model.detect([original_image], verbose=1)\n",
        "\n",
        "r = results[0]\n",
        "visualize.display_instances(original_image, r['rois'], r['masks'], r['class_ids'], \n",
        "                            dataset_val.class_names, r['scores'], ax=get_ax())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TTnWf9DxKLTF"
      },
      "source": [
        "## Evaluation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1zLKwzJEKMxl"
      },
      "outputs": [],
      "source": [
        "# Compute VOC-Style mAP @ IoU=0.5\n",
        "# Running on 10 images. Increase for better accuracy.\n",
        "image_ids = dataset_val.image_ids\n",
        "APs = []\n",
        "for image_id in image_ids:\n",
        "    # Load image and ground truth data\n",
        "    image, image_meta, gt_class_id, gt_bbox, gt_mask =\\\n",
        "        modellib.load_image_gt(dataset_val, inference_config,\n",
        "                               image_id, use_mini_mask=False)\n",
        "    molded_images = np.expand_dims(modellib.mold_image(image, inference_config), 0)\n",
        "    # Run object detection\n",
        "    results = model.detect([image], verbose=0)\n",
        "    r = results[0]\n",
        "    # Compute AP\n",
        "    AP, precisions, recalls, overlaps =\\\n",
        "        utils.compute_ap(gt_bbox, gt_class_id, gt_mask,\n",
        "                         r[\"rois\"], r[\"class_ids\"], r[\"scores\"], r['masks'])\n",
        "    APs.append(AP)\n",
        "    \n",
        "print(\"mAP: \", np.mean(APs))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RgJXzf9lKPCd"
      },
      "outputs": [],
      "source": [
        "# Draw precision-recall curve\n",
        "AP, precisions, recalls, overlaps = utils.compute_ap(gt_bbox, gt_class_id, gt_mask,\n",
        "                                          r['rois'], r['class_ids'], r['scores'], r['masks'])\n",
        "visualize.plot_precision_recall(AP, precisions, recalls)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e2GxSi5TKQeA"
      },
      "outputs": [],
      "source": [
        "# Grid of ground truth objects and their predictions\n",
        "visualize.plot_overlaps(gt_class_id, r['class_ids'], r['scores'],\n",
        "                        overlaps, dataset_val.class_names)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Testing"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Read an image from local\n",
        "image = cv2.imread(\"sample2.jpg\") \n",
        "image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) \n",
        "\n",
        "# Detection\n",
        "r = model.detect([image], verbose=0) \n",
        " \n",
        "r = r[0]\n",
        "\n",
        "# Visualization\n",
        "mrcnn.visualize.display_instances(image=image,  \n",
        "                                  boxes=r['rois'],  \n",
        "                                  masks=r['masks'],  \n",
        "                                  class_ids=r['class_ids'],  \n",
        "                                  class_names=CLASS_NAMES,  \n",
        "                                  scores=r['scores'])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": []
    }
  ],
  "metadata": {
    "colab": {
      "name": "Mask R-CNN for multi-class custom dataset.ipynb",
      "private_outputs": true,
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}