notebooks/209-handwritten-ocr/209-handwritten-ocr.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0293f65e",
   "metadata": {},
   "source": [
    "# Handwritten Chinese and Japanese OCR\n",
    "\n",
    "In this tutorial, we perform optical character recognition (OCR) for handwritten Chinese (simplified) and Japanese. An OCR tutorial using the Latin alphabet is available in [notebook 208](../208-optical-character-recognition/208-optical-character-recognition.ipynb). This model is capable of processing only one line of symbols at a time. \n",
    "\n",
    "The models used in this notebook are [handwritten-japanese-recognition](https://docs.openvinotoolkit.org/latest/omz_models_model_handwritten_japanese_recognition_0001.html) and [handwritten-simplified-chinese](https://docs.openvinotoolkit.org/latest/omz_models_model_handwritten_simplified_chinese_recognition_0001.html). To decode model outputs as readable text [kondate_nakayosi](https://github.com/openvinotoolkit/open_model_zoo/blob/master/data/dataset_classes/kondate_nakayosi.txt) and [scut_ept](https://github.com/openvinotoolkit/open_model_zoo/blob/master/data/dataset_classes/scut_ept.txt) charlists are used. Both models are available on [Open Model Zoo](https://github.com/openvinotoolkit/open_model_zoo/)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fda2e1e0",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0a6a0d5",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "from collections import namedtuple\n",
    "from itertools import groupby\n",
    "from pathlib import Path\n",
    "\n",
    "import cv2\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from openvino.runtime import Core"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a46517a",
   "metadata": {},
   "source": [
    "## Settings\n",
    "\n",
    "Set up all constants and folders used in this notebook"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b50b2bad",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Directories where data will be placed\n",
    "model_folder = \"model\"\n",
    "data_folder = \"data\"\n",
    "charlist_folder = f\"{data_folder}/charlists\"\n",
    "\n",
    "# Precision used by model\n",
    "precision = \"FP16\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f3c93fed",
   "metadata": {},
   "source": [
    "To group files, you have to define the collection. In this case, you can use `namedtuple`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17402fcc",
   "metadata": {},
   "outputs": [],
   "source": [
    "Language = namedtuple(\n",
    "    typename=\"Language\", field_names=[\"model_name\", \"charlist_name\", \"demo_image_name\"]\n",
    ")\n",
    "chinese_files = Language(\n",
    "    model_name=\"handwritten-simplified-chinese-recognition-0001\",\n",
    "    charlist_name=\"chinese_charlist.txt\",\n",
    "    demo_image_name=\"handwritten_chinese_test.jpg\",\n",
    ")\n",
    "japanese_files = Language(\n",
    "    model_name=\"handwritten-japanese-recognition-0001\",\n",
    "    charlist_name=\"japanese_charlist.txt\",\n",
    "    demo_image_name=\"handwritten_japanese_test.png\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c61827c",
   "metadata": {},
   "source": [
    "## Select Language\n",
    "\n",
    "Depending on your choice you will need to change a line of code in the cell below.\n",
    "\n",
    "If you want to do Japanese OCR, this line should be set to ```language = 'japanese'``` for Chinese set ```language = 'chinese'```."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d3b1190",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Select language by using either language='chinese' or language='japanese'\n",
    "language = \"chinese\"\n",
    "\n",
    "languages = {\"chinese\": chinese_files, \"japanese\": japanese_files}\n",
    "\n",
    "selected_language = languages.get(language)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "750355d4",
   "metadata": {},
   "source": [
    "## Download Model\n",
    "\n",
    "In addition to images and charlists, we need to download the model file. In the sections below there are cells for downloading either the Chinese or Japanese model.\n",
    " \n",
    "If it is your first time running the notebook, the model will download. It may take a few minutes. \n",
    "\n",
    "We use `omz_downloader`, which is a command-line tool from the `openvino-dev` package. `omz_downloader` automatically creates a directory structure and downloads the selected model. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8f266fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "path_to_model_weights = Path(f'{model_folder}/intel/{selected_language.model_name}/{precision}/{selected_language.model_name}.bin')\n",
    "if not path_to_model_weights.is_file():\n",
    "    download_command = f'omz_downloader --name {selected_language.model_name} --output_dir {model_folder} --precision {precision}'\n",
    "    print(download_command)\n",
    "    ! $download_command"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d026308",
   "metadata": {},
   "source": [
    "## Load Network and Execute\n",
    "\n",
    "When all files are downloaded and language is selected, you need to read and compile the network to run inference. The path to the model is defined based on the selected language."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7114e467",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "ie = Core()\n",
    "path_to_model = path_to_model_weights.with_suffix(\".xml\")\n",
    "model = ie.read_model(model=path_to_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ecb4551",
   "metadata": {},
   "source": [
    "### Select Device Name\n",
    "\n",
    "You may choose to run the network on multiple devices by default it will load the model on the CPU (you can choose manually CPU, GPU etc.) or let the engine choose the best available device (AUTO).\n",
    "\n",
    "To list all available devices that you can use, uncomment and run the line ```print(ie.available_devices)```."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "024ad3a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# To check available device names run the line below\n",
    "# print(ie.available_devices)\n",
    "\n",
    "compiled_model = ie.compile_model(model=model, device_name=\"CPU\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3df5b652",
   "metadata": {},
   "source": [
    "## Fetch Information About Input and Output Layers \n",
    "\n",
    "Now that the model is loaded, you need to fetch information about input and output layers. This is information about the input shape and the output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d31411a",
   "metadata": {},
   "outputs": [],
   "source": [
    "recognition_output_layer = compiled_model.output(0)\n",
    "recognition_input_layer = compiled_model.input(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41307bc7",
   "metadata": {},
   "source": [
    "## Load an Image\n",
    "\n",
    "The next step is to load an image. \n",
    "\n",
    "The model expects a single-channel image as input, which is why we read the image in grayscale.\n",
    "\n",
    "After loading the input image, the next step is getting information that you will use for calculating the scale ratio. This describes the ratio between required input layer height and the current image height. In the cell below, the image will be resized and padded to keep letters proportional and meet input shape."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "afa1d51a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Read file name of demo file based on the selected model\n",
    "\n",
    "file_name = selected_language.demo_image_name\n",
    "\n",
    "# Text detection models expects an image in grayscale format\n",
    "# IMPORTANT!!! This model allows to read only one line at time\n",
    "\n",
    "# Read image\n",
    "image = cv2.imread(filename=f\"{data_folder}/{file_name}\", flags=cv2.IMREAD_GRAYSCALE)\n",
    "\n",
    "# Fetch shape\n",
    "image_height, _ = image.shape\n",
    "\n",
    "# B,C,H,W = batch size, number of channels, height, width\n",
    "_, _, H, W = recognition_input_layer.shape\n",
    "\n",
    "# Calculate scale ratio between input shape height and image height to resize image\n",
    "scale_ratio = H / image_height\n",
    "\n",
    "# Resize image to expected input sizes\n",
    "resized_image = cv2.resize(\n",
    "    image, None, fx=scale_ratio, fy=scale_ratio, interpolation=cv2.INTER_AREA\n",
    ")\n",
    "\n",
    "# Pad image to match input size, without changing aspect ratio\n",
    "resized_image = np.pad(\n",
    "    resized_image, ((0, 0), (0, W - resized_image.shape[1])), mode=\"edge\"\n",
    ")\n",
    "\n",
    "# Reshape to network the input shape\n",
    "input_image = resized_image[None, None, :, :]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba364bee",
   "metadata": {},
   "source": [
    "## Visualise Input Image\n",
    "\n",
    "After preprocessing you can display the image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7a8e437",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(20, 1))\n",
    "plt.axis(\"off\")\n",
    "plt.imshow(resized_image, cmap=\"gray\", vmin=0, vmax=255);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6d4c642",
   "metadata": {},
   "source": [
    "## Prepare Charlist\n",
    "\n",
    "The model is loaded and the image is ready. The only element left is the charlist which is downloaded, but before we use it, there is one more step. You must add a blank symbol at the beginning of the charlist. This is expected for both the Chinese and Japanese models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2cb18186",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get dictionary to encode output, based on model documentation\n",
    "used_charlist = selected_language.charlist_name\n",
    "\n",
    "# With both models, there should be blank symbol added at index 0 of each charlist\n",
    "blank_char = \"~\"\n",
    "\n",
    "with open(f\"{charlist_folder}/{used_charlist}\", \"r\", encoding=\"utf-8\") as charlist:\n",
    "    letters = blank_char + \"\".join(line.strip() for line in charlist)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6a4da65",
   "metadata": {},
   "source": [
    "## Run Inference\n",
    "\n",
    "Now run inference. `compiled_model()` takes a list with input(s) in the same order as model input(s). Then we can fetch the output from output tensors.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e44a9de4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run inference on the model\n",
    "predictions = compiled_model([input_image])[recognition_output_layer]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "03a40d22",
   "metadata": {},
   "source": [
    "## Process Output Data\n",
    "\n",
    "The output of model format is W x B x L, where:\n",
    "\n",
    "* W - output sequence length\n",
    "* B - batch size\n",
    "* L - confidence distribution across the supported symbols in Kondate and Nakayosi.\n",
    "\n",
    "To get a more human-readable format, select a symbol with the highest probability. When you hold a list of indexes that are predicted to have the highest probability, due to limitations in [CTC Decoding](https://towardsdatascience.com/beam-search-decoding-in-ctc-trained-neural-networks-5a889a3d85a7), you will remove concurrent symbols and then remove the blanks.\n",
    "\n",
    "The last step is getting the symbols from corresponding indexes in the charlist."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da1b8bc5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Remove batch dimension\n",
    "predictions = np.squeeze(predictions)\n",
    "\n",
    "# Run argmax to pick the symbols with the highest probability\n",
    "predictions_indexes = np.argmax(predictions, axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6730159",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use groupby to remove concurrent letters, as required by CTC greedy decoding\n",
    "output_text_indexes = list(groupby(predictions_indexes))\n",
    "\n",
    "# Remove grouper objects\n",
    "output_text_indexes, _ = np.transpose(output_text_indexes, (1, 0))\n",
    "\n",
    "# Remove blank symbols\n",
    "output_text_indexes = output_text_indexes[output_text_indexes != 0]\n",
    "\n",
    "# Assign letters to indexes from output array\n",
    "output_text = [letters[letter_index] for letter_index in output_text_indexes]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac88f622",
   "metadata": {},
   "source": [
    "## Print Output\n",
    "\n",
    "Now you have a list of letters predicted by the model. The only thing left to do is display the image with predicted text printed below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f8a90f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(20, 1))\n",
    "plt.axis(\"off\")\n",
    "plt.imshow(resized_image, cmap=\"gray\", vmin=0, vmax=255)\n",
    "\n",
    "print(\"\".join(output_text))"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "ae617ccb002f72b3ab6d0069d721eac67ac2a969e83c083c4321cfcab0437cd1"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "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.9.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}