Mask_RCNN_on_OpenImagesV6_Subset.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Mask-RCNN_on_OpenImagesV6_Subset.ipynb",
      "private_outputs": true,
      "provenance": [],
      "collapsed_sections": [],
      "authorship_tag": "ABX9TyP4T/mgpl2T8HQ5ujVFRKid",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/GiX7000/object-detection/blob/main/Mask_RCNN_on_OpenImagesV6_Subset.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eQ81LNqhxhWu"
      },
      "source": [
        "# Implementation of Mask-RCNN on a subset of OpenImagesV6"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "17ATe5Huex2y"
      },
      "source": [
        "Use of the official [matterport's](https://github.com/matterport/Mask_RCNN) implementation."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3GlNyUL--oIw"
      },
      "source": [
        "# to run the model, we should use the TensorFlow 1.x version\n",
        "%tensorflow_version 1.x\n",
        "\n",
        "!pip3 uninstall -y keras\n",
        "!pip3 install keras==2.1.0"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "39OJJnttxf8P"
      },
      "source": [
        "# get the model\n",
        "!git clone https://github.com/matterport/Mask_RCNN\n",
        "%cd Mask_RCNN\n",
        "!pip3 install -r requirements.txt\n",
        "!python3 setup.py install"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "htJX6Z5exteH"
      },
      "source": [
        "# imports \n",
        "import os\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 xml.etree.ElementTree as ET\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",
        "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, \"logs\")\n",
        "\n",
        "# Local path to trained weights file\n",
        "COCO_MODEL_PATH = os.path.join(ROOT_DIR, \"mask_rcnn_coco.h5\")\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)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "648BcZ5Fxf5k"
      },
      "source": [
        "class DetectorConfig(Config):\n",
        "    # Give the configuration a recognizable name\n",
        "    NAME = \"object_detection\"\n",
        "    \n",
        "    # number of classes is 4\n",
        "    NUM_CLASSES = 1 + 4\n",
        "\n",
        "    GPU_COUNT = 1\n",
        "    IMAGES_PER_GPU = 1\n",
        "\n",
        "config = DetectorConfig()\n",
        "config.display()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9SWz99-GoXg6"
      },
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "W-3f9YRs2TEQ"
      },
      "source": [
        "Images_set1b is comprised of a specific number of images of specific classes from OpenImagesV6 dataset(45 images for train set and 10 for test set for 4 different classes). See more at OpenImagesV6_Subset.ipynb. To use the notebook on the other subsets, some little pre-process is needed to the datasets."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "K9qoABgIohPQ"
      },
      "source": [
        "# coppy images_set1b to workspace\n",
        "%cd /content/drive/MyDrive/\n",
        "!cp -r images_set1b /content/Mask_RCNN"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Kvxjoj0Trhst"
      },
      "source": [
        "# resize images. Use of resize_dataset_pascalvoc. Use this script to manually resize images and their corresponding xmls.\n",
        "\n",
        "#%cd /content/\n",
        "#!git clone https://github.com/italojs/resize_dataset_pascalvoc.git\n",
        "\n",
        "#%cd /content/resize_dataset_pascalvoc\n",
        "#!pip install -r requirements.txt\n",
        "\n",
        "#!python3 main.py -p /content/Mask_RCNN/images_set1b --output ./new_images_set1b --new_x 1024 --new_y 1024 --save_box_images 1\n",
        "\n",
        "#!mv new_images_set1b/ /content/Mask_RCNN"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "h8Hd6czr0yqa"
      },
      "source": [
        "class MyDataset(utils.Dataset):\n",
        "    def load_dataset(self, dataset_dir):\n",
        "        self.add_class('dataset', 1, 'Cattle')\n",
        "        self.add_class('dataset', 2, 'Horse')\n",
        "        self.add_class('dataset', 3, 'Sheep')\n",
        "        self.add_class('dataset', 4, 'Person')\n",
        "        \n",
        "        # find all images\n",
        "        for i, filename in enumerate(os.listdir(dataset_dir)):\n",
        "            if '.jpg' in filename:\n",
        "                self.add_image('dataset', \n",
        "                               image_id=i, \n",
        "                               path=os.path.join(dataset_dir, filename), \n",
        "                               annotation=os.path.join(dataset_dir, filename.replace('.jpg', '.xml')))\n",
        "    \n",
        "    # extract bounding boxes from an annotation file\n",
        "    def extract_boxes(self, filename):\n",
        "        # load and parse the file\n",
        "        tree = ET.parse(filename)\n",
        "        # get the root of the document\n",
        "        root = tree.getroot()\n",
        "        # extract each bounding box\n",
        "        boxes = []\n",
        "        classes = []\n",
        "        classes_list = ['cow', 'horse', 'sheep', 'person']\n",
        "        for member in root.findall('object'):\n",
        "            xmin = int(member[4][0].text)\n",
        "            ymin = int(member[4][1].text)\n",
        "            xmax = int(member[4][2].text)\n",
        "            ymax = int(member[4][3].text)\n",
        "            boxes.append([xmin, ymin, xmax, ymax])\n",
        "            classes.append(self.class_names.index(member[0].text))\n",
        "        # extract image dimensions\n",
        "        width = int(root.find('size')[0].text)\n",
        "        height = int(root.find('size')[1].text)\n",
        "        return boxes, classes, width, height\n",
        " \n",
        "    # load the masks for an image\n",
        "    def load_mask(self, image_id):\n",
        "        # get details of image\n",
        "        info = self.image_info[image_id]\n",
        "        # define box file location\n",
        "        path = info['annotation']\n",
        "        # load XML\n",
        "        boxes, classes, w, h = self.extract_boxes(path)\n",
        "        # create one array for all masks, each on a different channel\n",
        "        masks = np.zeros([h, w, len(boxes)], dtype='uint8')\n",
        "        # create masks\n",
        "        for i in range(len(boxes)):\n",
        "            box = boxes[i]\n",
        "            row_s, row_e = box[1], box[3]\n",
        "            col_s, col_e = box[0], box[2]\n",
        "            masks[row_s:row_e, col_s:col_e, i] = 1\n",
        "        return masks, np.asarray(classes, dtype='int32')\n",
        "    \n",
        "    def image_reference(self, image_id):\n",
        "        info = self.image_info[image_id]\n",
        "        return info['path']"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "v-kP7Yho07Z5"
      },
      "source": [
        "# create training and testing sets\n",
        "dataset_train = MyDataset()\n",
        "dataset_train.load_dataset('/content/Mask_RCNN/images_set1b/train')\n",
        "dataset_train.prepare()\n",
        "print('Train: %d' % len(dataset_train.image_ids))\n",
        " \n",
        "# test/val set\n",
        "dataset_val = MyDataset()\n",
        "dataset_val.load_dataset('/content/Mask_RCNN/images_set1b/test')\n",
        "dataset_val.prepare()\n",
        "print('Test: %d' % len(dataset_val.image_ids))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "23Po1AmV31hJ"
      },
      "source": [
        "# Load and display random samples\n",
        "image_ids = np.random.choice(dataset_train.image_ids, 4)\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)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "np6XCdT34hKU"
      },
      "source": [
        "# Create model in training mode\n",
        "model = modellib.MaskRCNN(mode=\"training\", config=config,\n",
        "                          model_dir=MODEL_DIR)\n",
        "\n",
        "# Which weights to start with?\n",
        "init_with = \"imagenet\"  # 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)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "N7viIl8S7xxC"
      },
      "source": [
        "# train the model, stage 1:\n",
        "\n",
        "# train the head branches \n",
        "model.train(dataset_train, dataset_val, \n",
        "            learning_rate=config.LEARNING_RATE, \n",
        "            epochs=5, \n",
        "            layers='heads')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "V7AcvA6PAXyq"
      },
      "source": [
        "# train the model, stage 2:\n",
        "\n",
        "# Fine tune all layers\n",
        "model.train(dataset_train, dataset_val, \n",
        "            learning_rate=config.LEARNING_RATE / 10,\n",
        "            epochs=10, \n",
        "            layers=\"all\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8VqMxoy8PS4C"
      },
      "source": [
        "# detection\n",
        "\n",
        "class InferenceConfig(DetectorConfig): \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)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2UJrx73DPSkt"
      },
      "source": [
        "def get_ax(rows=1, cols=1, size=8):\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"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wVWGTPo4CL5a"
      },
      "source": [
        "# Test on a random image\n",
        "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=(8, 8))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Avsik1XFCeZs"
      },
      "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())"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8XRLTKUaOkI3"
      },
      "source": [
        "# evaluation\n",
        "\n",
        "# evaluate on dataset_val images\n",
        "image_ids = dataset_val.image_ids\n",
        "APs = []\n",
        "ARs = []\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, r[\"rois\"], r[\"class_ids\"], r[\"scores\"], r['masks'], iou_threshold=0.5)\n",
        "    \n",
        "    AR, _ =\\\n",
        "        utils.compute_recall(r[\"rois\"], gt_bbox, iou=0.5)\n",
        "    ARs.append(AR)\n",
        "    APs.append(AP)\n",
        "\n",
        "mAP = np.mean(APs)\n",
        "mAR = np.mean(ARs)\n",
        "\n",
        "print(\"mAP: \", mAP)\n",
        "print(\"mAR: \", mAR)\n",
        "print(\"f1 score: \", 2 * ((mAP * mAR) / (mAP + mAR)))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Nau6Xx26DCrc"
      },
      "source": [
        "# start Tensorboard\n",
        "%load_ext tensorboard\n",
        "%tensorboard --logdir=/content/logs"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}