notebooks/meter-reader/meter-reader.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Industrial Meter Reader\n",
    "This notebook shows how to create a industrial meter reader with OpenVINO Runtime. We use the pre-trained [PPYOLOv2](https://github.com/PaddlePaddle/PaddleDetection/tree/release/2.4/configs/ppyolo) PaddlePaddle model and [DeepLabV3P](https://github.com/PaddlePaddle/PaddleSeg/tree/release/2.5/configs/deeplabv3p) to build up a multiple inference task pipeline:\n",
    "\n",
    "1. Run detection model to find the meters, and crop them from the origin photo.\n",
    "2. Run segmentation model on these cropped meters to get the pointer and scale instance.\n",
    "3. Find the location of the pointer in scale map.\n",
    "\n",
    "![workflow](https://user-images.githubusercontent.com/91237924/166137115-67284fa5-f703-4468-98f4-c43d2c584763.png)\n",
    "\n",
    "\n",
    "\n",
    "#### Table of contents:\n",
    "\n",
    "- [Import](#Import)\n",
    "- [Prepare the Model and Test Image](#Prepare-the-Model-and-Test-Image)\n",
    "- [Configuration](#Configuration)\n",
    "- [Load the Models](#Load-the-Models)\n",
    "- [Data Process](#Data-Process)\n",
    "- [Main Function](#Main-Function)\n",
    "    - [Initialize the model and parameters.](#Initialize-the-model-and-parameters.)\n",
    "    - [Run meter detection model](#Run-meter-detection-model)\n",
    "    - [Run meter segmentation model](#Run-meter-segmentation-model)\n",
    "    - [Postprocess the models result and calculate the final readings](#Postprocess-the-models-result-and-calculate-the-final-readings)\n",
    "    - [Get the reading result on the meter picture](#Get-the-reading-result-on-the-meter-picture)\n",
    "- [Try it with your meter photos!](#Try-it-with-your-meter-photos!)\n",
    "\n",
    "\n",
    "### Installation Instructions\n",
    "\n",
    "This is a self-contained example that relies solely on its own code.\n",
    "\n",
    "We recommend  running the notebook in a virtual environment. You only need a Jupyter server to start.\n",
    "For details, please refer to [Installation Guide](https://github.com/openvinotoolkit/openvino_notebooks/blob/latest/README.md#-installation-guide).\n",
    "\n",
    "<img referrerpolicy=\"no-referrer-when-downgrade\" src=\"https://static.scarf.sh/a.png?x-pxid=5b5a4db0-7875-4bfb-bdbd-01698b5b1a77&file=notebooks/meter-reader/meter-reader.ipynb\" />\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import platform\n",
    "\n",
    "# Install openvino package\n",
    "%pip install -q \"openvino>=2023.1.0\" opencv-python tqdm\n",
    "\n",
    "if platform.system() != \"Windows\":\n",
    "    %pip install -q \"matplotlib>=3.4\"\n",
    "else:\n",
    "    %pip install -q \"matplotlib>=3.4,<3.7\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import\n",
    "[back to top ⬆️](#Table-of-contents:)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from pathlib import Path\n",
    "import numpy as np\n",
    "import math\n",
    "import cv2\n",
    "import tarfile\n",
    "import matplotlib.pyplot as plt\n",
    "import openvino as ov\n",
    "\n",
    "# Fetch `notebook_utils` module\n",
    "import requests\n",
    "\n",
    "r = requests.get(\n",
    "    url=\"https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/utils/notebook_utils.py\",\n",
    ")\n",
    "\n",
    "open(\"notebook_utils.py\", \"w\").write(r.text)\n",
    "from notebook_utils import download_file, segmentation_map_to_image"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare the Model and Test Image\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "Download PPYOLOv2 and DeepLabV3P pre-trained models from PaddlePaddle community."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "MODEL_DIR = \"model\"\n",
    "DATA_DIR = \"data\"\n",
    "DET_MODEL_LINK = \"https://storage.openvinotoolkit.org/repositories/openvino_notebooks/models/meter-reader/meter_det_model.tar.gz\"\n",
    "SEG_MODEL_LINK = \"https://storage.openvinotoolkit.org/repositories/openvino_notebooks/models/meter-reader/meter_seg_model.tar.gz\"\n",
    "DET_FILE_NAME = DET_MODEL_LINK.split(\"/\")[-1]\n",
    "SEG_FILE_NAME = SEG_MODEL_LINK.split(\"/\")[-1]\n",
    "IMG_LINK = \"https://user-images.githubusercontent.com/91237924/170696219-f68699c6-1e82-46bf-aaed-8e2fc3fa5f7b.jpg\"\n",
    "IMG_FILE_NAME = IMG_LINK.split(\"/\")[-1]\n",
    "IMG_PATH = Path(f\"{DATA_DIR}/{IMG_FILE_NAME}\")\n",
    "\n",
    "os.makedirs(MODEL_DIR, exist_ok=True)\n",
    "\n",
    "download_file(DET_MODEL_LINK, directory=MODEL_DIR, show_progress=True)\n",
    "file = tarfile.open(f\"model/{DET_FILE_NAME}\")\n",
    "res = file.extractall(\"model\")\n",
    "if not res:\n",
    "    print(f'Detection Model Extracted to \"./{MODEL_DIR}\".')\n",
    "else:\n",
    "    print(\"Error Extracting the Detection model. Please check the network.\")\n",
    "\n",
    "download_file(SEG_MODEL_LINK, directory=MODEL_DIR, show_progress=True)\n",
    "file = tarfile.open(f\"model/{SEG_FILE_NAME}\")\n",
    "res = file.extractall(\"model\")\n",
    "if not res:\n",
    "    print(f'Segmentation Model Extracted to \"./{MODEL_DIR}\".')\n",
    "else:\n",
    "    print(\"Error Extracting the Segmentation model. Please check the network.\")\n",
    "\n",
    "download_file(IMG_LINK, directory=DATA_DIR, show_progress=True)\n",
    "if IMG_PATH.is_file():\n",
    "    print(f'Test Image Saved to \"./{DATA_DIR}\".')\n",
    "else:\n",
    "    print(\"Error Downloading the Test Image. Please check the network.\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configuration\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "Add parameter configuration for reading calculation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "METER_SHAPE = [512, 512]\n",
    "CIRCLE_CENTER = [256, 256]\n",
    "CIRCLE_RADIUS = 250\n",
    "PI = math.pi\n",
    "RECTANGLE_HEIGHT = 120\n",
    "RECTANGLE_WIDTH = 1570\n",
    "TYPE_THRESHOLD = 40\n",
    "COLORMAP = np.array([[28, 28, 28], [238, 44, 44], [250, 250, 250]])\n",
    "\n",
    "# There are 2 types of meters in test image datasets\n",
    "METER_CONFIG = [\n",
    "    {\"scale_interval_value\": 25.0 / 50.0, \"range\": 25.0, \"unit\": \"(MPa)\"},\n",
    "    {\"scale_interval_value\": 1.6 / 32.0, \"range\": 1.6, \"unit\": \"(MPa)\"},\n",
    "]\n",
    "\n",
    "SEG_LABEL = {\"background\": 0, \"pointer\": 1, \"scale\": 2}"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load the Models\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "Define a common class for model loading and inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize OpenVINO Runtime\n",
    "core = ov.Core()\n",
    "\n",
    "\n",
    "class Model:\n",
    "    \"\"\"\n",
    "    This class represents a OpenVINO model object.\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, model_path, new_shape, device=\"CPU\"):\n",
    "        \"\"\"\n",
    "        Initialize the model object\n",
    "\n",
    "        Param:\n",
    "            model_path (string): path of inference model\n",
    "            new_shape (dict): new shape of model input\n",
    "\n",
    "        \"\"\"\n",
    "        self.model = core.read_model(model=model_path)\n",
    "        self.model.reshape(new_shape)\n",
    "        self.compiled_model = core.compile_model(model=self.model, device_name=device)\n",
    "        self.output_layer = self.compiled_model.output(0)\n",
    "\n",
    "    def predict(self, input_image):\n",
    "        \"\"\"\n",
    "        Run inference\n",
    "\n",
    "        Param:\n",
    "            input_image (np.array): input data\n",
    "\n",
    "        Retuns:\n",
    "            result (np.array)): model output data\n",
    "        \"\"\"\n",
    "        result = self.compiled_model(input_image)[self.output_layer]\n",
    "        return result"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Process\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "Including the preprocessing and postprocessing tasks of each model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def det_preprocess(input_image, target_size):\n",
    "    \"\"\"\n",
    "    Preprocessing the input data for detection task\n",
    "\n",
    "    Param:\n",
    "        input_image (np.array): input data\n",
    "        size (int): the image size required by model input layer\n",
    "    Retuns:\n",
    "        img.astype (np.array): preprocessed image\n",
    "\n",
    "    \"\"\"\n",
    "    img = cv2.resize(input_image, (target_size, target_size))\n",
    "    img = np.transpose(img, [2, 0, 1]) / 255\n",
    "    img = np.expand_dims(img, 0)\n",
    "    img_mean = np.array([0.485, 0.456, 0.406]).reshape((3, 1, 1))\n",
    "    img_std = np.array([0.229, 0.224, 0.225]).reshape((3, 1, 1))\n",
    "    img -= img_mean\n",
    "    img /= img_std\n",
    "    return img.astype(np.float32)\n",
    "\n",
    "\n",
    "def filter_bboxes(det_results, score_threshold):\n",
    "    \"\"\"\n",
    "    Filter out the detection results with low confidence\n",
    "\n",
    "    Param：\n",
    "        det_results (list[dict]): detection results\n",
    "        score_threshold (float)： confidence threshold\n",
    "\n",
    "    Retuns：\n",
    "        filtered_results (list[dict]): filter detection results\n",
    "\n",
    "    \"\"\"\n",
    "    filtered_results = []\n",
    "    for i in range(len(det_results)):\n",
    "        if det_results[i, 1] > score_threshold:\n",
    "            filtered_results.append(det_results[i])\n",
    "    return filtered_results\n",
    "\n",
    "\n",
    "def roi_crop(image, results, scale_x, scale_y):\n",
    "    \"\"\"\n",
    "    Crop the area of detected meter of original image\n",
    "\n",
    "    Param：\n",
    "        img (np.array)：original image。\n",
    "        det_results (list[dict]): detection results\n",
    "        scale_x (float): the scale value in x axis\n",
    "        scale_y (float): the scale value in y axis\n",
    "\n",
    "    Retuns：\n",
    "        roi_imgs (list[np.array]): the list of meter images\n",
    "        loc (list[int]): the list of meter locations\n",
    "\n",
    "    \"\"\"\n",
    "    roi_imgs = []\n",
    "    loc = []\n",
    "    for result in results:\n",
    "        bbox = result[2:]\n",
    "        xmin, ymin, xmax, ymax = [\n",
    "            int(bbox[0] * scale_x),\n",
    "            int(bbox[1] * scale_y),\n",
    "            int(bbox[2] * scale_x),\n",
    "            int(bbox[3] * scale_y),\n",
    "        ]\n",
    "        sub_img = image[ymin : (ymax + 1), xmin : (xmax + 1), :]\n",
    "        roi_imgs.append(sub_img)\n",
    "        loc.append([xmin, ymin, xmax, ymax])\n",
    "    return roi_imgs, loc\n",
    "\n",
    "\n",
    "def roi_process(input_images, target_size, interp=cv2.INTER_LINEAR):\n",
    "    \"\"\"\n",
    "    Prepare the roi image of detection results data\n",
    "    Preprocessing the input data for segmentation task\n",
    "\n",
    "    Param：\n",
    "        input_images (list[np.array])：the list of meter images\n",
    "        target_size (list|tuple)： height and width of resized image， e.g [heigh,width]\n",
    "        interp (int)：the interp method for image reszing\n",
    "\n",
    "    Retuns：\n",
    "        img_list (list[np.array])：the list of processed images\n",
    "        resize_img (list[np.array]): for visualization\n",
    "\n",
    "    \"\"\"\n",
    "    img_list = list()\n",
    "    resize_list = list()\n",
    "    for img in input_images:\n",
    "        img_shape = img.shape\n",
    "        scale_x = float(target_size[1]) / float(img_shape[1])\n",
    "        scale_y = float(target_size[0]) / float(img_shape[0])\n",
    "        resize_img = cv2.resize(img, None, None, fx=scale_x, fy=scale_y, interpolation=interp)\n",
    "        resize_list.append(resize_img)\n",
    "        resize_img = resize_img.transpose(2, 0, 1) / 255\n",
    "        img_mean = np.array([0.5, 0.5, 0.5]).reshape((3, 1, 1))\n",
    "        img_std = np.array([0.5, 0.5, 0.5]).reshape((3, 1, 1))\n",
    "        resize_img -= img_mean\n",
    "        resize_img /= img_std\n",
    "        img_list.append(resize_img)\n",
    "    return img_list, resize_list\n",
    "\n",
    "\n",
    "def erode(seg_results, erode_kernel):\n",
    "    \"\"\"\n",
    "    Erode the segmentation result to get the more clear instance of pointer and scale\n",
    "\n",
    "    Param：\n",
    "        seg_results (list[dict])：segmentation results\n",
    "        erode_kernel (int): size of erode_kernel\n",
    "\n",
    "    Return：\n",
    "        eroded_results (list[dict])： the lab map of eroded_results\n",
    "\n",
    "    \"\"\"\n",
    "    kernel = np.ones((erode_kernel, erode_kernel), np.uint8)\n",
    "    eroded_results = seg_results\n",
    "    for i in range(len(seg_results)):\n",
    "        eroded_results[i] = cv2.erode(seg_results[i].astype(np.uint8), kernel)\n",
    "    return eroded_results\n",
    "\n",
    "\n",
    "def circle_to_rectangle(seg_results):\n",
    "    \"\"\"\n",
    "    Switch the shape of label_map from circle to rectangle\n",
    "\n",
    "    Param：\n",
    "        seg_results (list[dict])：segmentation results\n",
    "\n",
    "    Return：\n",
    "        rectangle_meters (list[np.array])：the rectangle of label map\n",
    "\n",
    "    \"\"\"\n",
    "    rectangle_meters = list()\n",
    "    for i, seg_result in enumerate(seg_results):\n",
    "        label_map = seg_result\n",
    "\n",
    "        # The size of rectangle_meter is determined by RECTANGLE_HEIGHT and RECTANGLE_WIDTH\n",
    "        rectangle_meter = np.zeros((RECTANGLE_HEIGHT, RECTANGLE_WIDTH), dtype=np.uint8)\n",
    "        for row in range(RECTANGLE_HEIGHT):\n",
    "            for col in range(RECTANGLE_WIDTH):\n",
    "                theta = PI * 2 * (col + 1) / RECTANGLE_WIDTH\n",
    "\n",
    "                # The radius of meter circle will be mapped to the height of rectangle image\n",
    "                rho = CIRCLE_RADIUS - row - 1\n",
    "                y = int(CIRCLE_CENTER[0] + rho * math.cos(theta) + 0.5)\n",
    "                x = int(CIRCLE_CENTER[1] - rho * math.sin(theta) + 0.5)\n",
    "                rectangle_meter[row, col] = label_map[y, x]\n",
    "        rectangle_meters.append(rectangle_meter)\n",
    "    return rectangle_meters\n",
    "\n",
    "\n",
    "def rectangle_to_line(rectangle_meters):\n",
    "    \"\"\"\n",
    "    Switch the dimension of rectangle label map from 2D to 1D\n",
    "\n",
    "    Param：\n",
    "        rectangle_meters (list[np.array])：2D rectangle OF label_map。\n",
    "\n",
    "    Return：\n",
    "        line_scales (list[np.array])： the list of scales value\n",
    "        line_pointers (list[np.array])：the list of pointers value\n",
    "\n",
    "    \"\"\"\n",
    "    line_scales = list()\n",
    "    line_pointers = list()\n",
    "    for rectangle_meter in rectangle_meters:\n",
    "        height, width = rectangle_meter.shape[0:2]\n",
    "        line_scale = np.zeros((width), dtype=np.uint8)\n",
    "        line_pointer = np.zeros((width), dtype=np.uint8)\n",
    "        for col in range(width):\n",
    "            for row in range(height):\n",
    "                if rectangle_meter[row, col] == SEG_LABEL[\"pointer\"]:\n",
    "                    line_pointer[col] += 1\n",
    "                elif rectangle_meter[row, col] == SEG_LABEL[\"scale\"]:\n",
    "                    line_scale[col] += 1\n",
    "        line_scales.append(line_scale)\n",
    "        line_pointers.append(line_pointer)\n",
    "    return line_scales, line_pointers\n",
    "\n",
    "\n",
    "def mean_binarization(data_list):\n",
    "    \"\"\"\n",
    "    Binarize the data\n",
    "\n",
    "    Param：\n",
    "        data_list (list[np.array])：input data\n",
    "\n",
    "    Return：\n",
    "        binaried_data_list (list[np.array])：output data。\n",
    "\n",
    "    \"\"\"\n",
    "    batch_size = len(data_list)\n",
    "    binaried_data_list = data_list\n",
    "    for i in range(batch_size):\n",
    "        mean_data = np.mean(data_list[i])\n",
    "        width = data_list[i].shape[0]\n",
    "        for col in range(width):\n",
    "            if data_list[i][col] < mean_data:\n",
    "                binaried_data_list[i][col] = 0\n",
    "            else:\n",
    "                binaried_data_list[i][col] = 1\n",
    "    return binaried_data_list\n",
    "\n",
    "\n",
    "def locate_scale(line_scales):\n",
    "    \"\"\"\n",
    "    Find location of center of each scale\n",
    "\n",
    "    Param：\n",
    "        line_scales (list[np.array])：the list of binaried scales value\n",
    "\n",
    "    Return：\n",
    "        scale_locations (list[list])：location of each scale\n",
    "\n",
    "    \"\"\"\n",
    "    batch_size = len(line_scales)\n",
    "    scale_locations = list()\n",
    "    for i in range(batch_size):\n",
    "        line_scale = line_scales[i]\n",
    "        width = line_scale.shape[0]\n",
    "        find_start = False\n",
    "        one_scale_start = 0\n",
    "        one_scale_end = 0\n",
    "        locations = list()\n",
    "        for j in range(width - 1):\n",
    "            if line_scale[j] > 0 and line_scale[j + 1] > 0:\n",
    "                if not find_start:\n",
    "                    one_scale_start = j\n",
    "                    find_start = True\n",
    "            if find_start:\n",
    "                if line_scale[j] == 0 and line_scale[j + 1] == 0:\n",
    "                    one_scale_end = j - 1\n",
    "                    one_scale_location = (one_scale_start + one_scale_end) / 2\n",
    "                    locations.append(one_scale_location)\n",
    "                    one_scale_start = 0\n",
    "                    one_scale_end = 0\n",
    "                    find_start = False\n",
    "        scale_locations.append(locations)\n",
    "    return scale_locations\n",
    "\n",
    "\n",
    "def locate_pointer(line_pointers):\n",
    "    \"\"\"\n",
    "    Find location of center of pointer\n",
    "\n",
    "    Param：\n",
    "        line_scales (list[np.array])：the list of binaried pointer value\n",
    "\n",
    "    Return：\n",
    "        scale_locations (list[list])：location of pointer\n",
    "\n",
    "    \"\"\"\n",
    "    batch_size = len(line_pointers)\n",
    "    pointer_locations = list()\n",
    "    for i in range(batch_size):\n",
    "        line_pointer = line_pointers[i]\n",
    "        find_start = False\n",
    "        pointer_start = 0\n",
    "        pointer_end = 0\n",
    "        location = 0\n",
    "        width = line_pointer.shape[0]\n",
    "        for j in range(width - 1):\n",
    "            if line_pointer[j] > 0 and line_pointer[j + 1] > 0:\n",
    "                if not find_start:\n",
    "                    pointer_start = j\n",
    "                    find_start = True\n",
    "            if find_start:\n",
    "                if line_pointer[j] == 0 and line_pointer[j + 1] == 0:\n",
    "                    pointer_end = j - 1\n",
    "                    location = (pointer_start + pointer_end) / 2\n",
    "                    find_start = False\n",
    "                    break\n",
    "        pointer_locations.append(location)\n",
    "    return pointer_locations\n",
    "\n",
    "\n",
    "def get_relative_location(scale_locations, pointer_locations):\n",
    "    \"\"\"\n",
    "    Match location of pointer and scales\n",
    "\n",
    "    Param：\n",
    "        scale_locations (list[list])：location of each scale\n",
    "        pointer_locations (list[list])：location of pointer\n",
    "\n",
    "    Return：\n",
    "        pointed_scales (list[dict])： a list of dict with:\n",
    "                                     'num_scales': total number of scales\n",
    "                                     'pointed_scale': predicted number of scales\n",
    "\n",
    "    \"\"\"\n",
    "    pointed_scales = list()\n",
    "    for scale_location, pointer_location in zip(scale_locations, pointer_locations):\n",
    "        num_scales = len(scale_location)\n",
    "        pointed_scale = -1\n",
    "        if num_scales > 0:\n",
    "            for i in range(num_scales - 1):\n",
    "                if scale_location[i] <= pointer_location < scale_location[i + 1]:\n",
    "                    pointed_scale = i + (pointer_location - scale_location[i]) / (scale_location[i + 1] - scale_location[i] + 1e-05) + 1\n",
    "        result = {\"num_scales\": num_scales, \"pointed_scale\": pointed_scale}\n",
    "        pointed_scales.append(result)\n",
    "    return pointed_scales\n",
    "\n",
    "\n",
    "def calculate_reading(pointed_scales):\n",
    "    \"\"\"\n",
    "    Calculate the value of meter according to the type of meter\n",
    "\n",
    "    Param：\n",
    "        pointed_scales (list[list])：predicted number of scales\n",
    "\n",
    "    Return：\n",
    "        readings (list[float])： the list of values read from meter\n",
    "\n",
    "    \"\"\"\n",
    "    readings = list()\n",
    "    batch_size = len(pointed_scales)\n",
    "    for i in range(batch_size):\n",
    "        pointed_scale = pointed_scales[i]\n",
    "        # find the type of meter according the total number of scales\n",
    "        if pointed_scale[\"num_scales\"] > TYPE_THRESHOLD:\n",
    "            reading = pointed_scale[\"pointed_scale\"] * METER_CONFIG[0][\"scale_interval_value\"]\n",
    "        else:\n",
    "            reading = pointed_scale[\"pointed_scale\"] * METER_CONFIG[1][\"scale_interval_value\"]\n",
    "        readings.append(reading)\n",
    "    return readings"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Main Function\n",
    "[back to top ⬆️](#Table-of-contents:)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initialize the model and parameters.\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "\n",
    "select device from dropdown list for running inference using OpenVINO"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ea1b8e76e34a4f5c88e13529e14766d3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Dropdown(description='Device:', index=2, options=('CPU', 'GPU', 'AUTO'), value='AUTO')"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import ipywidgets as widgets\n",
    "\n",
    "device = widgets.Dropdown(\n",
    "    options=core.available_devices + [\"AUTO\"],\n",
    "    value=\"AUTO\",\n",
    "    description=\"Device:\",\n",
    "    disabled=False,\n",
    ")\n",
    "\n",
    "device"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The number of detected meter from detection network can be arbitrary in some scenarios, which means the batch size of segmentation network input is a [dynamic dimension](https://docs.openvino.ai/2024/openvino-workflow/running-inference/dynamic-shapes.html), and it should be specified as ```-1``` or the ```ov::Dimension()``` instead of a positive number used for static dimensions. In this case, for memory consumption optimization, we can specify the lower and/or upper bounds of input batch size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "img_file = f\"{DATA_DIR}/{IMG_FILE_NAME}\"\n",
    "det_model_path = f\"{MODEL_DIR}/meter_det_model/model.pdmodel\"\n",
    "det_model_shape = {\n",
    "    \"image\": [1, 3, 608, 608],\n",
    "    \"im_shape\": [1, 2],\n",
    "    \"scale_factor\": [1, 2],\n",
    "}\n",
    "seg_model_path = f\"{MODEL_DIR}/meter_seg_model/model.pdmodel\"\n",
    "seg_model_shape = {\"image\": [ov.Dimension(1, 2), 3, 512, 512]}\n",
    "\n",
    "erode_kernel = 4\n",
    "score_threshold = 0.5\n",
    "seg_batch_size = 2\n",
    "input_shape = 608\n",
    "\n",
    "# Intialize the model objects\n",
    "detector = Model(det_model_path, det_model_shape, device.value)\n",
    "segmenter = Model(seg_model_path, seg_model_shape, device.value)\n",
    "\n",
    "# Visulize a original input photo\n",
    "image = cv2.imread(img_file)\n",
    "rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
    "plt.imshow(rgb_image)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run meter detection model\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "Detect the location of the meter and prepare the ROI images for segmentation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Prepare the input data for meter detection model\n",
    "im_shape = np.array([[input_shape, input_shape]]).astype(\"float32\")\n",
    "scale_factor = np.array([[1, 2]]).astype(\"float32\")\n",
    "input_image = det_preprocess(image, input_shape)\n",
    "inputs_dict = {\"image\": input_image, \"im_shape\": im_shape, \"scale_factor\": scale_factor}\n",
    "\n",
    "# Run meter detection model\n",
    "det_results = detector.predict(inputs_dict)\n",
    "\n",
    "# Filter out the bounding box with low confidence\n",
    "filtered_results = filter_bboxes(det_results, score_threshold)\n",
    "\n",
    "# Prepare the input data for meter segmentation model\n",
    "scale_x = image.shape[1] / input_shape * 2\n",
    "scale_y = image.shape[0] / input_shape\n",
    "\n",
    "# Create the individual picture for each detected meter\n",
    "roi_imgs, loc = roi_crop(image, filtered_results, scale_x, scale_y)\n",
    "roi_imgs, resize_imgs = roi_process(roi_imgs, METER_SHAPE)\n",
    "\n",
    "# Create the pictures of detection results\n",
    "roi_stack = np.hstack(resize_imgs)\n",
    "\n",
    "if cv2.imwrite(f\"{DATA_DIR}/detection_results.jpg\", roi_stack):\n",
    "    print('The detection result image has been saved as \"detection_results.jpg\" in data')\n",
    "    plt.imshow(cv2.cvtColor(roi_stack, cv2.COLOR_BGR2RGB))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run meter segmentation model\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "Get the results of segmentation task on detected ROI."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "seg_results = list()\n",
    "mask_list = list()\n",
    "num_imgs = len(roi_imgs)\n",
    "\n",
    "# Run meter segmentation model on all detected meters\n",
    "for i in range(0, num_imgs, seg_batch_size):\n",
    "    batch = roi_imgs[i : min(num_imgs, i + seg_batch_size)]\n",
    "    seg_result = segmenter.predict({\"image\": np.array(batch)})\n",
    "    seg_results.extend(seg_result)\n",
    "results = []\n",
    "for i in range(len(seg_results)):\n",
    "    results.append(np.argmax(seg_results[i], axis=0))\n",
    "seg_results = erode(results, erode_kernel)\n",
    "\n",
    "# Create the pictures of segmentation results\n",
    "for i in range(len(seg_results)):\n",
    "    mask_list.append(segmentation_map_to_image(seg_results[i], COLORMAP))\n",
    "mask_stack = np.hstack(mask_list)\n",
    "\n",
    "if cv2.imwrite(f\"{DATA_DIR}/segmentation_results.jpg\", cv2.cvtColor(mask_stack, cv2.COLOR_RGB2BGR)):\n",
    "    print('The segmentation result image has been saved as \"segmentation_results.jpg\" in data')\n",
    "    plt.imshow(mask_stack)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Postprocess the models result and calculate the final readings\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "Use OpenCV function to find the location of the pointer in a scale map."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Find the pointer location in scale map and calculate the meters reading\n",
    "rectangle_meters = circle_to_rectangle(seg_results)\n",
    "line_scales, line_pointers = rectangle_to_line(rectangle_meters)\n",
    "binaried_scales = mean_binarization(line_scales)\n",
    "binaried_pointers = mean_binarization(line_pointers)\n",
    "scale_locations = locate_scale(binaried_scales)\n",
    "pointer_locations = locate_pointer(binaried_pointers)\n",
    "pointed_scales = get_relative_location(scale_locations, pointer_locations)\n",
    "meter_readings = calculate_reading(pointed_scales)\n",
    "\n",
    "rectangle_list = list()\n",
    "# Plot the rectangle meters\n",
    "for i in range(len(rectangle_meters)):\n",
    "    rectangle_list.append(segmentation_map_to_image(rectangle_meters[i], COLORMAP))\n",
    "rectangle_meters_stack = np.hstack(rectangle_list)\n",
    "\n",
    "if cv2.imwrite(\n",
    "    f\"{DATA_DIR}/rectangle_meters.jpg\",\n",
    "    cv2.cvtColor(rectangle_meters_stack, cv2.COLOR_RGB2BGR),\n",
    "):\n",
    "    print('The rectangle_meters result image has been saved as \"rectangle_meters.jpg\" in data')\n",
    "    plt.imshow(rectangle_meters_stack)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Get the reading result on the meter picture\n",
    "[back to top ⬆️](#Table-of-contents:)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a final result photo with reading\n",
    "for i in range(len(meter_readings)):\n",
    "    print(\"Meter {}: {:.3f}\".format(i + 1, meter_readings[i]))\n",
    "\n",
    "result_image = image.copy()\n",
    "for i in range(len(loc)):\n",
    "    cv2.rectangle(result_image, (loc[i][0], loc[i][1]), (loc[i][2], loc[i][3]), (0, 150, 0), 3)\n",
    "    font = cv2.FONT_HERSHEY_SIMPLEX\n",
    "    cv2.rectangle(\n",
    "        result_image,\n",
    "        (loc[i][0], loc[i][1]),\n",
    "        (loc[i][0] + 100, loc[i][1] + 40),\n",
    "        (0, 150, 0),\n",
    "        -1,\n",
    "    )\n",
    "    cv2.putText(\n",
    "        result_image,\n",
    "        \"#{:.3f}\".format(meter_readings[i]),\n",
    "        (loc[i][0], loc[i][1] + 25),\n",
    "        font,\n",
    "        0.8,\n",
    "        (255, 255, 255),\n",
    "        2,\n",
    "        cv2.LINE_AA,\n",
    "    )\n",
    "if cv2.imwrite(f\"{DATA_DIR}/reading_results.jpg\", result_image):\n",
    "    print('The reading results image has been saved as \"reading_results.jpg\" in data')\n",
    "    plt.imshow(cv2.cvtColor(result_image, cv2.COLOR_BGR2RGB))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Try it with your meter photos!\n",
    "[back to top ⬆️](#Table-of-contents:)\n"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "e7370f93d1d0cde622a1f8e1c04877d8463912d04d973331ad4851f04de6915a"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  },
  "openvino_notebooks": {
   "imageUrl": "https://user-images.githubusercontent.com/91237924/166135627-194405b0-6c25-4fd8-9ad1-83fb3a00a081.jpg",
   "tags": {
    "categories": [
     "Model Demos"
    ],
    "libraries": [],
    "other": [
     "YOLO"
    ],
    "tasks": [
     "Object Detection",
     "Image Segmentation"
    ]
   }
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}