detectron2-balloons.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9e26548e",
   "metadata": {},
   "source": [
    "# Balloons Toy Dataset + Detectron2 + 3LC Tutorial"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59632039",
   "metadata": {},
   "source": [
    "<div style=\"display: inline-flex; align-items: center; gap: 10px;\">\n",
    "        <a href=\"https://colab.research.google.com/github/3lc-ai/notebook-examples/blob/main/detectron2-balloons.ipynb\"\n",
    "        target=\"_blank\"\n",
    "            style=\"background-color: transparent; text-decoration: none; display: inline-flex; align-items: center;\n",
    "            padding: 5px 10px; font-family: Arial, sans-serif;\"> <img\n",
    "            src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\" style=\"height: 30px;\n",
    "            vertical-align: middle;box-shadow: none;\"/>\n",
    "        </a> <a href=\"https://github.com/3lc-ai/notebook-examples/blob/main/detectron2-balloons.ipynb\"\n",
    "            style=\"text-decoration: none; display: inline-flex; align-items: center; background-color: #ffffff; border:\n",
    "            1px solid #d1d5da; border-radius: 8px; padding: 2px 10px; color: #333; font-family: Arial, sans-serif;\">\n",
    "            <svg aria-hidden=\"true\" focusable=\"false\" role=\"img\" class=\"octicon octicon-mark-github\" viewBox=\"0 0 16 16\"\n",
    "            width=\"20\" height=\"20\" fill=\"#333\"\n",
    "            style=\"display:inline-block;user-select:none;vertical-align:text-bottom;overflow:visible; margin-right:\n",
    "            8px;\">\n",
    "                <path d=\"M8 0c4.42 0 8 3.58 8 8a8.013 8.013 0 0 1-5.45 7.59c-.4.08-.55-.17-.55-.38 0-.27.01-1.13.01-2.2\n",
    "                0-.75-.25-1.23-.54-1.48 1.78-.2 3.65-.88 3.65-3.95 0-.88-.31-1.59-.82-2.15.08-.2.36-1.02-.08-2.12 0\n",
    "                0-.67-.22-2.2.82-.64-.18-1.32-.27-2-.27-.68 0-1.36.09-2 .27-1.53-1.03-2.2-.82-2.2-.82-.44 1.1-.16\n",
    "                1.92-.08 2.12-.51.56-.82 1.28-.82 2.15 0 3.06 1.86 3.75 3.64 3.95-.23.2-.44.55-.51\n",
    "                1.07-.46.21-1.61.55-2.33-.66-.15-.24-.6-.83-1.23-.82-.67.01-.27.38.01.53.34.19.73.9.82 1.13.16.45.68\n",
    "                1.31 2.69.94 0 .67.01 1.3.01 1.49 0 .21-.15.45-.55.38A7.995 7.995 0 0 1 0 8c0-4.42 3.58-8 8-8Z\"></path>\n",
    "            </svg> <span style=\"vertical-align: middle; color: #333;\">Open in GitHub</span>\n",
    "        </a>\n",
    "</div>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "fb9fde68",
   "metadata": {
    "id": "QHnVupBBn9eR"
   },
   "source": [
    "This notebook is a modified version of the official colab tutorial of detectron which can be found [here](https://colab.research.google.com/drive/16jcaJoc6bCFAQ96jDe2HwtXj7BMD_-m5).\n",
    "\n",
    "In this tutorial we will see how to fine-tune a pre-trained detectron model for object detection on a custom dataset in the COCO format.\n",
    "We will integrate with 3LC by creating a training run, registering 3LC datasets, and collecting per-sample bounding box metrics.\n",
    "\n",
    "This notebook demonstrates:\n",
    "\n",
    "+ Training a detectron2 model on a custom dataset.\n",
    "+ Integrating a COCO dataset with 3LC using `register_coco_instances()`.\n",
    "+ Collecting per-sample bounding box metrics using `BoundingBoxMetricsCollector`.\n",
    "+ Registering a custom per-sample metrics collection callback."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "72511eb4",
   "metadata": {},
   "source": [
    "## Setup Project"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c363054",
   "metadata": {
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "PROJECT_NAME = \"Balloons\"\n",
    "RUN_NAME = \"Train Balloon Detector\"\n",
    "DESCRIPTION = \"Train a balloon detector using detectron2\"\n",
    "TRAIN_DATASET_NAME = \"balloons-train\"\n",
    "VAL_DATASET_NAME = \"balloons-val\"\n",
    "TRANSIENT_DATA_PATH = \"../transient_data\"\n",
    "TEST_DATA_PATH = \"./data\"\n",
    "MODEL_CONFIG = \"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml\"\n",
    "MAX_ITERS = 200\n",
    "BATCH_SIZE = 2\n",
    "MAX_DETECTIONS_PER_IMAGE = 30\n",
    "SCORE_THRESH_TEST = 0.5\n",
    "TLC_PUBLIC_EXAMPLES_DEVELOPER_MODE = True\n",
    "EPOCHS = 1\n",
    "INSTALL_DEPENDENCIES = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2798b98",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "if INSTALL_DEPENDENCIES:\n",
    "    # NOTE: There is no single version of detectron2 that is appropriate for all users and all systems.\n",
    "    #       This notebook uses a particular prebuilt version of detectron2 that is only available for\n",
    "    #       Linux and for specific versions of torch, torchvision, and CUDA. It may not be appropriate\n",
    "    #       for your system. See https://detectron2.readthedocs.io/en/latest/tutorials/install.html for\n",
    "    #       instructions on how to install or build a version of detectron2 for your system.\n",
    "    %pip install torch==1.10.1+cu111 torchvision==0.11.2+cu111 -f https://download.pytorch.org/whl/cu111/torch_stable.html\n",
    "    %pip install detectron2 -f \"https://dl.fbaipublicfiles.com/detectron2/wheels/cu111/torch1.10/index.html\"\n",
    "    %pip install 3lc\n",
    "    %pip install opencv-python\n",
    "    %pip install matplotlib"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b24a3322",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "471922b0",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "0d288Z2mF5dC",
    "outputId": "c47c5426-64d6-4632-f868-e2f14dfe39be"
   },
   "outputs": [],
   "source": [
    "import detectron2\n",
    "import tlc\n",
    "import torch\n",
    "\n",
    "TORCH_VERSION = \".\".join(torch.__version__.split(\".\")[:2])\n",
    "CUDA_VERSION = torch.__version__.split(\"+\")[-1]\n",
    "print(\"tlc: \", tlc.__version__)\n",
    "print(\"torch: \", TORCH_VERSION, \"; cuda: \", CUDA_VERSION)\n",
    "print(\"detectron2:\", detectron2.__version__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c7aa9314",
   "metadata": {
    "id": "ZyAvNCJMmvFF"
   },
   "outputs": [],
   "source": [
    "# Some basic setup:\n",
    "from __future__ import annotations\n",
    "\n",
    "# Setup detectron2 logger\n",
    "from detectron2.utils.logger import setup_logger\n",
    "\n",
    "logger = setup_logger()\n",
    "logger.setLevel(\"ERROR\")\n",
    "\n",
    "# import some common libraries\n",
    "import os\n",
    "import cv2\n",
    "import random\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# import some common detectron2 utilities\n",
    "from detectron2 import model_zoo\n",
    "from detectron2.config import get_cfg\n",
    "from detectron2.utils.visualizer import Visualizer\n",
    "from detectron2.data import MetadataCatalog, DatasetCatalog"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "9643de5c",
   "metadata": {
    "id": "b2bjrfb2LDeo"
   },
   "source": [
    "## Prepare the dataset\n",
    "\n",
    "In this section, we show how to train an existing detectron2 model on a custom dataset in the COCO format.\n",
    "\n",
    "We use [the balloon segmentation dataset](https://github.com/matterport/Mask_RCNN/tree/master/samples/balloon)\n",
    "which only has one class: balloon. \n",
    "\n",
    "You can find a [modified COCO version](https://github.com/3lc-ai/notebook-examples/tree/main/data/balloons) of this dataset inside the \"data\" directory provided while cloning our [repository](https://github.com/3lc-ai/notebook-examples/).\n",
    "\n",
    "We'll train a balloon segmentation model from an existing model pre-trained on COCO dataset, available in detectron2's model zoo.\n",
    "\n",
    "Note that COCO dataset does not have the \"balloon\" category. We'll be able to recognize this new class in a few minutes."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "23d0cce3",
   "metadata": {},
   "source": [
    "## Register the dataset with 3LC\n",
    "\n",
    "Now that we have the dataset in the COCO format, we can register it with 3LC."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5cb4a75",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tlc.integration.detectron2 import register_coco_instances\n",
    "\n",
    "register_coco_instances(\n",
    "    TRAIN_DATASET_NAME,\n",
    "    {},\n",
    "    train_json_path.to_str(),\n",
    "    train_image_folder.to_str(),\n",
    "    project_name=PROJECT_NAME,\n",
    ")\n",
    "\n",
    "register_coco_instances(\n",
    "    VAL_DATASET_NAME,\n",
    "    {},\n",
    "    val_json_path.to_str(),\n",
    "    val_image_folder.to_str(),\n",
    "    project_name=PROJECT_NAME,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c72ccde9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# The detectron2 dataset dicts and dataset metadata can be read from the DatasetCatalog and\n",
    "# MetadataCatalog, respectively.\n",
    "dataset_metadata = MetadataCatalog.get(TRAIN_DATASET_NAME)\n",
    "dataset_dicts = DatasetCatalog.get(TRAIN_DATASET_NAME)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "2a9a4656",
   "metadata": {
    "id": "6ljbWTX0Wi8E"
   },
   "source": [
    "To verify the dataset is in correct format, let's visualize the annotations of randomly selected samples in the training set:\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44d3f524",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "UkNbUzUOLYf0",
    "outputId": "4f5ed932-624a-4ede-9d5b-22371569fe1d"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from detectron2.utils.file_io import PathManager\n",
    "\n",
    "for d in random.sample(dataset_dicts, 3):\n",
    "    filename = tlc.Url(d[\"file_name\"]).to_absolute().to_str()\n",
    "    if \"s3://\" in filename:\n",
    "        with PathManager.open(filename, \"rb\") as f:\n",
    "            img = np.asarray(bytearray(f.read()), dtype=\"uint8\")\n",
    "            img = cv2.imdecode(img, cv2.IMREAD_COLOR)\n",
    "    else:\n",
    "        img = cv2.imread(filename)\n",
    "    visualizer = Visualizer(img[:, :, ::-1], metadata=dataset_metadata, scale=0.5)\n",
    "    out = visualizer.draw_dataset_dict(d)\n",
    "    out_rgb = cv2.cvtColor(out.get_image(), cv2.COLOR_BGR2RGB)\n",
    "    plt.imshow(out_rgb[:, :, ::-1])\n",
    "    plt.title(filename.split(\"/\")[-1])\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8aab54a2",
   "metadata": {},
   "source": [
    "## Create a custom metrics collection function\n",
    "\n",
    "We will use a BoundingBoxMetricsCollection to collect per-sample bounding box metrics.\n",
    "This allows users to supply a custom function to collect the metrics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae579d17",
   "metadata": {},
   "outputs": [],
   "source": [
    "def custom_bbox_metrics_collector(\n",
    "    gts: list[tlc.COCOGroundTruth], preds: list[tlc.COCOPrediction], metrics: dict[str, list] | None = None\n",
    ") -> None:\n",
    "    \"\"\"Example function that computes custom metrics for bounding box detection.\"\"\"\n",
    "\n",
    "    # Lets just return the number of ground truth boxes and predictions\n",
    "    num_gts = [len(gt[\"annotations\"]) for gt in gts]\n",
    "    num_preds = [len(pred[\"annotations\"]) for pred in preds]\n",
    "\n",
    "    metrics[\"num_gts\"] = num_gts\n",
    "    metrics[\"num_preds\"] = num_preds"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "0cfca07d",
   "metadata": {
    "id": "wlqXIXXhW8dA"
   },
   "source": [
    "## Train!\n",
    "\n",
    "Now, let's fine-tune a COCO-pretrained R50-FPN Mask R-CNN model on the balloon dataset. It takes ~2 minutes to train 300 iterations on a P100 GPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ac782fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tlc\n",
    "\n",
    "run = tlc.init(\n",
    "    project_name=PROJECT_NAME,\n",
    "    run_name=RUN_NAME,\n",
    "    description=DESCRIPTION,\n",
    "    if_exists=\"overwrite\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e7d2947",
   "metadata": {},
   "outputs": [],
   "source": [
    "# For a full list of config values: https://github.com/facebookresearch/detectron2/blob/main/detectron2/config/defaults.py\n",
    "cfg = get_cfg()\n",
    "cfg.merge_from_file(model_zoo.get_config_file(MODEL_CONFIG))\n",
    "cfg.DATASETS.TRAIN = (TRAIN_DATASET_NAME,)\n",
    "cfg.DATASETS.TEST = (VAL_DATASET_NAME,)\n",
    "cfg.DATALOADER.NUM_WORKERS = 0\n",
    "cfg.OUTPUT_DIR = TRANSIENT_DATA_PATH\n",
    "cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(MODEL_CONFIG)  # Let training initialize from model zoo\n",
    "cfg.SOLVER.IMS_PER_BATCH = BATCH_SIZE  # This is the real \"batch size\" commonly known to deep learning people\n",
    "cfg.SOLVER.BASE_LR = 0.00025  # pick a good LR\n",
    "cfg.SOLVER.MAX_ITER = (\n",
    "    MAX_ITERS  # Seems good enough for this toy dataset; you will need to train longer for a practical dataset\n",
    ")\n",
    "cfg.SOLVER.STEPS = []  # Do not decay learning rate\n",
    "cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = (\n",
    "    128  # The \"RoIHead batch size\". 128 is faster, and good enough for this toy dataset (default: 512)\n",
    ")\n",
    "cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1  # Only has one class (balloon).\n",
    "\n",
    "cfg.TEST.DETECTIONS_PER_IMAGE = MAX_DETECTIONS_PER_IMAGE\n",
    "cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = SCORE_THRESH_TEST\n",
    "cfg.MODEL.DEVICE = \"cuda\"\n",
    "cfg.DATALOADER.FILTER_EMPTY_ANNOTATIONS = False\n",
    "\n",
    "os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)\n",
    "\n",
    "config = {\n",
    "    \"model_config\": MODEL_CONFIG,\n",
    "    \"solver.ims_per_batch\": BATCH_SIZE,\n",
    "    \"test.detections_per_image\": MAX_DETECTIONS_PER_IMAGE,\n",
    "    \"model.roi_heads.score_thresh_test\": SCORE_THRESH_TEST,\n",
    "}\n",
    "\n",
    "run.set_parameters(config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80e9d027",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "7unkuuiqLdqd",
    "outputId": "ba1716cd-3f3b-401d-bae5-8fbbd2199d9c"
   },
   "outputs": [],
   "source": [
    "from detectron2.engine import DefaultTrainer\n",
    "from tlc.integration.detectron2 import DetectronMetricsCollectionHook, MetricsCollectionHook\n",
    "\n",
    "trainer = DefaultTrainer(cfg)\n",
    "\n",
    "metrics_collector = tlc.BoundingBoxMetricsCollector(\n",
    "    classes=dataset_metadata.thing_classes,\n",
    "    label_mapping=dataset_metadata.thing_dataset_id_to_contiguous_id,\n",
    "    iou_threshold=0.5,\n",
    "    compute_derived_metrics=True,\n",
    "    extra_metrics_fn=custom_bbox_metrics_collector,\n",
    ")\n",
    "\n",
    "# Add schemas for the custom metrics defined above\n",
    "metrics_collector.add_schema(\n",
    "    \"num_gts\",\n",
    "    tlc.Schema(value=tlc.Int32Value(value_min=0), description=\"The number of ground truth boxes\"),\n",
    ")\n",
    "metrics_collector.add_schema(\n",
    "    \"num_preds\",\n",
    "    tlc.Schema(value=tlc.Int32Value(value_min=0), description=\"The number of predicted boxes\"),\n",
    ")\n",
    "\n",
    "# Register the metrics collector with the trainer;\n",
    "# + Collect metrics on the training set every 50 iterations starting at iteration 0\n",
    "# + Collect metrics on the validation set after training\n",
    "# + Collect default detectron2 metrics every 5 iterations\n",
    "trainer.register_hooks(\n",
    "    [\n",
    "        MetricsCollectionHook(\n",
    "            dataset_name=TRAIN_DATASET_NAME,\n",
    "            metrics_collectors=[metrics_collector],\n",
    "            collection_frequency=50,\n",
    "            collection_start_iteration=0,\n",
    "            collect_metrics_after_train=True,\n",
    "        ),\n",
    "        MetricsCollectionHook(\n",
    "            dataset_name=VAL_DATASET_NAME,\n",
    "            metrics_collectors=[metrics_collector],\n",
    "            collect_metrics_after_train=True,\n",
    "        ),\n",
    "        DetectronMetricsCollectionHook(\n",
    "            run_url=run.url,\n",
    "            collection_frequency=5,\n",
    "        ),\n",
    "    ]\n",
    ")\n",
    "trainer.resume_or_load(resume=False)\n",
    "trainer.train()"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "provenance": [],
   "toc_visible": true
  },
  "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}