CV_11_detectron2_custom_data.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Copy of AMMI_Thesis_Experiment_1.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "authorship_tag": "ABX9TyNBGqLB9pNgwQ/Uz7SqiIKJ",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Gedeon-m-gedus/Image_Processing/blob/master/CV_11_detectron2_custom_data.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R27O82HBmKt9",
        "colab_type": "text"
      },
      "source": [
        "# Training object detection model on my custom dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GsOKkvVUPW-n",
        "colab_type": "text"
      },
      "source": [
        "## 1. DTECTRON2"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nHtwVY56kYAb",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Downloading data\n",
        "from google.colab import drive\n",
        "drive.mount('/content/gdrive')#\n",
        "#root_path = 'gdrive/My Drive/AMMI_project/object_detector'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tIwnimpHeLSS",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "#Downloading data from my google drive\n",
        "!unzip -q /content/gdrive/My\\ Drive/AMMI_project/object_detector/data.zip -d dataset/\n",
        "#Downloading data from my google drive\n",
        "#!unzip -q /content/gdrive/My\\ Drive/AMMI_project/object_detector/data2.zip -d dataset2/"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZbqSKEGR-4kh",
        "colab_type": "text"
      },
      "source": [
        "### Imports"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "U6S5mCj-w1mb",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!pip install -U torch==1.5 torchvision==0.6 -f https://download.pytorch.org/whl/cu101/torch_stable.html >/dev/null\n",
        "!pip install cython pyyaml==5.1 >/dev/null\n",
        "!pip install -U 'git+https://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI' >/dev/null\n",
        "import torch, torchvision\n",
        "print(torch.__version__, torch.cuda.is_available())\n",
        "!gcc --version\n",
        "\n",
        "# install detectron2:\n",
        "!pip install detectron2==0.1.2 -f https://dl.fbaipublicfiles.com/detectron2/wheels/cu101/index.html >/dev/null\n",
        "\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 os\n",
        "import random\n",
        "import time\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 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, DatasetCatalog\n",
        "from detectron2.engine import DefaultTrainer\n",
        "from detectron2.config import get_cfg\n",
        "from detectron2.utils.visualizer import ColorMode\n",
        "from detectron2.evaluation import COCOEvaluator, inference_on_dataset\n",
        "from detectron2.data import build_detection_test_loader\n",
        "#from detectron2.data.datasets import register_pascal_voc\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ga00CFdUXeq3",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "name = 'mydata'\n",
        "dirname = './dataset/data/'\n",
        "dirname2 = './dataset2/data2/'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CZUPhZSJdXKv",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import numpy as np\n",
        "import os\n",
        "import xml.etree.ElementTree as ET\n",
        "from typing import List, Tuple, Union\n",
        "from fvcore.common.file_io import PathManager\n",
        "\n",
        "from detectron2.data import DatasetCatalog, MetadataCatalog\n",
        "from detectron2.structures import BoxMode"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sLhG55Esc2QJ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "\n",
        "__all__ = [\"load_voc_instances\", \"register_pascal_voc\"]\n",
        "\n",
        "\n",
        "# fmt: off\n",
        "CLASS_NAMES = ('giraffe',\n",
        "                'person',\n",
        "                'zebra',\n",
        "                'elephant',\n",
        "                'impala',\n",
        "                'monkey',\n",
        "                'lion',\n",
        "                'leopard',\n",
        "                'crocodile',\n",
        "                'buffalo',\n",
        "                'hyna',\n",
        "                'bird',\n",
        "                'gorilla')\n",
        "# fmt: on\n",
        "\n",
        "def load_voc_instances(dirname: str, split: str, class_names: Union[List[str], Tuple[str, ...]]):\n",
        "    \"\"\"\n",
        "    Load Pascal VOC detection annotations to Detectron2 format.\n",
        "    Args:\n",
        "        dirname: Contain \"Annotations\", \"ImageSets\", \"JPEGImages\"\n",
        "        split (str): one of \"train\", \"test\", \"val\", \"trainval\"\n",
        "        class_names: list or tuple of class names\n",
        "    \"\"\"\n",
        "    with PathManager.open(os.path.join(dirname, split + \".txt\")) as f:\n",
        "        fileids = np.loadtxt(f, dtype=np.str)\n",
        "\n",
        "    # Needs to read many small annotation files. Makes sense at local\n",
        "    annotation_dirname = PathManager.get_local_path(os.path.join(dirname, \"annotations/\"))\n",
        "    dicts = []\n",
        "    for fileid in fileids:\n",
        "        anno_file = os.path.join(annotation_dirname, fileid + \".xml\")\n",
        "        jpeg_file = os.path.join(dirname, \"images/\", fileid + \".jpg\")\n",
        "\n",
        "        with PathManager.open(anno_file) as f:\n",
        "            tree = ET.parse(f)\n",
        "\n",
        "        r = {\n",
        "            \"file_name\": jpeg_file,\n",
        "            \"image_id\": fileid,\n",
        "            \"height\": int(tree.findall(\"./size/height\")[0].text),\n",
        "            \"width\": int(tree.findall(\"./size/width\")[0].text),\n",
        "        }\n",
        "        instances = []\n",
        "\n",
        "        for obj in tree.findall(\"object\"):\n",
        "            cls = obj.find(\"name\").text\n",
        "            # We include \"difficult\" samples in training.\n",
        "            # Based on limited experiments, they don't hurt accuracy.\n",
        "            # difficult = int(obj.find(\"difficult\").text)\n",
        "            # if difficult == 1:\n",
        "            # continue\n",
        "            bbox = obj.find(\"bndbox\")\n",
        "            bbox = [float(bbox.find(x).text) for x in [\"xmin\", \"ymin\", \"xmax\", \"ymax\"]]\n",
        "            # Original annotations are integers in the range [1, W or H]\n",
        "            # Assuming they mean 1-based pixel indices (inclusive),\n",
        "            # a box with annotation (xmin=1, xmax=W) covers the whole image.\n",
        "            # In coordinate space this is represented by (xmin=0, xmax=W)\n",
        "            bbox[0] -= 1.0\n",
        "            bbox[1] -= 1.0\n",
        "            instances.append(\n",
        "                {\"category_id\": class_names.index(cls), \"bbox\": bbox, \"bbox_mode\": BoxMode.XYXY_ABS}\n",
        "            )\n",
        "        r[\"annotations\"] = instances\n",
        "        dicts.append(r)\n",
        "    return dicts\n",
        "\n",
        "\n",
        "def register_pascal_voc(name, dirname, split, year, class_names=CLASS_NAMES):\n",
        "    meta_data = DatasetCatalog.register(name, lambda: load_voc_instances(dirname, split, class_names))\n",
        "    catalog = MetadataCatalog.get(name).set(\n",
        "        thing_classes=list(class_names), dirname=dirname, year=year, split=split\n",
        "    )\n",
        "    #return meta_data, catalog"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GRDQ-5RueApm",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "#train_meta_data, train_catalog = \n",
        "register_pascal_voc('train_data', dirname, split = 'train', year='2020', class_names=CLASS_NAMES)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "s6vxuxfFBPlH",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "train_dataset_dicts = DatasetCatalog.get('train_data')\n",
        "train_metadata=MetadataCatalog.get('train_data')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "u0fewzc6Q_HM",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "register_pascal_voc('valid_data', dirname, split = 'valid', year='2020', class_names=CLASS_NAMES)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jh8YyhiRQ9EM",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "valid_dataset_dicts = DatasetCatalog.get('valid_data')\n",
        "valid_metadata=MetadataCatalog.get('valid_data')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BWGoU9W9yqqw",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "valid_dataset_dicts[0]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FoJR4X629Ria",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "images_index = [0,7,78,156,444]\n",
        "for index in images_index:\n",
        "    img = cv2.imread(train_dataset_dicts[index][\"file_name\"])\n",
        "    visualizer = Visualizer(img[:, :, ::-1], metadata=train_metadata, scale=0.5)\n",
        "    vis = visualizer.draw_dataset_dict(train_dataset_dicts[index])\n",
        "    print('visualizing ',train_dataset_dicts[index][\"file_name\"])\n",
        "    cv2_imshow(vis.get_image()[:, :, ::-1])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gXg7w-Ou9RmT",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def rcnn_R_50(init_weights_flag = True, path = '/COCO_init_model'):\n",
        "\n",
        "\n",
        "      \"\"\"\n",
        "      This function takes two parameters:\n",
        "        1.  init_weights_flag : This is a flag that defines the initial weights of the backbone network of our R-CNN\n",
        "        2. path : this is the path where I need to save my model\n",
        "      \n",
        "      This function returns the predictor, after training the model it will loads the weights from the directory and the model configurations(cfg)\n",
        "\n",
        "      Parameters are defined in the same way as our lecturer said:\n",
        "        train for 300 iterations, a start learning rate of 0.02, 2 images per batch, and 128 regions per batch\n",
        "      \"\"\"\n",
        "\n",
        "\n",
        "      cfg = get_cfg()\n",
        "      cfg.merge_from_file(model_zoo.get_config_file(\"COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml\"))\n",
        "      cfg.DATASETS.TRAIN = (\"train_data\",)\n",
        "      cfg.DATASETS.TEST = ()\n",
        "      cfg.DATALOADER.NUM_WORKERS = 2\n",
        "      cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(\"COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml\")\n",
        "      cfg.SOLVER.IMS_PER_BATCH = 8 # images per batch\n",
        "      cfg.SOLVER.BASE_LR = 0.02  # Learning rate\n",
        "      cfg.SOLVER.MAX_ITER = 500    # number of iterations \n",
        "      cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128 #regions per batch\n",
        "      cfg.MODEL.ROI_HEADS.NUM_CLASSES = 13  #three classes\n",
        "      cfg.with_coco_init = init_weights_flag # Flag to initialize weight to COCO or IMAGENET\n",
        "\n",
        "      #create a path where I will store my model, this path defers for each model\n",
        "      os.makedirs(path, exist_ok=True)\n",
        "      cfg.OUTPUT_DIR = path\n",
        "      trainer = DefaultTrainer(cfg) \n",
        "      trainer.resume_or_load(resume=False)\n",
        "      trainer.train()\n",
        "\n",
        "      ## After training lets return the predictor\n",
        "      cfg.MODEL.WEIGHTS = os.path.join(path, \"model_final.pth\")\n",
        "      cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7   # set the testing threshold for this model\n",
        "      cfg.DATASETS.TEST = (\"valid_data\", )\n",
        "\n",
        "      return DefaultPredictor(cfg),cfg,trainer"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cJWTP7nPHP7Q",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "coco_init_predictor, coco_init_cfg, coco_init_trainer = rcnn_R_50()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZnHthAe-SGRa",
        "colab_type": "text"
      },
      "source": [
        "### Prediction"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "G8zRWceEbegb",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from google.colab.patches import cv2_imshow\n",
        "for d in valid_dataset_dicts:\n",
        "  img = cv2.imread(d[\"file_name\"])\n",
        "  prediction = coco_init_predictor(img)\n",
        "  visualizer = Visualizer(img[:, :, ::-1],\n",
        "                                metadata=valid_metadata,\n",
        "                                scale=0.5)\n",
        "  vis = visualizer.draw_instance_predictions(prediction[\"instances\"].to(\"cpu\"))\n",
        "  cv2_imshow( vis.get_image()[:, :, ::-1])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HEjoAuxdWivs",
        "colab_type": "text"
      },
      "source": [
        "### Evalution"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "iXolKAfFbeto",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "evaluator = COCOEvaluator(\"valid_data\", coco_init_cfg, False, output_dir=\"./dataset\")\n",
        "val_loader = build_detection_test_loader(coco_init_cfg, \"valid_data\")\n",
        "inference_on_dataset(coco_init_trainer.model, val_loader, evaluator)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_V5c9HLubeaQ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "y6NUsBUZNWbD",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fKFciMp-NWCR",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yyIQGMSyyLDT",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}