PlateCount_FasterRCNN.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "Check the torch version compatibility with detectron2"
      ],
      "metadata": {
        "id": "biry6xuCmYXn"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "aNYKF6ezpbKl",
        "outputId": "f807d5e3-98cf-41bc-d6b2-9a8e4aa5301b"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "2.0.1+cu118\n"
          ]
        }
      ],
      "source": [
        "!python3 -c \"import torch; print(torch.__version__)\""
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Install the detectron2 using pip"
      ],
      "metadata": {
        "id": "XpMDOQ7lmgNg"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ffu4GdHzpGaD"
      },
      "outputs": [],
      "source": [
        "!python -m pip install 'git+https://github.com/facebookresearch/detectron2.git'"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Import the required libraries"
      ],
      "metadata": {
        "id": "JXM_L1Xbmk8Y"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Sb2CKqP9oDZD"
      },
      "outputs": [],
      "source": [
        "import detectron2\n",
        "from detectron2.utils.logger import setup_logger\n",
        "setup_logger()\n",
        "\n",
        "# import some common libraries\n",
        "import numpy as np\n",
        "import cv2\n",
        "import random\n",
        "from google.colab.patches import cv2_imshow\n",
        "\n",
        "# import some common detectron2 utilities\n",
        "from detectron2 import model_zoo\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\n",
        "from detectron2.data.catalog import DatasetCatalog"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Import the data for training as well as evaluation"
      ],
      "metadata": {
        "id": "va1vBcM5mpVV"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "slYH9q2JxghG"
      },
      "outputs": [],
      "source": [
        "!curl -L \"https://app.roboflow.com/ds/3XWUSp6duU?key=eT77LEswhs\" > roboflow.zip; unzip roboflow.zip; rm roboflow.zip"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Register the dataset as instances of detectron2"
      ],
      "metadata": {
        "id": "2LZ_Q3usmwKZ"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CZfbcStMumHH"
      },
      "outputs": [],
      "source": [
        "from detectron2.data.datasets import register_coco_instances\n",
        "\n",
        "register_coco_instances(\"my_dataset_train\", {}, \"/content/train/_annotations.coco.json\", \"/content/train\")\n",
        "\n",
        "register_coco_instances(\"my_dataset_val\", {}, \"/content/valid/_annotations.coco.json\", \"/content/valid\")\n",
        "\n",
        "register_coco_instances(\"my_dataset_test\", {}, \"/content/test/_annotations.coco.json\", \"/content/test\")"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Visualize the Dataset"
      ],
      "metadata": {
        "id": "cuBXsEMwm3Mj"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_LU6JlBayHFM"
      },
      "outputs": [],
      "source": [
        "#visualize training data\n",
        "my_dataset_train_metadata = MetadataCatalog.get(\"my_dataset_train\")\n",
        "dataset_dicts = DatasetCatalog.get(\"my_dataset_train\")\n",
        "\n",
        "import random\n",
        "from detectron2.utils.visualizer import Visualizer\n",
        "\n",
        "for d in random.sample(dataset_dicts, 3):\n",
        "    img = cv2.imread(d[\"file_name\"])\n",
        "    visualizer = Visualizer(img[:, :, ::-1], metadata=my_dataset_train_metadata, scale=0.5)\n",
        "    vis = visualizer.draw_dataset_dict(d)\n",
        "    cv2_imshow(vis.get_image()[:, :, ::-1])"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Configure the CFG file for training"
      ],
      "metadata": {
        "id": "jfl6dhXtm_Ii"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "G45WIOFzyL2O"
      },
      "outputs": [],
      "source": [
        "cfg = get_cfg()\n",
        "cfg.merge_from_file(model_zoo.get_config_file(\"COCO-Detection/faster_rcnn_X_101_32x8d_FPN_3x.yaml\"))\n",
        "cfg.DATASETS.TRAIN = (\"my_dataset_train\",)\n",
        "cfg.DATASETS.TEST = (\"my_dataset_val\",)\n",
        "\n",
        "cfg.DATALOADER.NUM_WORKERS = 4\n",
        "cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(\"COCO-Detection/faster_rcnn_X_101_32x8d_FPN_3x.yaml\")  # Let training initialize from model zoo\n",
        "cfg.SOLVER.IMS_PER_BATCH = 4\n",
        "cfg.SOLVER.BASE_LR = 0.0001\n",
        "\n",
        "\n",
        "cfg.SOLVER.WARMUP_ITERS = 1000\n",
        "cfg.SOLVER.MAX_ITER = 500 #adjust up if val mAP is still rising, adjust down if overfit\n",
        "cfg.SOLVER.STEPS = (1000, 1500)\n",
        "cfg.SOLVER.GAMMA = 0.05\n",
        "\n",
        "cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 64\n",
        "cfg.MODEL.ROI_HEADS.NUM_CLASSES = 2\n",
        "\n",
        "cfg.TEST.EVAL_PERIOD = 500"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Define the training class"
      ],
      "metadata": {
        "id": "3wSNg7q6nE0t"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DgEZtJsPyTLs"
      },
      "outputs": [],
      "source": [
        "from detectron2.engine import DefaultTrainer\n",
        "from detectron2.evaluation import COCOEvaluator\n",
        "\n",
        "class CocoTrainer(DefaultTrainer):\n",
        "\n",
        "  @classmethod\n",
        "  def build_evaluator(cls, cfg, dataset_name, output_folder=None):\n",
        "\n",
        "    if output_folder is None:\n",
        "        os.makedirs(\"coco_eval\", exist_ok=True)\n",
        "        output_folder = \"coco_eval\"\n",
        "\n",
        "    return COCOEvaluator(dataset_name, cfg, False, output_folder)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Start the Training. Finally!"
      ],
      "metadata": {
        "id": "UuyNWQCUnKlE"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RJ0ZPZe9yW13"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "from IPython.lib.display import exists\n",
        "os.makedirs(cfg.OUTPUT_DIR, exist_ok = True)\n",
        "trainer = CocoTrainer(cfg)\n",
        "trainer.resume_or_load(resume=False)\n",
        "trainer.train()"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Evaluate the model performance on validation dataset"
      ],
      "metadata": {
        "id": "23qAIGGMnOtn"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "csT8Xysq_wT7"
      },
      "outputs": [],
      "source": [
        "from detectron2.data import DatasetCatalog, MetadataCatalog, build_detection_test_loader\n",
        "from detectron2.evaluation import COCOEvaluator, inference_on_dataset\n",
        "\n",
        "# cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, \"model_final.pth\")\n",
        "cfg.MODEL.WEIGHTS = '/content/model_final.pth'\n",
        "cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.85\n",
        "predictor = DefaultPredictor(cfg)\n",
        "evaluator = COCOEvaluator(\"my_dataset_test\", cfg, False, output_dir=\"./output/\")\n",
        "val_loader = build_detection_test_loader(cfg, \"my_dataset_test\")\n",
        "inference_on_dataset(trainer.model, val_loader, evaluator)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Predict on test dataset from the model trained. Congratulations!"
      ],
      "metadata": {
        "id": "TNVZW93dnUCB"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_5v-I5kb_6OJ"
      },
      "outputs": [],
      "source": [
        "cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, \"model_final.pth\")\n",
        "cfg.DATASETS.TEST = (\"my_dataset_test1\", )\n",
        "cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5   # set the testing threshold for this model\n",
        "predictor = DefaultPredictor(cfg)\n",
        "test_metadata = MetadataCatalog.get(\"my_dataset_test1\")\n",
        "\n",
        "from detectron2.utils.visualizer import ColorMode\n",
        "import glob\n",
        "\n",
        "for imageName in glob.glob('/content/test/*jpg'):\n",
        "  im = cv2.imread(imageName)\n",
        "  outputs = predictor(im)\n",
        "  print(f'output : {outputs}')\n",
        "  v = Visualizer(im[:, :, ::-1],\n",
        "                metadata=test_metadata,\n",
        "                scale=0.8\n",
        "                 )\n",
        "  out = v.draw_instance_predictions(outputs[\"instances\"].to(\"cpu\"))\n",
        "  cv2_imshow(out.get_image()[:, :, ::-1])"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}