Detectron2_custom_coco_data_instance_segmentation_marine_ver5_TRAIN_Colab.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Detectron2_custom_coco_data_instance_segmentation_marine_ver5_TRAIN.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QHnVupBBn9eR",
        "colab_type": "text"
      },
      "source": [
        "# [How to train Detectron2 with Custom COCO Datasets](https://www.dlology.com/blog/how-to-train-detectron2-with-custom-coco-datasets/) | DLology\n",
        "\n",
        "<img src=\"https://dl.fbaipublicfiles.com/detectron2/Detectron2-Logo-Horz.png\" width=\"500\">\n",
        "\n",
        "This notebook will help you get started with this framwork by training a instance segmentation model with your custom COCO datasets."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vM54r6jlKTII",
        "colab_type": "text"
      },
      "source": [
        "# Install detectron2"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9_FzH13EjseR",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "#!pip install -U torch torchvision\n",
        "!pip install git+https://github.com/facebookresearch/fvcore.git\n",
        "import torch, torchvision\n",
        "torch.__version__"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "b-i4hmGYk1dL",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "#!git clone https://github.com/facebookresearch/detectron2 detectron2_repo\n",
        "#!pip install -e detectron2_repo\n",
        "#!pip install detectron2 -f https://github.com/facebookresearch/detectron2\n",
        "# install detectron2:\n",
        "!pip install detectron2 -f https://dl.fbaipublicfiles.com/detectron2/wheels/cu100/index.html"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZyAvNCJMmvFF",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# You may need to restart your runtime prior to this, to let your installation take effect\n",
        "# Some basic setup\n",
        "# Setup detectron2 logger\n",
        "import detectron2\n",
        "from detectron2.utils.logger import setup_logger\n",
        "setup_logger()\n",
        "\n",
        "# import some common libraries\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import cv2\n",
        "from google.colab.patches import cv2_imshow\n",
        "\n",
        "# import some common detectron2 utilities\n",
        "from detectron2.engine import DefaultPredictor\n",
        "from detectron2.config import get_cfg\n",
        "from detectron2.utils.visualizer import Visualizer\n",
        "from detectron2.data import MetadataCatalog, DatasetCatalog"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b2bjrfb2LDeo",
        "colab_type": "text"
      },
      "source": [
        "# Train on a custom COCO dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tjbUIhSxUdm_",
        "colab_type": "text"
      },
      "source": [
        "In this section, we show how to train an existing detectron2 model on a custom dataset in a new format.\n",
        "\n",
        "We use [the marine segmentation dataset](https://github.com/Allopart/Maritme_Mask_RCNN)\n",
        "which only has 5 classes: data, buoy, land, sea, ship and sky.\n",
        "We'll train a segmentation model from an existing model pre-trained on the COCO dataset, available in detectron2's model zoo.\n",
        "\n",
        "Note that the COCO dataset does not have the \"data\", \"buoy\" and \"sea\" categories."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4Qg7zSVOulkb",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# download, decompress the data\n",
        "!wget https://cbpetro.s3.us-east-2.amazonaws.com/api/download/data.zip\n",
        "!unzip data.zip > /dev/null"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tVJoOm6LVJwW",
        "colab_type": "text"
      },
      "source": [
        "Register the fruits_nuts dataset to detectron2, following the [detectron2 custom dataset tutorial](https://detectron2.readthedocs.io/tutorials/datasets.html).\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Lnkg1PByUjGQ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from detectron2.data.datasets import register_coco_instances\n",
        "register_coco_instances(\"nautical_ecp\", {}, \"./data/trainval.json\", \"./data/images\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MWknKqWTWIw9",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "nautical_metadata = MetadataCatalog.get(\"nautical_ecp\")\n",
        "dataset_dicts = DatasetCatalog.get(\"nautical_ecp\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6ljbWTX0Wi8E",
        "colab_type": "text"
      },
      "source": [
        "To verify the data loading is correct, let's visualize the annotations of randomly selected samples in the training set:\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UkNbUzUOLYf0",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import random\n",
        "\n",
        "for d in random.sample(dataset_dicts, 1):\n",
        "    img = cv2.imread(d[\"file_name\"])\n",
        "    visualizer = Visualizer(img[:, :, ::-1], metadata=nautical_metadata, scale=0.5)\n",
        "    vis = visualizer.draw_dataset_dict(d)\n",
        "    cv2_imshow(vis.get_image()[:, :, ::-1])"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wlqXIXXhW8dA",
        "colab_type": "text"
      },
      "source": [
        "Now, let's fine-tune a coco-pretrained R50-FPN Mask R-CNN model on the Nautical dataset. It takes ~6 minutes to train 300 iterations on Colab's K80 GPU.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7unkuuiqLdqd",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from detectron2.engine import DefaultTrainer\n",
        "from detectron2.config import get_cfg\n",
        "from detectron2 import model_zoo\n",
        "import os\n",
        "\n",
        "cfg = get_cfg()\n",
        "cfg.merge_from_file(model_zoo.get_config_file(\"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml\"))\n",
        "cfg.DATASETS.TRAIN = (\"nautical_ecp\",)\n",
        "cfg.DATASETS.TEST = ()   # no metrics implemented for this dataset\n",
        "cfg.DATALOADER.NUM_WORKERS = 2\n",
        "cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(\"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml\")  # Let training initialize from model zoo\n",
        "cfg.SOLVER.IMS_PER_BATCH = 2\n",
        "cfg.SOLVER.BASE_LR = 0.002\n",
        "cfg.SOLVER.MAX_ITER = 300    # 300 iterations seems good enough, but you can certainly train longer\n",
        "cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128   # faster, and good enough for this toy dataset\n",
        "cfg.MODEL.ROI_HEADS.NUM_CLASSES = 3  # 3 classes (buoy, land, ship)\n",
        "\n",
        "#os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)\n",
        "os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)\n",
        "trainer = DefaultTrainer(cfg)\n",
        "trainer.resume_or_load(resume=False)\n",
        "trainer.train()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JY67ava3nsnh",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Look at training curves in tensorboard:\n",
        "%load_ext tensorboard\n",
        "%tensorboard --logdir output"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0e4vdDIOXyxF",
        "colab_type": "text"
      },
      "source": [
        "Now, we perform inference with the trained model on the nautical dataset. First, let's create a predictor using the model we just trained:\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ya5nEuMELeq8",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Directory where output files are written\n",
        "cfg.OUTPUT_DIR = \"./output\"\n",
        "\n",
        "cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, \"model_final.pth\")\n",
        "cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7   # set the testing threshold for this model\n",
        "cfg.DATASETS.TEST = (\"nautical_ecp\", )\n",
        "predictor = DefaultPredictor(cfg)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qWq1XHfDWiXO",
        "colab_type": "text"
      },
      "source": [
        "Then, we randomly select several samples to visualize the prediction results."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "U5LhISJqWXgM",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from detectron2.utils.visualizer import ColorMode\n",
        "\n",
        "for d in random.sample(dataset_dicts, 1):    \n",
        "    im = cv2.imread(d[\"file_name\"])\n",
        "    outputs = predictor(im)\n",
        "    v = Visualizer(im[:, :, ::-1],\n",
        "                   metadata=nautical_metadata, \n",
        "                   scale=0.8, \n",
        "                   instance_mode=ColorMode.IMAGE_BW   # remove the colors of unsegmented pixels\n",
        "    )\n",
        "    v = v.draw_instance_predictions(outputs[\"instances\"].to(\"cpu\"))\n",
        "    cv2_imshow(v.get_image()[:, :, ::-1])\n",
        "\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "t_bo0cypwllj",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "nautical_metadata"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wZ6lYrCqLLLW",
        "colab_type": "text"
      },
      "source": [
        "## Benchmark inference speed"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PxRHYcAC_Z0f",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import time\n",
        "times = []\n",
        "for i in range(20):\n",
        "    start_time = time.time()\n",
        "    outputs = predictor(im)\n",
        "    delta = time.time() - start_time\n",
        "    times.append(delta)\n",
        "mean_delta = np.array(times).mean()\n",
        "fps = 1 / mean_delta\n",
        "print(\"Average(sec):{:.2f},fps:{:.2f}\".format(mean_delta, fps))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WFMOqBbWEh5v",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Download the results\n",
        "#from google.colab import files\n",
        "#files.download('./output/metrics.json')\n",
        "#files.download('./output/model_final.pth')\n",
        "#files.download('./output/last_checkpoint')"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}