model_test.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# How to link Colab and google Drive\n",
        "\n",
        "After executing this cell, your Colab will be linked to your Google Drive account.\n",
        "\n",
        "Create a folder on Google Drive and upload your test data. In my scenario, /datasets/test\n",
        "\n",
        "Also, upload the model to the same path as the test folder."
      ],
      "metadata": {
        "id": "P4aWg-DJEKWv"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/gdrive')\n",
        "\n",
        "!cp -r /content/gdrive/MyDrive/datasets/custom_model_lite /content\n",
        "!cp -r /content/gdrive/MyDrive/datasets/test /content"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "YGkGfDmKsw6v",
        "outputId": "377b6f58-3aeb-4825-d15f-4ea1e76d4ed4"
      },
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Mounted at /content/gdrive\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!sudo apt install tesseract-ocr\n",
        "!pip install pytesseract"
      ],
      "metadata": {
        "id": "Dqe7R5c1ofKZ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import pytesseract\n",
        "import cv2 as cv\n",
        "\n",
        "\n",
        "\n",
        "def extract_text(xmin, xmax, ymin, ymax, image_path):\n",
        "  img = cv.imread(image_path)\n",
        "\n",
        "  blank = np.zeros(img.shape[:2], dtype='uint8')\n",
        "\n",
        "  blank_rect = cv.rectangle(blank, (xmin, ymin), (xmax, ymax), (255, 0, 255), -1)\n",
        "\n",
        "  detected_img = cv.bitwise_and(img, img, mask=blank_rect)\n",
        "\n",
        "  text = pytesseract.image_to_string(detected_img)\n",
        "  return text"
      ],
      "metadata": {
        "id": "f0kKWg_eoCg8"
      },
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "import cv2\n",
        "import numpy as np\n",
        "import sys\n",
        "import glob\n",
        "import random\n",
        "import importlib.util\n",
        "from tensorflow.lite.python.interpreter import Interpreter\n",
        "import matplotlib\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "%matplotlib inline\n",
        "\n",
        "### Define function for inferencing with TFLite model and displaying results\n",
        "\n",
        "def tflite_detect_images(modelpath, imgpath, lblpath, min_conf=0.5, num_test_images=10, savepath='/content/results', txt_only=False):\n",
        "\n",
        "  # Grab filenames of all images in test folder\n",
        "  images = glob.glob(imgpath + '/*.jpg') + glob.glob(imgpath + '/*.JPG') + glob.glob(imgpath + '/*.png') + glob.glob(imgpath + '/*.bmp')\n",
        "\n",
        "  # Load the label map into memory\n",
        "  with open(lblpath, 'r') as f:\n",
        "      labels = [line.strip() for line in f.readlines()]\n",
        "\n",
        "  # Load the Tensorflow Lite model into memory\n",
        "  interpreter = Interpreter(model_path=modelpath)\n",
        "  interpreter.allocate_tensors()\n",
        "\n",
        "  # Get model details\n",
        "  input_details = interpreter.get_input_details()\n",
        "  output_details = interpreter.get_output_details()\n",
        "  height = input_details[0]['shape'][1]\n",
        "  width = input_details[0]['shape'][2]\n",
        "\n",
        "  float_input = (input_details[0]['dtype'] == np.float32)\n",
        "\n",
        "  input_mean = 127.5\n",
        "  input_std = 127.5\n",
        "\n",
        "  # Randomly select test images\n",
        "  images_to_test = random.sample(images, num_test_images)\n",
        "\n",
        "  # Loop over every image and perform detection\n",
        "  for image_path in images_to_test:\n",
        "\n",
        "      # Load image and resize to expected shape [1xHxWx3]\n",
        "      image = cv2.imread(image_path)\n",
        "      image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
        "      imH, imW, _ = image.shape\n",
        "      image_resized = cv2.resize(image_rgb, (width, height))\n",
        "      input_data = np.expand_dims(image_resized, axis=0)\n",
        "\n",
        "      # Normalize pixel values if using a floating model (i.e. if model is non-quantized)\n",
        "      if float_input:\n",
        "          input_data = (np.float32(input_data) - input_mean) / input_std\n",
        "\n",
        "      # Perform the actual detection by running the model with the image as input\n",
        "      interpreter.set_tensor(input_details[0]['index'],input_data)\n",
        "      interpreter.invoke()\n",
        "\n",
        "      # Retrieve detection results\n",
        "      boxes = interpreter.get_tensor(output_details[1]['index'])[0] # Bounding box coordinates of detected objects\n",
        "      classes = interpreter.get_tensor(output_details[3]['index'])[0] # Class index of detected objects\n",
        "      scores = interpreter.get_tensor(output_details[0]['index'])[0] # Confidence of detected objects\n",
        "\n",
        "      detections = []\n",
        "\n",
        "      # Loop over all detections and draw detection box if confidence is above minimum threshold\n",
        "      for i in range(len(scores)):\n",
        "          if ((scores[i] > min_conf) and (scores[i] <= 1.0)):\n",
        "\n",
        "              # Get bounding box coordinates and draw box\n",
        "              # Interpreter can return coordinates that are outside of image dimensions, need to force them to be within image using max() and min()\n",
        "              ymin = int(max(1,(boxes[i][0] * imH)))\n",
        "              xmin = int(max(1,(boxes[i][1] * imW)))\n",
        "              ymax = int(min(imH,(boxes[i][2] * imH)))\n",
        "              xmax = int(min(imW,(boxes[i][3] * imW)))\n",
        "\n",
        "              cv2.rectangle(image, (xmin,ymin), (xmax,ymax), (10, 255, 0), 2)\n",
        "\n",
        "              # Draw label\n",
        "              object_name = labels[int(classes[i])] # Look up object name from \"labels\" array using class index\n",
        "              label = '%s: %d%%' % (object_name, int(scores[i]*100)) # Example: 'person: 72%'\n",
        "              labelSize, baseLine = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.7, 2) # Get font size\n",
        "              label_ymin = max(ymin, labelSize[1] + 10) # Make sure not to draw label too close to top of window\n",
        "              cv2.rectangle(image, (xmin, label_ymin-labelSize[1]-10), (xmin+labelSize[0], label_ymin+baseLine-10), (255, 255, 255), cv2.FILLED) # Draw white box to put label text in\n",
        "              cv2.putText(image, label, (xmin, label_ymin-7), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 0), 2) # Draw label text\n",
        "\n",
        "\n",
        "              detections.append([object_name, scores[i], xmin, ymin, xmax, ymax])\n",
        "\n",
        "              print(extract_text(xmin, xmax, ymin, ymax, image_path))\n",
        "\n",
        "      # All the results have been drawn on the image, now display the image\n",
        "      if txt_only == False: # \"text_only\" controls whether we want to display the image results or just save them in .txt files\n",
        "        image = cv2.cvtColor(image,cv2.COLOR_BGR2RGB)\n",
        "        plt.figure(figsize=(12,16))\n",
        "        plt.imshow(image)\n",
        "        plt.show()\n",
        "\n",
        "      # Save detection results in .txt files (for calculating mAP)\n",
        "      elif txt_only == True:\n",
        "\n",
        "        # Get filenames and paths\n",
        "        image_fn = os.path.basename(image_path)\n",
        "        base_fn, ext = os.path.splitext(image_fn)\n",
        "        txt_result_fn = base_fn +'.txt'\n",
        "        txt_savepath = os.path.join(savepath, txt_result_fn)\n",
        "\n",
        "        # Write results to text file\n",
        "        # (Using format defined by https://github.com/Cartucho/mAP, which will make it easy to calculate mAP)\n",
        "        with open(txt_savepath,'w') as f:\n",
        "            for detection in detections:\n",
        "                f.write('%s %.4f %d %d %d %d\\n' % (detection[0], detection[1], detection[2], detection[3], detection[4], detection[5]))\n",
        "\n",
        "  return"
      ],
      "metadata": {
        "id": "b3PxLWItoEZF"
      },
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "PATH_TO_IMAGES='/content/test'\n",
        "PATH_TO_MODEL='/content/custom_model_lite/detect.tflite'   # Path to .tflite model file\n",
        "PATH_TO_LABELS='/content/custom_model_lite/labelmap.txt'   # Path to labelmap.txt file\n",
        "min_conf_threshold=0.32   # Confidence threshold\n",
        "images_to_test = 3\n",
        "\n",
        "\n",
        "tflite_detect_images(PATH_TO_MODEL, PATH_TO_IMAGES, PATH_TO_LABELS, min_conf_threshold, images_to_test)"
      ],
      "metadata": {
        "id": "a5-DWatYoL_X"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}