VietOCR_Detectron2.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "VietOCR-Detectron2.ipynb",
      "private_outputs": true,
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Yd6NebZ5hP0A"
      },
      "source": [
        "# Install Detectron2 Dependencies"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tI5iqHcyLS7I"
      },
      "source": [
        "# memory footprint support libraries/code\n",
        "!ln -sf /opt/bin/nvidia-smi /usr/bin/nvidia-smi\n",
        "!pip install gputil\n",
        "!pip install psutil\n",
        "!pip install humanize\n",
        "import psutil\n",
        "import humanize\n",
        "import os\n",
        "import GPUtil as GPU\n",
        "GPUs = GPU.getGPUs()\n",
        "# XXX: only one GPU on Colab and isn’t guaranteed\n",
        "gpu = GPUs[0]\n",
        "def printm():\n",
        " process = psutil.Process(os.getpid())\n",
        " print(\"Gen RAM Free: \" + humanize.naturalsize( psutil.virtual_memory().available ), \" | Proc size: \" + humanize.naturalsize( process.memory_info().rss))\n",
        " print(\"GPU RAM Free: {0:.0f}MB | Used: {1:.0f}MB | Util {2:3.0f}% | Total {3:.0f}MB\".format(gpu.memoryFree, gpu.memoryUsed, gpu.memoryUtil*100, gpu.memoryTotal))\n",
        "printm()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wXisIbT1Zqou"
      },
      "source": [
        "# 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": "code",
      "metadata": {
        "id": "wW8A0IHVZ_MR"
      },
      "source": [
        "# install detectron2:\n",
        "!pip install detectron2==0.1.3 -f https://dl.fbaipublicfiles.com/detectron2/wheels/cu101/torch1.5/index.html"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2rCUZZnbhcyl"
      },
      "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 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"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F0evgr7zrxQF"
      },
      "source": [
        "# Mount To My Drive"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uwZSoMun2A9Z"
      },
      "source": [
        "#mount to my drive\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Kd0NJiPJhiu7"
      },
      "source": [
        "# Import and Register Custom Detectron2 Data"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QvbWgc7HhfWi"
      },
      "source": [
        "# !curl -L \"/content/drive/MyDrive/\" > roboflow.zip; unzip roboflow.zip; rm roboflow.zip\n",
        "!unzip /content/drive/MyDrive/infor_card_json.zip > /dev/null\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7Mi9gsZzhokl"
      },
      "source": [
        "from detectron2.data.datasets import register_coco_instances\n",
        "register_coco_instances(\"my_dataset_train\", {}, \"/content/coco_json/train/_annotations.coco.json\", \"/content/coco_json/train\")\n",
        "register_coco_instances(\"my_dataset_val\", {}, \"/content/coco_json/valid/_annotations.coco.json\", \"/content/coco_json/valid\")\n",
        "register_coco_instances(\"my_dataset_test\", {}, \"/content/coco_json/test/_annotations.coco.json\", \"/content/coco_json/test\")\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BdTAusKE9zUQ"
      },
      "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])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "okQbhIYIh_CL"
      },
      "source": [
        "# Train Custom Detectron2 Detector"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "B4UESNQ4tyVm"
      },
      "source": [
        "#We are importing our own Trainer Module here to use the COCO validation evaluation during training. Otherwise no validation eval occurs.\n",
        "\n",
        "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)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NgnVuGA4BG3S"
      },
      "source": [
        "# import torch\n",
        "# torch.cuda.empty_cache()\n",
        "# torch.cuda.memory_summary(device=None, abbreviated=False)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TPc8yVBVh52F"
      },
      "source": [
        "from detectron2.config import get_cfg\n",
        "import os\n",
        "\n",
        "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.OUTPUT_DIR=(\"/content/drive/MyDrive/ID_CARD/Detect_infor/model_coco_X101\")\n",
        "cfg.DATASETS.TRAIN = (\"my_dataset_train\",)\n",
        "cfg.DATASETS.TEST = (\"my_dataset_val\",)\n",
        "\n",
        "cfg.DATALOADER.NUM_WORKERS = 2\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 = 1\n",
        "cfg.SOLVER.BASE_LR = 0.00025\n",
        "\n",
        "\n",
        "# cfg.SOLVER.WARMUP_ITERS = 1000\n",
        "cfg.SOLVER.MAX_ITER = 1500 #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",
        "\n",
        "cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 32\n",
        "cfg.MODEL.ROI_HEADS.NUM_CLASSES = 9 #your number of classes + 1\n",
        "\n",
        "# cfg.TEST.EVAL_PERIOD = 500\n",
        "\n",
        "\n",
        "os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)\n",
        "trainer = CocoTrainer(cfg)\n",
        "# trainer.resume_or_load(resume=False)\n",
        "# trainer.train()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oca9rEQKif1h"
      },
      "source": [
        "# Look at training curves in tensorboard:\n",
        "%load_ext tensorboard\n",
        "%tensorboard --logdir output"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HBUdNVhn1rHh"
      },
      "source": [
        "#evaluate my model\n",
        "\n",
        "#test evaluation\n",
        "from detectron2.data import DatasetCatalog, MetadataCatalog, build_detection_test_loader\n",
        "from detectron2.evaluation import COCOEvaluator, inference_on_dataset\n",
        "from detectron2.config import get_cfg\n",
        "#from detectron2.evaluation.coco_evaluation import COCOEvaluator\n",
        "import os\n",
        "\n",
        "# cfg = get_cfg()\n",
        "# cfg.OUTPUT_DIR=(\"/content/drive/MyDrive/ID_CARD/Detect_infor/model_coco_X101\")\n",
        "cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, \"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)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IiM5jqwLrQmz"
      },
      "source": [
        "# Download VietOCR"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EQGW9_YZro-9"
      },
      "source": [
        "! pip install --quiet vietocr==0.3.2"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-yqkEq6WrpUx"
      },
      "source": [
        "import matplotlib.pyplot as plt\n",
        "from PIL import Image\n",
        "from vietocr.tool.predictor import Predictor\n",
        "from vietocr.tool.config import Cfg"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kYO8R_gXsQNL"
      },
      "source": [
        "config = Cfg.load_config_from_name('vgg_transformer')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YAzk-K8KsQYQ"
      },
      "source": [
        "# config['weights'] = './weights/transformerocr.pth'\n",
        "config['weights'] = 'https://drive.google.com/uc?id=13327Y1tz1ohsm5YZMyXVMPIOjoOA0OaA'\n",
        "config['cnn']['pretrained']=False\n",
        "config['device'] = 'cuda:0'\n",
        "config['predictor']['beamsearch']=False"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9VqrWfGosScE"
      },
      "source": [
        "detector = Predictor(config)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "A3BX34iXw1gU"
      },
      "source": [
        "# Inference with Detectron2 Saved Weights and VietOCR\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zVBjf0DE7HEW"
      },
      "source": [
        "#load model weights\n",
        "cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, \"/content/drive/MyDrive/ID_CARD/Detect_infor/model_coco_X101/model_final.pth\")\n",
        "cfg.DATASETS.TEST = (\"my_dataset_test\", )\n",
        "cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7   # set the testing threshold for this model\n",
        "predictor = DefaultPredictor(cfg)\n",
        "test_metadata = MetadataCatalog.get(\"my_dataset_test\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZrjBxdv4stN4"
      },
      "source": [
        "from detectron2.utils.visualizer import ColorMode\n",
        "import glob\n",
        "import numpy as np\n",
        "import matplotlib\n",
        "import matplotlib.pyplot as plt\n",
        "from google.colab.patches import cv2_imshow\n",
        "import cv2"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fD8GfUcFdao8"
      },
      "source": [
        "from PIL import Image\n",
        "import numpy \n",
        "import string\n",
        "\n",
        "def inference_OCR(img):\n",
        "  cv2_imshow(img)\n",
        "  image=Image.fromarray(numpy.uint8(img))\n",
        "  s=detector.predict(image)\n",
        "  print(s)\n",
        "\n",
        "\n",
        "def cropped_images(imageName):\n",
        "  im = cv2.imread(imageName)\n",
        "  outputs = predictor(im)\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])\n",
        "\n",
        "  #shape of original_image\n",
        "  original_h, original_w = im.shape[:2]\n",
        "\n",
        "  #extract 4-points and categories ID\n",
        "  pred_box=outputs[\"instances\"].get('pred_boxes').tensor\n",
        "  pred_class=outputs[\"instances\"].get('pred_classes')\n",
        "\n",
        "  box=[]\n",
        "  for i in range(len(pred_box)):\n",
        "    temp=[]\n",
        "    for j in range(0,4):\n",
        "      temp.append(int(pred_box[i][j].item()))\n",
        "    box.append(temp)\n",
        "\n",
        "  #extract identify number classes\n",
        "  class_id=[]\n",
        "  for i in range(len(pred_class)):\n",
        "    class_id.append(int(pred_class[i].item()))\n",
        "\n",
        "  #sorted identify number classes\n",
        "  list_id=[]\n",
        "  count_1=0\n",
        "  count_5=0\n",
        "  for index,element in enumerate(class_id):\n",
        "    list_id.append([element,index])\n",
        "    if element==1: count_1+=1\n",
        "    if element==5: count_5+=1\n",
        "  list_id= sorted(list_id[:])\n",
        "\n",
        "  #append sub-image from image id-card\n",
        "  list_cropped=[]\n",
        "  for i in range(len(class_id)):\n",
        "      [xmin, ymin, xmax, ymax]= box[list_id[i][1]]\n",
        "      cropped = im[int(ymin):int(ymax),int(xmin):int(xmax)]\n",
        "      list_cropped.append(cropped)\n",
        "\n",
        "  #show out sub-image results\n",
        "  print(\"______________________________________________________________________________________________\\n\")\n",
        "\n",
        "  if count_1==2: pred_1=1\n",
        "  else: pred_1=0\n",
        "  if count_5==2: pred_5=1\n",
        "  else: pred_5=0\n",
        "\n",
        "  print(\"Identify number: \")\n",
        "  inference_OCR(list_cropped[7+pred_1+pred_5])\n",
        "\n",
        "  print('-----------------------------------------------')\n",
        "  print(\"Name: \")\n",
        "  inference_OCR(list_cropped[5+pred_1+pred_5])\n",
        "\n",
        "  print('-----------------------------------------------')\n",
        "  print(\"Date of Birth: \")\n",
        "  inference_OCR(list_cropped[1+pred_1])\n",
        "\n",
        "  print('-----------------------------------------------')\n",
        "  print(\"Gender: \")\n",
        "  inference_OCR(list_cropped[3+pred_1])\n",
        "\n",
        "  print('-----------------------------------------------')\n",
        "  print(\"Nationality: \")\n",
        "  inference_OCR(list_cropped[6+pred_1+pred_5])\n",
        "\n",
        "  print('-----------------------------------------------')\n",
        "  print(\"Native land: \")\n",
        "  if count_1==2:\n",
        "    # im=cv2.hconcat([list_cropped[0+pred_1],list_cropped[0]])\n",
        "    inference_OCR(list_cropped[0+pred_1])\n",
        "    inference_OCR(list_cropped[0])\n",
        " \n",
        "  else:\n",
        "    inference_OCR(list_cropped[0])\n",
        "\n",
        "  print('-----------------------------------------------')\n",
        "  print(\"Permanent address: \")\n",
        "  if count_5==2:\n",
        "    # im=cv2.hconcat([list_cropped[4+pred_1+pred_5]],list_cropped[4+pred_1])\n",
        "    inference_OCR(list_cropped[4+pred_1+pred_5])\n",
        "    inference_OCR(list_cropped[4+pred_1])\n",
        "  else:\n",
        "    inference_OCR(list_cropped[4+pred_1])\n",
        " \n",
        "  print('-----------------------------------------------')\n",
        "  print(\"Expiry date: \")\n",
        "  inference_OCR(list_cropped[2+pred_1])\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8a0JcYbM7MfM"
      },
      "source": [
        "!pip install opencv-contrib-python"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WiJ0Ylc_XAUa"
      },
      "source": [
        "from detectron2.utils.visualizer import ColorMode\n",
        "import glob\n",
        "import numpy as np\n",
        "import matplotlib\n",
        "import matplotlib.pyplot as plt\n",
        "from google.colab.patches import cv2_imshow\n",
        "import cv2\n",
        "import os\n",
        "\n",
        "for imageName in glob.glob('/content/drive/MyDrive/ID_CARD/images_infor/*jpg'):\n",
        "  cropped_images(imageName)"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}