008/008-plate-ocr.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Пример распознавания автомобильных номеров"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[Building a Captcha OCR in TF2.0](https://www.kaggle.com/aakashnain/building-a-captcha-ocr-in-tf2-0)\n",
    "\n",
    "[Deep Learning OCR with Keras and Supervisely](https://hackernoon.com/latest-deep-learning-ocr-with-keras-and-supervisely-in-15-minutes-34aecd630ed8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.2.0\n"
     ]
    }
   ],
   "source": [
    "import cv2\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "\n",
    "print(tf.__version__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "PLATE_NUM_LENGTH = 8\n",
    "\n",
    "TEST_DATA_PATH = 'data/plate/test'\n",
    "TRAIN_DATA_PATH = 'data/plate/train'\n",
    "\n",
    "IMG_WIDTH = 128\n",
    "IMG_HEIGHT = 64\n",
    "\n",
    "POOL_NUM = 2\n",
    "DOWNSAMPLE = POOL_NUM ** 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Загружаем алфавит"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "from collections import Counter\n",
    "from glob import glob\n",
    "\n",
    "\n",
    "def load_char_counts(path):\n",
    "    \"\"\"Загружаем словарь символов на основе данных разметки\"\"\"\n",
    "    counts = Counter()\n",
    "    for filename in glob(os.path.join(path, 'ann', '*.json')):\n",
    "        with open(filename, 'r') as src:\n",
    "            annotation = json.load(src)\n",
    "        label = annotation['description']\n",
    "        assert len(label) == PLATE_NUM_LENGTH\n",
    "        counts.update(label)\n",
    "    return counts\n",
    "\n",
    "LETTERS = sorted(load_char_counts(TRAIN_DATA_PATH).keys())\n",
    "print('Letters:', ' '.join(LETTERS))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Утилиты подготовки входных данных"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "import functools\n",
    "\n",
    "def indices_to_text(indicies, letters):\n",
    "    \"\"\"Возвращает текст для заданных индексов\"\"\"\n",
    "    return ''.join(map(letters.__getitem__, indicies))\n",
    "\n",
    "def text_to_indices(text, letters):\n",
    "    \"\"\"Конвертирует текст в список индексов\"\"\"\n",
    "    return list(map(letters.index, text))\n",
    "\n",
    "class DataGenerator:\n",
    "    \"\"\"Загрузка и подготвка данных для обучения/тестирования\"\"\"\n",
    "    def __init__(self, \n",
    "                 path, \n",
    "                 img_width,\n",
    "                 img_height, \n",
    "                 batch_size,\n",
    "                 letters,\n",
    "                 max_label_len):\n",
    "        \n",
    "        self._img_height = img_height\n",
    "        self._img_width = img_width\n",
    "        self._batch_size = batch_size\n",
    "        self._letters = letters\n",
    "        self._max_label_len = max_label_len\n",
    "        \n",
    "        self._files = []\n",
    "        for img_path in glob(os.path.join(path, 'img', '*.png')):\n",
    "            basename, _ = os.path.splitext(os.path.basename(img_path))\n",
    "            annotation_path = os.path.join(path, 'ann', '%s.json' % basename)\n",
    "            with open(annotation_path, 'r') as src:\n",
    "                annotation = json.load(src)\n",
    "                self._files.append((img_path, annotation['description']))\n",
    "    \n",
    "    @property\n",
    "    def steps_per_epoch(self):\n",
    "        return len(self._files) // self._batch_size\n",
    "    \n",
    "    @functools.lru_cache(maxsize=1024)\n",
    "    def _load_img(self, path):\n",
    "        img = cv2.imread(path, cv2.IMREAD_GRAYSCALE)\n",
    "        img = cv2.resize(img, (self._img_width, self._img_height))\n",
    "        img = np.float32(img) / 255.\n",
    "        return img.reshape((self._img_height, self._img_width, 1))\n",
    "    \n",
    "    def __call__(self):\n",
    "        while True:\n",
    "            assert self.steps_per_epoch > 0\n",
    "    \n",
    "            batch_img = np.zeros(\n",
    "                [self._batch_size, \n",
    "                 self._img_height, self._img_width, 1], np.float32)\n",
    "            \n",
    "            random.shuffle(self._files)\n",
    "            for i in range(0, len(self._files), self._batch_size):\n",
    "                k = i + self._batch_size\n",
    "                if k > len(self._files):\n",
    "                    break\n",
    "\n",
    "                batch_labels = np.zeros([self._batch_size, \n",
    "                                         self._max_label_len], np.int32)\n",
    "            \n",
    "                batch_labels_length = np.zeros([self._batch_size, 1])\n",
    "                \n",
    "                for j, (img_path, label) in enumerate(self._files[i:k]):\n",
    "                    batch_img[j] = self._load_img(img_path)\n",
    "                    batch_labels[j] = text_to_indices(label, self._letters)\n",
    "                    batch_labels_length[j] = len(label)\n",
    "                                \n",
    "                outputs = {\n",
    "                    'labels': batch_labels,\n",
    "                    'labels_length': batch_labels_length\n",
    "                }\n",
    "\n",
    "                yield batch_img, outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "data = DataGenerator(TRAIN_DATA_PATH,\n",
    "                     batch_size=1,\n",
    "                     img_width=128,\n",
    "                     img_height=64,\n",
    "                     letters=LETTERS,\n",
    "                     max_label_len=PLATE_NUM_LENGTH)\n",
    "\n",
    "x, y_true = next(data())\n",
    "label_indices = y_true['labels'][0]\n",
    "label = indices_to_text(label_indices, LETTERS)\n",
    "label_length = y_true['labels_length'][0]\n",
    "\n",
    "print('Text generator output (data which will be fed into the neutral network):')\n",
    "print('1) input (image) %dx%d' % x[0].shape[:2])\n",
    "plt.imshow(np.squeeze(x[0]), cmap='gray')\n",
    "plt.show()\n",
    "\n",
    "print('2) labels (plate number): %s is encoded as %s' % (label, label_indices))\n",
    "print('3) label_length (length of plate number): %d' % label_length)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Архитектура сети и функция потерь"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from itertools import groupby\n",
    "from tensorflow.keras import backend as K\n",
    "\n",
    "class PlateOCR(tf.keras.Model):\n",
    "    \n",
    "    def __init__(self, input_shape,\n",
    "                 num_classes, name='plate_ocr'):\n",
    "        super(PlateOCR, self).__init__(name=name)\n",
    "        \n",
    "        img_height, img_width, img_channels = input_shape\n",
    "        \n",
    "        self._permute = tf.keras.layers.Permute(dims=(2, 1, 3),\n",
    "                                                input_shape=input_shape,\n",
    "                                                name='permute')\n",
    "        \n",
    "        self._conv1 = tf.keras.layers.Conv2D(16,\n",
    "                                             kernel_size=3,\n",
    "                                             padding='same',\n",
    "                                             activation='relu',\n",
    "                                             name='conv_1')\n",
    "        self._pool1 = tf.keras.layers.MaxPooling2D(pool_size=2, name='pool_1')\n",
    "        \n",
    "        self._conv2 = tf.keras.layers.Conv2D(16,\n",
    "                                             kernel_size=3,\n",
    "                                             padding='same',\n",
    "                                             activation='relu',\n",
    "                                             name='conv_2')\n",
    "        self._pool2 = tf.keras.layers.MaxPooling2D(pool_size=2, name='pool_2')\n",
    "        \n",
    "        self._rnn_input_length = img_width\n",
    "        self._rnn_input_length //= self._pool1.pool_size[0]\n",
    "        self._rnn_input_length //= self._pool2.pool_size[0]\n",
    "        \n",
    "        self._rnn_input_features = img_height\n",
    "        self._rnn_input_features //= self._pool1.pool_size[1]\n",
    "        self._rnn_input_features //= self._pool2.pool_size[1]\n",
    "        self._rnn_input_features *= self._conv2.filters\n",
    "        \n",
    "        self._reshape = tf.keras.layers.Reshape(\n",
    "            target_shape=(self._rnn_input_length,\n",
    "                          self._rnn_input_features),\n",
    "            name='reshape')\n",
    "        \n",
    "        self._dense = tf.keras.layers.Dense(32, \n",
    "                                            activation='relu',\n",
    "                                            name='dense_1')\n",
    "        \n",
    "        self._gru_1_fw = tf.keras.layers.GRU(512,\n",
    "                                             return_sequences=True,\n",
    "                                             name='gru_1_fw')\n",
    "        self._gru_1_bw = tf.keras.layers.GRU(512,\n",
    "                                             go_backwards=True,\n",
    "                                             return_sequences=True,\n",
    "                                             name='gru_1_bw')\n",
    "        \n",
    "        self._gru_2_fw = tf.keras.layers.GRU(512,\n",
    "                                             return_sequences=True,\n",
    "                                             name='gru_2_fw')\n",
    "        self._gru_2_bw = tf.keras.layers.GRU(512,\n",
    "                                             go_backwards=True,\n",
    "                                             return_sequences=True,\n",
    "                                             name='gru_2_bw')\n",
    "        \n",
    "        self._output = tf.keras.layers.Dense(num_classes + 1, \n",
    "                                             activation='softmax',\n",
    "                                             name='output')\n",
    "        \n",
    "        self.build(input_shape=(None, img_height, img_width, img_channels))\n",
    "\n",
    "    def call(self, x):\n",
    "        x = self._conv1(self._permute(x))\n",
    "        x = self._pool1(x)\n",
    "        \n",
    "        x = self._conv2(x)\n",
    "        x = self._pool2(x)\n",
    "        \n",
    "        x = self._dense(self._reshape(x))\n",
    "        \n",
    "        x = tf.keras.layers.add([self._gru_1_fw(x),\n",
    "                                 self._gru_1_bw(x)])\n",
    "        \n",
    "        x = tf.keras.layers.add([self._gru_2_fw(x),\n",
    "                                 self._gru_2_bw(x)])\n",
    "        \n",
    "        return self._output(x)\n",
    "\n",
    "    def train_step(self, data):\n",
    "        x, y_true = data\n",
    "        \n",
    "        labels = y_true['labels']\n",
    "        labels_length = y_true['labels_length']\n",
    "        \n",
    "        input_length = self._rnn_input_length * tf.ones_like(labels_length)\n",
    "        \n",
    "        with tf.GradientTape() as tape:\n",
    "            y_pred = self(x, training=True)\n",
    "            loss = K.ctc_batch_cost(labels,\n",
    "                                    y_pred,\n",
    "                                    input_length,\n",
    "                                    labels_length)\n",
    "\n",
    "        gradients = tape.gradient(\n",
    "            loss, self.trainable_variables)\n",
    "        self.optimizer.apply_gradients(\n",
    "            zip(gradients, self.trainable_variables))\n",
    "        \n",
    "        return {'loss': loss}\n",
    "    \n",
    "    def decode(self, y_pred, letters):\n",
    "        result = []\n",
    "        for prediction in y_pred:\n",
    "            indices = [i for i, _ in groupby(np.argmax(prediction, -1))]\n",
    "            indices = [i for i in indices if i < len(letters)]\n",
    "            result.append(indices_to_text(indices, letters))\n",
    "        return result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Обучение модели"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "model = PlateOCR(num_classes=len(LETTERS),\n",
    "                 input_shape=(IMG_HEIGHT, IMG_WIDTH, 1))\n",
    "\n",
    "model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "EPOCHS = 200\n",
    "BATCH_SIZE = 8\n",
    "\n",
    "data = DataGenerator(TRAIN_DATA_PATH,\n",
    "                     img_width=IMG_WIDTH,\n",
    "                     img_height=IMG_HEIGHT,\n",
    "                     batch_size=BATCH_SIZE,\n",
    "                     letters=LETTERS,\n",
    "                     max_label_len=PLATE_NUM_LENGTH)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.compile(optimizer=tf.keras.optimizers.Adam(lr=3e-4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "model.fit(data(),\n",
    "          epochs=EPOCHS,\n",
    "          steps_per_epoch=data.steps_per_epoch)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Загружаем обученную ранее модель"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.load_weights('data/plate/plate-ocr.hdf5')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Визуализация результата"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "import matplotlib.gridspec as gridspec\n",
    "\n",
    "TEST_BATCH_SIZE = 8\n",
    "\n",
    "test_data = DataGenerator(TEST_DATA_PATH,\n",
    "                          img_width=IMG_WIDTH,\n",
    "                          img_height=IMG_HEIGHT,\n",
    "                          batch_size=TEST_BATCH_SIZE,\n",
    "                          letters=LETTERS,\n",
    "                          max_label_len=PLATE_NUM_LENGTH)\n",
    "\n",
    "for batch_img, y_true in test_data():\n",
    "    y_pred = model.predict(batch_img)\n",
    "    labels_pred = model.decode(y_pred, LETTERS)\n",
    "    labels = [indices_to_text(i, LETTERS) for i in y_true['labels']]\n",
    "    \n",
    "    for i in range(TEST_BATCH_SIZE):\n",
    "        fig = plt.figure(figsize=(10, 10))\n",
    "        outer = gridspec.GridSpec(2, 1, wspace=10, hspace=0.1)\n",
    "        ax1 = plt.Subplot(fig, outer[0])\n",
    "        fig.add_subplot(ax1)\n",
    "        ax2 = plt.Subplot(fig, outer[1])\n",
    "        fig.add_subplot(ax2)\n",
    "        print('Predicted: %s\\nExpected: %s' % (labels_pred[i], labels[i]))\n",
    "        ax1.set_title('Input img')\n",
    "        ax1.imshow(np.squeeze(batch_img[i]), cmap='gray')\n",
    "        ax1.set_xticks([])\n",
    "        ax1.set_yticks([])\n",
    "        ax2.set_title('Acrtivations')\n",
    "        ax2.imshow(y_pred[i].T, cmap='binary', interpolation='nearest')\n",
    "        ax2.set_yticks(list(range(len(LETTERS) + 1)))\n",
    "        ax2.set_yticklabels(LETTERS + ['blank'])\n",
    "        ax2.grid(False)\n",
    "        for h in np.arange(-0.5, len(LETTERS) + 1 + 0.5, 1):\n",
    "            ax2.axhline(h, linestyle='-', color='k', alpha=0.5, linewidth=1)\n",
    "        \n",
    "        plt.show()\n",
    "    break"
   ]
  }
 ],
 "metadata": {
  "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}