Colab-Detectron2.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Colab-Detectron2.ipynb",
      "private_outputs": true,
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QHnVupBBn9eR"
      },
      "source": [
        "# Detectron2 (Original+ResNeSt+centermask2)\n",
        "\n",
        "<img src=\"https://dl.fbaipublicfiles.com/detectron2/Detectron2-Logo-Horz.png\" width=\"500\">\n",
        "\n",
        "This is a modification of the offiical detectron2 colab that can be found [here](https://colab.research.google.com/drive/16jcaJoc6bCFAQ96jDe2HwtXj7BMD_-m5 ).\n",
        "\n",
        "This colab is a combination of [facebookresearch/detectron2](https://github.com/facebookresearch/detectron2), [zhanghang1989/detectron2-ResNeSt](https://github.com/zhanghang1989/detectron2-ResNeSt) and [youngwanLEE/centermask2](https://github.com/youngwanLEE/centermask2).\n",
        "\n",
        "My fork is located in [styler00dollar/Colab-Detectron2](https://github.com/styler00dollar/Colab-Detectron2).\n",
        "\n",
        "You might need to edit some paths or parameter to fit your needs. This assumes you know what you are doing.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_MZbDQNfx2DB",
        "cellView": "form"
      },
      "source": [
        "#@title check gpu\n",
        "!nvidia-smi"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zljLQ5Ij1VQo",
        "cellView": "form"
      },
      "source": [
        "#@title check amount of files in folder (specify path)\n",
        "%cd \"/path/\"\n",
        "!ls | wc -l"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2s7SaYMYYQHF",
        "cellView": "form"
      },
      "source": [
        "#@title mount Google Drive\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vM54r6jlKTII"
      },
      "source": [
        "# Install detectron2"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9_FzH13EjseR",
        "cellView": "form"
      },
      "source": [
        "#@title Install dependencies\n",
        "# install dependencies: (use cu101 because colab has CUDA 10.1)\n",
        "#!pip install -U torch==1.5 torchvision==0.6 -f https://download.pytorch.org/whl/cu101/torch_stable.html \n",
        "!pip install cython pyyaml==5.1\n",
        "!pip install -U 'git+https://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI'\n",
        "import torch, torchvision\n",
        "print(torch.__version__, torch.cuda.is_available())\n",
        "!gcc --version\n",
        "# opencv is pre-installed on colab"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zuPhNcoNc6bG"
      },
      "source": [
        "# Install detectron2 or detectron2-ResNeSt"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "b-i4hmGYk1dL",
        "cellView": "form"
      },
      "source": [
        "#@title Install detectron2\n",
        "%cd /content/\n",
        "# install detectron2:\n",
        "!git clone https://github.com/facebookresearch/detectron2\n",
        "%cd /content/detectron2\n",
        "!pip install -r requirements.txt\n",
        "!python setup.py install\n",
        "!pip install git+https://github.com/facebookresearch/fvcore.git\n",
        "\"\"\"\n",
        "%cd /content/\n",
        "!git clone https://github.com/facebookresearch/iopath\n",
        "%cd /content/iopath\n",
        "!python setup.py install\n",
        "\n",
        "%cd /content/\n",
        "!git clone https://github.com/WoLpH/portalocker\n",
        "%cd /content/portalocker\n",
        "!python setup.py install\n",
        "\"\"\""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ImwNQz6BvWsk",
        "cellView": "form"
      },
      "source": [
        "#@title Install detectron2-ResNeSt\n",
        "%cd /content/\n",
        "# install detectron resnest\n",
        "!python -m pip install 'git+https://github.com/zhanghang1989/detectron2-ResNeSt.git'\n",
        "import detectron2\n",
        "\n",
        "# clone detectron2 resnest\n",
        "%cd /content/\n",
        "!git clone https://github.com/zhanghang1989/detectron2-ResNeSt"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vE9W1BCvZuKw",
        "cellView": "form"
      },
      "source": [
        "#@title Install detectron2 / clone centermask2\n",
        "%cd /content/\n",
        "# install detectron2:\n",
        "!pip install detectron2==0.3 -f https://dl.fbaipublicfiles.com/detectron2/wheels/cu101/torch1.5/index.html\n",
        "# git clone detectron2\n",
        "!git clone https://github.com/youngwanLEE/centermask2"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZyAvNCJMmvFF",
        "cellView": "form"
      },
      "source": [
        "#@title Import detectron2\n",
        "\n",
        "#@markdown 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 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"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b2bjrfb2LDeo"
      },
      "source": [
        "# Train on a custom coco dataset\n",
        "\n",
        "Copy a .7z file and extract its contents.\n",
        "\n",
        "The datastructure should be like this:\n",
        " - some_name.7z\n",
        "  - train\n",
        "    - image1.png\n",
        "    - lksjfdl.png\n",
        "    - ...\n",
        "    - train.json\n",
        "  - test\n",
        "    - test1.png\n",
        "    - lksdjfl.png\n",
        "    - ...\n",
        "    - test.json"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UIpcAq0PYUVT",
        "cellView": "form"
      },
      "source": [
        "#@title copy to local storage and extract\n",
        "!cp -r '/content/drive/My Drive/dataset.7z' /content/dataset.7z\n",
        "%cd /content/\n",
        "!7z x /content/dataset.7z -oc:/content/"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wlqXIXXhW8dA"
      },
      "source": [
        "## Train config\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mxSE5pLoiBgQ"
      },
      "source": [
        "\n",
        "[TannerGilbert/Tensorflow-Object-Detection-API-train-custom-Mask-R-CNN-model](https://github.com/TannerGilbert/Tensorflow-Object-Detection-API-train-custom-Mask-R-CNN-model) contains a small coco dataset, which is used as a test."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZfPwp1_zhdqX",
        "cellView": "form"
      },
      "source": [
        "#@title Downloading demo dataset\n",
        "%cd /content/\n",
        "!git clone https://github.com/TannerGilbert/Tensorflow-Object-Detection-API-train-custom-Mask-R-CNN-model object_detection"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GDJ2NHbnIPQB"
      },
      "source": [
        "!7z x dataset.7z"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PIbAM2pv-urF",
        "cellView": "form"
      },
      "source": [
        "#@title Register dataset\n",
        "from detectron2.data.datasets import register_coco_instances\n",
        "register_coco_instances(\"my_dataset_train2\", {}, \"/content/dataset/coco.json\", \"/content/dataset/output/\")\n",
        "register_coco_instances(\"my_dataset_val2\", {}, \"/content/dataset/coco.json\", \"/content/dataset/output/\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hBXeH8UXFcqU",
        "cellView": "form"
      },
      "source": [
        "#@title Activate tensorboard\n",
        "# Look at training curves in tensorboard:\n",
        "%load_ext tensorboard\n",
        "%tensorboard --logdir \"/content/drive/My Drive/detectron2\""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Hn3Y61EZocl2"
      },
      "source": [
        "# detectron2"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7unkuuiqLdqd",
        "cellView": "form"
      },
      "source": [
        "#@title (detectron2) mask_rcnn_X_101_32x8d_FPN_3x.yaml\n",
        "import os\n",
        "import numpy as np\n",
        "import json\n",
        "from detectron2.structures import BoxMode\n",
        "from detectron2.engine import DefaultTrainer\n",
        "from detectron2.config import get_cfg\n",
        "\n",
        "cfg = get_cfg()\n",
        "\n",
        "os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)\n",
        "\n",
        "cfg.merge_from_file(\"/content/detectron2/configs/COCO-InstanceSegmentation/mask_rcnn_X_101_32x8d_FPN_3x.yaml\")\n",
        "\n",
        "cfg.DATASETS.TRAIN = (\"my_dataset_train2\",)\n",
        "cfg.DATASETS.TEST = (\"my_dataset_val2\",)\n",
        "cfg.DATALOADER.NUM_WORKERS = 2\n",
        "\n",
        "cfg.MODEL.WEIGHTS = (\"/content/detectron2/configs/COCO-InstanceSegmentation/mask_rcnn_X_101_32x8d_FPN_3x.yaml\")\n",
        "\n",
        "cfg.INPUT.MAX_SIZE_TEST = 1700\n",
        "cfg.INPUT.MAX_SIZE_TRAIN = 1700\n",
        "cfg.SOLVER.IMS_PER_BATCH = 2   # increase this for more speed, but it will need more vram\n",
        "cfg.SOLVER.BASE_LR = 0.00025  # pick a good LR\n",
        "#cfg.SOLVER.BASE_LR = 0.000001  # pick a good LR /content/detectron2/configs/COCO-InstanceSegmentation/mask_rcnn_R_101_FPN_3x.yaml\n",
        "cfg.SOLVER.MAX_ITER = 500000    # 300 iterations seems good enough for this toy dataset; you may need to train longer for a practical dataset\n",
        "cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 4096   # or 128 for example\n",
        "cfg.MODEL.ROI_HEADS.NUM_CLASSES = 2  # only has one class (ballon)\n",
        "#cfg.SNAPSHOT_ITERS = 50 # when checkpoint gets created\n",
        "cfg.SOLVER.CHECKPOINT_PERIOD = 5000\n",
        "\n",
        "cfg.OUTPUT_DIR = \"/content/drive/MyDrive/detectron2\"\n",
        "trainer = DefaultTrainer(cfg) \n",
        "#trainer.resume_or_load(resume=True)\n",
        "trainer.train()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TmKSHXdKH1P6",
        "cellView": "form"
      },
      "source": [
        "#@title (detectron2) faster_rcnn_R_50_FPN_3x.yaml\n",
        "import os\n",
        "import numpy as np\n",
        "import json\n",
        "from detectron2.structures import BoxMode\n",
        "from detectron2.engine import DefaultTrainer\n",
        "from detectron2.config import get_cfg\n",
        "\n",
        "cfg = get_cfg()\n",
        "\n",
        "os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)\n",
        "\n",
        "cfg.merge_from_file(\"/content/detectron2/configs/COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml\")\n",
        "\n",
        "cfg.DATASETS.TRAIN = (\"my_dataset_train2\",)\n",
        "cfg.DATASETS.TEST = (\"my_dataset_val2\",)\n",
        "cfg.DATALOADER.NUM_WORKERS = 2\n",
        "\n",
        "cfg.MODEL.WEIGHTS = (\"/content/detectron2/configs/COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml\")\n",
        "\n",
        "cfg.INPUT.MAX_SIZE_TEST = 1300\n",
        "cfg.INPUT.MAX_SIZE_TRAIN = 1300\n",
        "cfg.SOLVER.IMS_PER_BATCH = 2   # increase this for more speed, but it will need more vram\n",
        "cfg.SOLVER.BASE_LR = 0.00025  # pick a good LR\n",
        "#cfg.SOLVER.BASE_LR = 0.000001  # pick a good LR /content/detectron2/configs/COCO-InstanceSegmentation/mask_rcnn_R_101_FPN_3x.yaml\n",
        "cfg.SOLVER.MAX_ITER = 5000    # 300 iterations seems good enough for this toy dataset; you may need to train longer for a practical dataset\n",
        "cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128   # or 128 for example / 4096?\n",
        "cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1  # only has one class (ballon)\n",
        "#cfg.SNAPSHOT_ITERS = 50 # when checkpoint gets created\n",
        "cfg.SOLVER.CHECKPOINT_PERIOD = 5000\n",
        "\n",
        "cfg.OUTPUT_DIR = \"/content/drive/MyDrive/detectron2\"\n",
        "trainer = DefaultTrainer(cfg) \n",
        "#trainer.resume_or_load(resume=True)\n",
        "trainer.train()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "miqwBt0BogBc"
      },
      "source": [
        "# detectron2-ResNeSt"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Es9qFt5LJnC3",
        "cellView": "form"
      },
      "source": [
        "#@title (detectron2-ResNeSt) mask_cascade_rcnn_R_50_FPN_syncbn_1x.yaml\n",
        "\n",
        "import os\n",
        "import numpy as np\n",
        "import json\n",
        "from detectron2.structures import BoxMode\n",
        "from detectron2.engine import DefaultTrainer\n",
        "from detectron2.config import get_cfg\n",
        "\n",
        "cfg = get_cfg()\n",
        "cfg.OUTPUT_DIR = '/content/drive/MyDrive/detectron2'\n",
        "os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)\n",
        "\n",
        "cfg.merge_from_file(\"/content/detectron2-ResNeSt/configs/COCO-InstanceSegmentation/mask_cascade_rcnn_R_50_FPN_syncbn_1x.yaml\")\n",
        "\n",
        "cfg.DATASETS.TRAIN = (\"my_dataset_train2\",)\n",
        "cfg.DATASETS.TEST = (\"my_dataset_val2\",)\n",
        "cfg.DATALOADER.NUM_WORKERS = 2\n",
        "\n",
        "cfg.MODEL.WEIGHTS = (\"/content/detectron2-ResNeSt/configs/COCO-InstanceSegmentation/mask_cascade_rcnn_R_50_FPN_syncbn_1x.yaml\")\n",
        "\n",
        "#1700x1700 ims_per_batch 2 max_mem: 12522M\n",
        "#1800x1800 ims_per_batch 2 max_mem: 13303M\n",
        "#3000x3000 ims_per_batch 1 max_mem: 7989M\n",
        "cfg.INPUT.MAX_SIZE_TEST = 3000\n",
        "cfg.INPUT.MAX_SIZE_TRAIN = 3000\n",
        "cfg.SOLVER.IMS_PER_BATCH = 1   # increase this for more speed, but it will need more vram\n",
        "cfg.SOLVER.BASE_LR = 0.00025  # pick a good LR\n",
        "\n",
        "cfg.SOLVER.MAX_ITER = 500000    # 300 iterations seems good enough for this toy dataset; you may need to train longer for a practical dataset\n",
        "cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 4096   # or 128 for example\n",
        "cfg.MODEL.ROI_HEADS.NUM_CLASSES = 2  # only has one class (ballon)\n",
        "#cfg.SNAPSHOT_ITERS = 50 # when checkpoint gets created\n",
        "cfg.SOLVER.CHECKPOINT_PERIOD = 5000\n",
        "\n",
        "trainer = DefaultTrainer(cfg) \n",
        "#trainer.resume_or_load(resume=True)\n",
        "trainer.train()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hqkpyPJUssUL",
        "cellView": "form"
      },
      "source": [
        "#@title (detectron2-ResNeSt) mask_cascade_rcnn_R_101_FPN_syncbn_1x.yaml\n",
        "\n",
        "import os\n",
        "import numpy as np\n",
        "import json\n",
        "from detectron2.structures import BoxMode\n",
        "from detectron2.engine import DefaultTrainer\n",
        "from detectron2.config import get_cfg\n",
        "\n",
        "cfg = get_cfg()\n",
        "cfg.OUTPUT_DIR = '/content/drive/MyDrive/detectron2'\n",
        "os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)\n",
        "\n",
        "cfg.merge_from_file(\"/content/detectron2-ResNeSt/configs/COCO-InstanceSegmentation/mask_cascade_rcnn_R_101_FPN_syncbn_1x.yaml\")\n",
        "\n",
        "cfg.DATASETS.TRAIN = (\"my_dataset_train2\",)\n",
        "cfg.DATASETS.TEST = (\"my_dataset_val2\",)\n",
        "cfg.DATALOADER.NUM_WORKERS = 2\n",
        "\n",
        "cfg.MODEL.WEIGHTS = (\"/content/detectron2-ResNeSt/configs/COCO-InstanceSegmentation/mask_cascade_rcnn_R_101_FPN_syncbn_1x.yaml\")\n",
        "\n",
        "#1900x1900 ims_per_batch = 1 max_mem: 7539M\n",
        "#2500x2500 ims_per_batch = 1 max_mem: 9193M\n",
        "#3000x3000 ims_per_batch = 1 max_mem: 9236M\n",
        "cfg.INPUT.MAX_SIZE_TEST = 3000\n",
        "cfg.INPUT.MAX_SIZE_TRAIN = 3000\n",
        "cfg.SOLVER.IMS_PER_BATCH = 1   # increase this for more speed, but it will need more vram\n",
        "cfg.SOLVER.BASE_LR = 0.00025  # pick a good LR\n",
        "\n",
        "cfg.SOLVER.MAX_ITER = 500000    # 300 iterations seems good enough for this toy dataset; you may need to train longer for a practical dataset\n",
        "cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 4096   # or 128 for example\n",
        "cfg.MODEL.ROI_HEADS.NUM_CLASSES = 2  # only has one class (ballon)\n",
        "#cfg.SNAPSHOT_ITERS = 50 # when checkpoint gets created\n",
        "cfg.SOLVER.CHECKPOINT_PERIOD = 5000\n",
        "\n",
        "trainer = DefaultTrainer(cfg) \n",
        "#trainer.resume_or_load(resume=True)\n",
        "trainer.train()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uHjME3iWoi9M"
      },
      "source": [
        "# centermask2"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "ZiiMRRFgoO8s"
      },
      "source": [
        "#@title centermask2/train_net.py (registering dataset)\n",
        "%%writefile /content/centermask2/train_net.py\n",
        "# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved\n",
        "# Modified by Youngwan Lee (ETRI), 2020. All Rights Reserved.\n",
        "from detectron2.data.datasets import register_coco_instances\n",
        "register_coco_instances(\"my_dataset_train2\", {}, \"/content/object_detection/images/train.json\", \"/content/object_detection/images/train/\")\n",
        "register_coco_instances(\"my_dataset_val2\", {}, \"/content/object_detection/images/test.json\", \"/content/object_detection/images/test/\")\n",
        "\n",
        "\n",
        "\n",
        "import logging\n",
        "import os\n",
        "from collections import OrderedDict\n",
        "import torch\n",
        "\n",
        "import detectron2.utils.comm as comm\n",
        "from detectron2.data import MetadataCatalog\n",
        "from detectron2.engine import DefaultTrainer, default_argument_parser, default_setup, hooks, launch\n",
        "from detectron2.evaluation import (\n",
        "    # CityscapesInstanceEvaluator,\n",
        "    # CityscapesSemSegEvaluator,\n",
        "    # COCOEvaluator,\n",
        "    COCOPanopticEvaluator,\n",
        "    DatasetEvaluators,\n",
        "    LVISEvaluator,\n",
        "    PascalVOCDetectionEvaluator,\n",
        "    SemSegEvaluator,\n",
        "    verify_results,\n",
        ")\n",
        "from centermask.evaluation import (\n",
        "    COCOEvaluator,\n",
        "    CityscapesInstanceEvaluator,\n",
        "    CityscapesSemSegEvaluator\n",
        ")\n",
        "from detectron2.modeling import GeneralizedRCNNWithTTA\n",
        "from detectron2.checkpoint import DetectionCheckpointer\n",
        "from centermask.config import get_cfg\n",
        "\n",
        "\n",
        "class Trainer(DefaultTrainer):\n",
        "    \"\"\"\n",
        "    This is the same Trainer except that we rewrite the\n",
        "    `build_train_loader` method.\n",
        "    \"\"\"\n",
        "\n",
        "\n",
        "\n",
        "    @classmethod\n",
        "    def build_evaluator(cls, cfg, dataset_name, output_folder=None):\n",
        "        \"\"\"\n",
        "        Create evaluator(s) for a given dataset.\n",
        "        This uses the special metadata \"evaluator_type\" associated with each builtin dataset.\n",
        "        For your own dataset, you can simply create an evaluator manually in your\n",
        "        script and do not have to worry about the hacky if-else logic here.\n",
        "        \"\"\"\n",
        "        if output_folder is None:\n",
        "            output_folder = os.path.join(cfg.OUTPUT_DIR, \"inference\")\n",
        "        evaluator_list = []\n",
        "        evaluator_type = MetadataCatalog.get(dataset_name).evaluator_type\n",
        "        if evaluator_type in [\"sem_seg\", \"coco_panoptic_seg\"]:\n",
        "            evaluator_list.append(\n",
        "                SemSegEvaluator(\n",
        "                    dataset_name,\n",
        "                    distributed=True,\n",
        "                    output_dir=output_folder,\n",
        "                )\n",
        "            )\n",
        "        if evaluator_type in [\"coco\", \"coco_panoptic_seg\"]:\n",
        "            evaluator_list.append(COCOEvaluator(dataset_name, output_dir=output_folder))\n",
        "        if evaluator_type == \"coco_panoptic_seg\":\n",
        "            evaluator_list.append(COCOPanopticEvaluator(dataset_name, output_folder))\n",
        "        if evaluator_type == \"cityscapes_instance\":\n",
        "            assert (\n",
        "                torch.cuda.device_count() >= comm.get_rank()\n",
        "            ), \"CityscapesEvaluator currently do not work with multiple machines.\"\n",
        "            return CityscapesInstanceEvaluator(dataset_name)\n",
        "        if evaluator_type == \"cityscapes_sem_seg\":\n",
        "            assert (\n",
        "                torch.cuda.device_count() >= comm.get_rank()\n",
        "            ), \"CityscapesEvaluator currently do not work with multiple machines.\"\n",
        "            return CityscapesSemSegEvaluator(dataset_name)\n",
        "        elif evaluator_type == \"pascal_voc\":\n",
        "            return PascalVOCDetectionEvaluator(dataset_name)\n",
        "        elif evaluator_type == \"lvis\":\n",
        "            return LVISEvaluator(dataset_name, output_dir=output_folder)\n",
        "        if len(evaluator_list) == 0:\n",
        "            raise NotImplementedError(\n",
        "                \"no Evaluator for the dataset {} with the type {}\".format(\n",
        "                    dataset_name, evaluator_type\n",
        "                )\n",
        "            )\n",
        "        elif len(evaluator_list) == 1:\n",
        "            return evaluator_list[0]\n",
        "        return DatasetEvaluators(evaluator_list)\n",
        "\n",
        "    @classmethod\n",
        "    def test_with_TTA(cls, cfg, model):\n",
        "        logger = logging.getLogger(\"detectron2.trainer\")\n",
        "        # In the end of training, run an evaluation with TTA\n",
        "        # Only support some R-CNN models.\n",
        "        logger.info(\"Running inference with test-time augmentation ...\")\n",
        "        model = GeneralizedRCNNWithTTA(cfg, model)\n",
        "        evaluators = [\n",
        "            cls.build_evaluator(\n",
        "                cfg, name, output_folder=os.path.join(cfg.OUTPUT_DIR, \"inference_TTA\")\n",
        "            )\n",
        "            for name in cfg.DATASETS.TEST\n",
        "        ]\n",
        "        res = cls.test(cfg, model, evaluators)\n",
        "        res = OrderedDict({k + \"_TTA\": v for k, v in res.items()})\n",
        "        return res\n",
        "\n",
        "\n",
        "\n",
        "def setup(args):\n",
        "    \"\"\"\n",
        "    Create configs and perform basic setups.\n",
        "    \"\"\"\n",
        "    cfg = get_cfg()\n",
        "    cfg.merge_from_file(args.config_file)\n",
        "    cfg.merge_from_list(args.opts)\n",
        "    cfg.freeze()\n",
        "    default_setup(cfg, args)\n",
        "    return cfg\n",
        "\n",
        "\n",
        "def main(args):\n",
        "    cfg = setup(args)\n",
        "\n",
        "    if args.eval_only:\n",
        "        model = Trainer.build_model(cfg)\n",
        "        DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load(\n",
        "            cfg.MODEL.WEIGHTS, resume=args.resume\n",
        "        )\n",
        "        res = Trainer.test(cfg, model)\n",
        "        if cfg.TEST.AUG.ENABLED:\n",
        "            res.update(Trainer.test_with_TTA(cfg, model))\n",
        "        if comm.is_main_process():\n",
        "            verify_results(cfg, res)\n",
        "        return res\n",
        "\n",
        "    \"\"\"\n",
        "    If you'd like to do anything fancier than the standard training logic,\n",
        "    consider writing your own training loop or subclassing the trainer.\n",
        "    \"\"\"\n",
        "    trainer = Trainer(cfg)\n",
        "    trainer.resume_or_load(resume=args.resume)\n",
        "    if cfg.TEST.AUG.ENABLED:\n",
        "        trainer.register_hooks(\n",
        "            [hooks.EvalHook(0, lambda: trainer.test_with_TTA(cfg, trainer.model))]\n",
        "        )\n",
        "    return trainer.train()\n",
        "\n",
        "\n",
        "if __name__ == \"__main__\":\n",
        "    args = default_argument_parser().parse_args()\n",
        "    print(\"Command Line Args:\", args)\n",
        "    launch(\n",
        "        main,\n",
        "        args.num_gpus,\n",
        "        num_machines=args.num_machines,\n",
        "        machine_rank=args.machine_rank,\n",
        "        dist_url=args.dist_url,\n",
        "        args=(args,),\n",
        "    )\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zr3HExSTZ73t",
        "cellView": "form"
      },
      "source": [
        "#@title (centermask2) Base-CenterMask-ResNet.yaml\n",
        "%%writefile /content/centermask2/configs/centermask/Base-CenterMask-ResNet.yaml\n",
        "MODEL:\n",
        "  META_ARCHITECTURE: \"GeneralizedRCNN\"\n",
        "  BACKBONE:\n",
        "    NAME: \"build_fcos_resnet_fpn_backbone\"\n",
        "  RESNETS:\n",
        "    OUT_FEATURES: [\"res3\", \"res4\", \"res5\"]\n",
        "  FPN:\n",
        "    IN_FEATURES: [\"res3\", \"res4\", \"res5\"]\n",
        "  PROPOSAL_GENERATOR:\n",
        "    NAME: \"FCOS\"\n",
        "  FCOS:\n",
        "    POST_NMS_TOPK_TEST: 50\n",
        "  # PIXEL_MEAN: [102.9801, 115.9465, 122.7717]\n",
        "  MASK_ON: True\n",
        "  MASKIOU_ON: True\n",
        "  ROI_HEADS:\n",
        "    NAME: \"CenterROIHeads\"\n",
        "    IN_FEATURES: [\"p3\", \"p4\", \"p5\"]\n",
        "  ROI_MASK_HEAD:\n",
        "    NAME: \"SpatialAttentionMaskHead\"\n",
        "    ASSIGN_CRITERION: \"ratio\"\n",
        "    NUM_CONV: 4\n",
        "    POOLER_RESOLUTION: 14\n",
        "DATASETS:\n",
        "  TRAIN: (\"my_dataset_train2\",)\n",
        "  TEST: (\"my_dataset_val2\",)\n",
        "SOLVER:\n",
        "  IMS_PER_BATCH: 8\n",
        "  BASE_LR: 0.01  # Note that RetinaNet uses a different default learning rate\n",
        "  STEPS: (60000, 80000)\n",
        "  MAX_ITER: 90000\n",
        "INPUT:\n",
        "  MIN_SIZE_TRAIN: (640, 672, 704, 736, 768, 800)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6pj_4sxBl-Xl",
        "cellView": "form"
      },
      "source": [
        "#@title (centermask2) Training (Base-CenterMask-ResNet.yaml)\n",
        "%cd /content/centermask2/\n",
        "!python train_net.py --num-gpus 1 --config /content/centermask2/configs/centermask/centermask_R_50_FPN_ms_3x.yaml"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0e4vdDIOXyxF"
      },
      "source": [
        "## Run tests with model\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ya5nEuMELeq8",
        "cellView": "form"
      },
      "source": [
        "#@title load model\n",
        "import iopath\n",
        "from detectron2.engine import DefaultPredictor\n",
        "cfg = get_cfg()\n",
        "cfg.MODEL.WEIGHTS = '/content/drive/MyDrive/detectron2-test/model_0039999.pth'\n",
        "cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.00  # set the testing threshold for this model\n",
        "#cfg.DATASETS.TEST = (\"my_dataset_val2\", )\n",
        "predictor = DefaultPredictor(cfg)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "U5LhISJqWXgM",
        "cellView": "form"
      },
      "source": [
        "#@title do detection and display result\n",
        "from detectron2.utils.visualizer import ColorMode\n",
        "from detectron2.data import MetadataCatalog\n",
        "from detectron2.utils.visualizer import Visualizer\n",
        "import cv2\n",
        "from google.colab.patches import cv2_imshow\n",
        "\n",
        "#balloon_metadata = MetadataCatalog.get(\"my_dataset_train2\")\n",
        "\n",
        "im = cv2.imread(\"/content/data/test.png\")\n",
        "outputs = predictor(im)\n",
        "print(outputs)\n",
        "v = Visualizer(im[:, :, ::-1],\n",
        "                metadata=None, \n",
        "                scale=0.75, \n",
        "                instance_mode=ColorMode.IMAGE_BW   # remove the colors of unsegmented pixels. This option is only available for segmentation models\n",
        ")\n",
        "v = v.draw_instance_predictions(outputs[\"instances\"].to(\"cpu\"))\n",
        "cv2_imshow(v.get_image()[:, :, ::-1])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "w1N4fITDdqnh"
      },
      "source": [
        "#@title print results for a folder\n",
        "import glob\n",
        "\n",
        "rootdir = \"/content/data/\"\n",
        "files = glob.glob(rootdir + '/**/*.png', recursive=True)\n",
        "files_jpg = glob.glob(rootdir + '/**/*.jpg', recursive=True)\n",
        "files.extend(files_jpg)\n",
        "\n",
        "\n",
        "for f in files:\n",
        "  print(f)\n",
        "  im = cv2.imread(f)\n",
        "  outputs = predictor(im)\n",
        "  print(outputs)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "h9tECBQCvMv3",
        "cellView": "form"
      },
      "source": [
        "#@title AP metric\n",
        "from detectron2.evaluation import COCOEvaluator, inference_on_dataset\n",
        "from detectron2.data import build_detection_test_loader\n",
        "evaluator = COCOEvaluator(\"my_dataset_val2\", cfg, False, output_dir=\"./output/\")\n",
        "val_loader = build_detection_test_loader(cfg, \"my_dataset_val2\")\n",
        "inference_on_dataset(trainer.model, val_loader, evaluator)\n",
        "# another equivalent way is to use trainer.test"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d6n37oWDqhL7"
      },
      "source": [
        "# Misc"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xwJUcmnVAzhK",
        "cellView": "form"
      },
      "source": [
        "#@title restart kernel if something goes wrong\n",
        "import os\n",
        "os._exit(00)"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}