tf2image_ocr.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "tf2image_ocr.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "authorship_tag": "ABX9TyPJNOhDWOLMBfJx5J7BUZ3F",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/menon92/Tensorflow2ImageOCR/blob/main/tf2image_ocr.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "E6dYrB38sCIb",
        "outputId": "b73159b5-0c98-4b9d-ff50-8d4499f685d6"
      },
      "source": [
        "!pip install -qq cairocffi editdistance\n",
        "!apt install -qq graphviz\n",
        "!pip install -qq pydot\n",
        "!pip install -qq matplotlib graphviz pydot"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[?25l\r\u001b[K     |████▋                           | 10 kB 27.3 MB/s eta 0:00:01\r\u001b[K     |█████████▎                      | 20 kB 25.0 MB/s eta 0:00:01\r\u001b[K     |██████████████                  | 30 kB 13.0 MB/s eta 0:00:01\r\u001b[K     |██████████████████▋             | 40 kB 10.5 MB/s eta 0:00:01\r\u001b[K     |███████████████████████▎        | 51 kB 5.5 MB/s eta 0:00:01\r\u001b[K     |████████████████████████████    | 61 kB 5.9 MB/s eta 0:00:01\r\u001b[K     |████████████████████████████████| 70 kB 3.5 MB/s \n",
            "\u001b[?25h  Building wheel for cairocffi (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "graphviz is already the newest version (2.40.1-2).\n",
            "0 upgraded, 0 newly installed, 0 to remove and 37 not upgraded.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VCq0-o1PsHoS"
      },
      "source": [
        "import os\n",
        "import itertools\n",
        "import codecs\n",
        "import re\n",
        "import datetime\n",
        "import cairocffi as cairo\n",
        "import editdistance\n",
        "import numpy as np\n",
        "from scipy import ndimage\n",
        "import matplotlib.pyplot as plt\n",
        "import pylab\n",
        "\n",
        "from tensorflow.keras import backend as K\n",
        "from tensorflow.keras.layers import Conv2D, MaxPooling2D\n",
        "from tensorflow.keras.layers import Input, Dense, Activation\n",
        "from tensorflow.keras.layers import Reshape, Lambda\n",
        "from tensorflow.keras.layers import add, concatenate, Bidirectional\n",
        "from tensorflow.keras.models import Model\n",
        "from tensorflow.keras.layers import GRU\n",
        "from tensorflow.keras.optimizers import SGD\n",
        "from tensorflow.keras.utils import get_file\n",
        "from tensorflow.keras.preprocessing import image\n",
        "from tensorflow.keras import layers\n",
        "import tensorflow.keras.callbacks\n",
        "from tensorflow import keras\n",
        "import tensorflow as tf"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zGj1Kl6TsHq6"
      },
      "source": [
        "OUTPUT_DIR = 'image_ocr'\n",
        "\n",
        "# character classes and matching regex filter\n",
        "regex = r'^[a-z ]+$'\n",
        "alphabet = u'abcdefghijklmnopqrstuvwxyz '\n",
        "\n",
        "np.random.seed(55)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5vzpOZM1sHsX"
      },
      "source": [
        "# this creates larger \"blotches\" of noise which look\n",
        "# more realistic than just adding gaussian noise\n",
        "# assumes greyscale with pixels ranging from 0 to 1\n",
        "\n",
        "def speckle(img):\n",
        "    severity = np.random.uniform(0, 0.6)\n",
        "    blur = ndimage.gaussian_filter(np.random.randn(*img.shape) * severity, 1)\n",
        "    img_speck = (img + blur)\n",
        "    img_speck[img_speck > 1] = 1\n",
        "    img_speck[img_speck <= 0] = 0\n",
        "    return img_speck\n",
        "\n",
        "\n",
        "# paints the string in a random location the bounding box\n",
        "# also uses a random font, a slight random rotation,\n",
        "# and a random amount of speckle noise\n",
        "\n",
        "def paint_text(text, w, h, rotate=False, ud=False, multi_fonts=False):\n",
        "    surface = cairo.ImageSurface(cairo.FORMAT_RGB24, w, h)\n",
        "    with cairo.Context(surface) as context:\n",
        "        context.set_source_rgb(1, 1, 1)  # White\n",
        "        context.paint()\n",
        "        # this font list works in CentOS 7\n",
        "        if multi_fonts:\n",
        "            fonts = ['Century Schoolbook', 'Courier', 'STIX', 'URW Chancery L', 'FreeMono']\n",
        "            context.select_font_face(np.random.choice(fonts), cairo.FONT_SLANT_NORMAL,\n",
        "                                     np.random.choice([cairo.FONT_WEIGHT_BOLD, cairo.FONT_WEIGHT_NORMAL]))\n",
        "        else:\n",
        "            context.select_font_face('Courier', cairo.FONT_SLANT_NORMAL, cairo.FONT_WEIGHT_BOLD)\n",
        "        context.set_font_size(25)\n",
        "        box = context.text_extents(text)\n",
        "        border_w_h = (4, 4)\n",
        "        if box[2] > (w - 2 * border_w_h[1]) or box[3] > (h - 2 * border_w_h[0]):\n",
        "            raise IOError('Could not fit string into image. Max char count is too large for given image width.')\n",
        "\n",
        "        # teach the RNN translational invariance by\n",
        "        # fitting text box randomly on canvas, with some room to rotate\n",
        "        max_shift_x = w - box[2] - border_w_h[0]\n",
        "        max_shift_y = h - box[3] - border_w_h[1]\n",
        "        top_left_x = np.random.randint(0, int(max_shift_x))\n",
        "        if ud:\n",
        "            top_left_y = np.random.randint(0, int(max_shift_y))\n",
        "        else:\n",
        "            top_left_y = h // 2\n",
        "        context.move_to(top_left_x - int(box[0]), top_left_y - int(box[1]))\n",
        "        context.set_source_rgb(0, 0, 0)\n",
        "        context.show_text(text)\n",
        "\n",
        "    buf = surface.get_data()\n",
        "    a = np.frombuffer(buf, np.uint8)\n",
        "    a.shape = (h, w, 4)\n",
        "    a = a[:, :, 0]  # grab single channel\n",
        "    a = a.astype(np.float32) / 255\n",
        "    a = np.expand_dims(a, 0)\n",
        "    if rotate:\n",
        "        a = image.random_rotation(a, 3 * (w - top_left_x) / w + 1)\n",
        "    a = speckle(a)\n",
        "\n",
        "    return a\n",
        "\n",
        "\n",
        "def shuffle_mats_or_lists(matrix_list, stop_ind=None):\n",
        "    ret = []\n",
        "    assert all([len(i) == len(matrix_list[0]) for i in matrix_list])\n",
        "    len_val = len(matrix_list[0])\n",
        "    if stop_ind is None:\n",
        "        stop_ind = len_val\n",
        "    assert stop_ind <= len_val\n",
        "\n",
        "    a = list(range(stop_ind))\n",
        "    np.random.shuffle(a)\n",
        "    a += list(range(stop_ind, len_val))\n",
        "    for mat in matrix_list:\n",
        "        if isinstance(mat, np.ndarray):\n",
        "            ret.append(mat[a])\n",
        "        elif isinstance(mat, list):\n",
        "            ret.append([mat[i] for i in a])\n",
        "        else:\n",
        "            raise TypeError('`shuffle_mats_or_lists` only supports '\n",
        "                            'numpy.array and list objects.')\n",
        "    return ret\n",
        "\n",
        "\n",
        "# Translation of characters to unique integer values\n",
        "def text_to_labels(text):\n",
        "    ret = []\n",
        "    for char in text:\n",
        "        ret.append(alphabet.find(char))\n",
        "    return ret\n",
        "\n",
        "\n",
        "# Reverse translation of numerical classes back to characters\n",
        "def labels_to_text(labels):\n",
        "    ret = []\n",
        "    for c in labels:\n",
        "        if c == len(alphabet):  # CTC Blank\n",
        "            ret.append(\"\")\n",
        "        else:\n",
        "            ret.append(alphabet[c])\n",
        "    return \"\".join(ret)\n",
        "\n",
        "\n",
        "# only a-z and space..probably not to difficult\n",
        "# to expand to uppercase and symbols\n",
        "\n",
        "def is_valid_str(in_str):\n",
        "    search = re.compile(regex, re.UNICODE).search\n",
        "    return bool(search(in_str))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xFR3rJJWsHu2"
      },
      "source": [
        "# Uses generator functions to supply train/test with\n",
        "# data. Image renderings are text are created on the fly\n",
        "# each time with random perturbations\n",
        "\n",
        "class TextImageGenerator(tensorflow.keras.callbacks.Callback):\n",
        "\n",
        "    def __init__(self, monogram_file, bigram_file, minibatch_size,\n",
        "                 img_w, img_h, downsample_factor, val_split,\n",
        "                 absolute_max_string_len=16):\n",
        "\n",
        "        self.minibatch_size = minibatch_size\n",
        "        self.img_w = img_w\n",
        "        self.img_h = img_h\n",
        "        self.monogram_file = monogram_file\n",
        "        self.bigram_file = bigram_file\n",
        "        self.downsample_factor = downsample_factor\n",
        "        self.val_split = val_split\n",
        "        self.blank_label = self.get_output_size() - 1\n",
        "        self.absolute_max_string_len = absolute_max_string_len\n",
        "\n",
        "        self.current_epoch = 0\n",
        "        self.cur_train_index = 0\n",
        "\n",
        "    def get_output_size(self):\n",
        "        return len(alphabet) + 1\n",
        "\n",
        "    # num_words can be independent of the epoch size due to the use of generators\n",
        "    # as max_string_len grows, num_words can grow\n",
        "    def build_word_list(self, num_words, max_string_len=None, mono_fraction=0.5):\n",
        "        assert max_string_len <= self.absolute_max_string_len\n",
        "        assert num_words % self.minibatch_size == 0\n",
        "        assert (self.val_split * num_words) % self.minibatch_size == 0\n",
        "        self.num_words = num_words\n",
        "        self.string_list = [''] * self.num_words\n",
        "        tmp_string_list = []\n",
        "        self.max_string_len = max_string_len\n",
        "        self.Y_data = np.ones([self.num_words, self.absolute_max_string_len]) * -1\n",
        "        self.X_text = []\n",
        "        self.Y_len = [0] * self.num_words\n",
        "\n",
        "        # monogram file is sorted by frequency in english speech\n",
        "        with codecs.open(self.monogram_file, mode='r', encoding='utf-8') as f:\n",
        "            for line in f:\n",
        "                if len(tmp_string_list) == int(self.num_words * mono_fraction):\n",
        "                    break\n",
        "                word = line.rstrip()\n",
        "                if max_string_len == -1 or max_string_len is None or len(word) <= max_string_len:\n",
        "                    tmp_string_list.append(word)\n",
        "\n",
        "        # bigram file contains common word pairings in english speech\n",
        "        with codecs.open(self.bigram_file, mode='r', encoding='utf-8') as f:\n",
        "            lines = f.readlines()\n",
        "            for line in lines:\n",
        "                if len(tmp_string_list) == self.num_words:\n",
        "                    break\n",
        "                columns = line.lower().split()\n",
        "                word = columns[0] + ' ' + columns[1]\n",
        "                if is_valid_str(word) and \\\n",
        "                        (max_string_len == -1 or max_string_len is None or len(word) <= max_string_len):\n",
        "                    tmp_string_list.append(word)\n",
        "        if len(tmp_string_list) != self.num_words:\n",
        "            raise IOError('Could not pull enough words from supplied monogram and bigram files. ')\n",
        "        # interlace to mix up the easy and hard words\n",
        "        self.string_list[::2] = tmp_string_list[:self.num_words // 2]\n",
        "        self.string_list[1::2] = tmp_string_list[self.num_words // 2:]\n",
        "\n",
        "        for i, word in enumerate(self.string_list):\n",
        "            self.Y_len[i] = len(word)\n",
        "            self.Y_data[i, 0:len(word)] = text_to_labels(word)\n",
        "            self.X_text.append(word)\n",
        "        self.Y_len = np.expand_dims(np.array(self.Y_len), 1)\n",
        "\n",
        "        self.cur_val_index = self.val_split\n",
        "        self.cur_train_index = 0\n",
        "\n",
        "    # each time an image is requested from train/val/test, a new random\n",
        "    # painting of the text is performed\n",
        "    def get_batch(self, index, size, train):\n",
        "\n",
        "        if self.current_epoch == 0:\n",
        "            # print(\"get batch 1st epoch\")\n",
        "            self.build_word_list(16000, 4, 1)\n",
        "            self.paint_func = lambda text: paint_text(text, self.img_w, self.img_h,\n",
        "                                                    rotate=False, ud=False, multi_fonts=False)\n",
        "            self.current_epoch += 1\n",
        "        # rebind the paint function to implement curriculum learning\n",
        "        if 3 <= self.current_epoch < 6:\n",
        "            self.paint_func = lambda text: paint_text(text, self.img_w, self.img_h,\n",
        "                                                      rotate=False, ud=True, multi_fonts=False)\n",
        "        elif 6 <= self.current_epoch < 9:\n",
        "            self.paint_func = lambda text: paint_text(text, self.img_w, self.img_h,\n",
        "                                                      rotate=False, ud=True, multi_fonts=True)\n",
        "        elif self.current_epoch >= 9:\n",
        "            self.paint_func = lambda text: paint_text(text, self.img_w, self.img_h,\n",
        "                                                      rotate=True, ud=True, multi_fonts=True)\n",
        "        if self.current_epoch >= 21 and self.max_string_len < 12:\n",
        "            self.build_word_list(32000, 12, 0.5)\n",
        "\n",
        "        # width and height are backwards from typical Keras convention\n",
        "        # because width is the time dimension when it gets fed into the RNN\n",
        "        if K.image_data_format() == 'channels_first':\n",
        "            X_data = np.ones([size, 1, self.img_w, self.img_h])\n",
        "        else:\n",
        "            X_data = np.ones([size, self.img_w, self.img_h, 1])\n",
        "\n",
        "        labels = np.ones([size, self.absolute_max_string_len]) * -1\n",
        "        input_length = np.zeros([size, 1])\n",
        "        label_length = np.zeros([size, 1])\n",
        "        source_str = []\n",
        "        for i in range(size):\n",
        "            # Mix in some blank inputs.  This seems to be important for\n",
        "            # achieving translational invariance\n",
        "            if train and i > size - 4:\n",
        "                if K.image_data_format() == 'channels_first':\n",
        "                    X_data[i, 0, 0:self.img_w, :] = self.paint_func('')[0, :, :].T\n",
        "                else:\n",
        "                    X_data[i, 0:self.img_w, :, 0] = self.paint_func('',)[0, :, :].T\n",
        "                labels[i, 0] = self.blank_label\n",
        "                input_length[i] = self.img_w // self.downsample_factor - 2\n",
        "                label_length[i] = 1\n",
        "                source_str.append('')\n",
        "            else:\n",
        "                if K.image_data_format() == 'channels_first':\n",
        "                    X_data[i, 0, 0:self.img_w, :] = self.paint_func(self.X_text[index + i])[0, :, :].T\n",
        "                else:\n",
        "                    X_data[i, 0:self.img_w, :, 0] = self.paint_func(self.X_text[index + i])[0, :, :].T\n",
        "                labels[i, :] = self.Y_data[index + i]\n",
        "                input_length[i] = self.img_w // self.downsample_factor - 2\n",
        "                label_length[i] = self.Y_len[index + i]\n",
        "                source_str.append(self.X_text[index + i])\n",
        "        inputs = {'the_input': X_data,\n",
        "                  'the_labels': labels,\n",
        "                  'input_length': input_length,\n",
        "                  'label_length': label_length,\n",
        "                  'source_str': np.array(source_str)  # used for visualization only\n",
        "                  }\n",
        "        outputs = {'ctc': np.zeros([size])}  # dummy data for dummy loss function\n",
        "        return (inputs, outputs)\n",
        "\n",
        "    def next_train(self):\n",
        "        while 1:\n",
        "            ret = self.get_batch(self.cur_train_index, self.minibatch_size, train=True)\n",
        "            self.cur_train_index += self.minibatch_size\n",
        "            if self.cur_train_index >= self.val_split:\n",
        "                self.cur_train_index = self.cur_train_index % 32\n",
        "                (self.X_text, self.Y_data, self.Y_len) = shuffle_mats_or_lists(\n",
        "                    [self.X_text, self.Y_data, self.Y_len], self.val_split)\n",
        "            yield ret\n",
        "\n",
        "    def next_val(self):\n",
        "        while 1:\n",
        "            ret = self.get_batch(self.cur_val_index, self.minibatch_size, train=False)\n",
        "            self.cur_val_index += self.minibatch_size\n",
        "            if self.cur_val_index >= self.num_words:\n",
        "                self.cur_val_index = self.val_split + self.cur_val_index % 32\n",
        "            yield ret\n",
        "\n",
        "    def on_train_begin(self, logs={}):\n",
        "        pass\n",
        "        # print(\"On train begin ...\")\n",
        "        # self.build_word_list(16000, 4, 1)\n",
        "        # self.paint_func = lambda text: paint_text(text, self.img_w, self.img_h,\n",
        "        #                                           rotate=False, ud=False, multi_fonts=False)\n",
        "\n",
        "    def on_epoch_begin(self, epoch, logs={}):\n",
        "        self.current_epoch = epoch"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "P3ANvevhsxPK"
      },
      "source": [
        "# the actual loss calc occurs here despite it not being\n",
        "# an internal Keras loss function\n",
        "\n",
        "def ctc_lambda_func(args):\n",
        "    y_pred, labels, input_length, label_length = args\n",
        "    # the 2 is critical here since the first couple outputs of the RNN\n",
        "    # tend to be garbage:\n",
        "    y_pred = y_pred[:, 2:, :]\n",
        "    return K.ctc_batch_cost(labels, y_pred, input_length, label_length)\n",
        "\n",
        "\n",
        "# For a real OCR application, this should be beam search with a dictionary\n",
        "# and language model.  For this example, best path is sufficient.\n",
        "\n",
        "def decode_batch(test_func, word_batch):\n",
        "    out = test_func([word_batch])[0]\n",
        "    ret = []\n",
        "    for j in range(out.shape[0]):\n",
        "        out_best = list(np.argmax(out[j, 2:], 1))\n",
        "        out_best = [k for k, g in itertools.groupby(out_best)]\n",
        "        outstr = labels_to_text(out_best)\n",
        "        ret.append(outstr)\n",
        "    return ret"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gRHDtadosyRo"
      },
      "source": [
        "class VizCallback(tensorflow.keras.callbacks.Callback):\n",
        "\n",
        "    def __init__(self, run_name, test_func, text_img_gen, num_display_words=6):\n",
        "        self.test_func = test_func\n",
        "        self.output_dir = os.path.join(\n",
        "            OUTPUT_DIR, run_name)\n",
        "        self.text_img_gen = text_img_gen\n",
        "        self.num_display_words = num_display_words\n",
        "        if not os.path.exists(self.output_dir):\n",
        "            os.makedirs(self.output_dir)\n",
        "\n",
        "    def show_edit_distance(self, num):\n",
        "        num_left = num\n",
        "        mean_norm_ed = 0.0\n",
        "        mean_ed = 0.0\n",
        "        while num_left > 0:\n",
        "            word_batch = next(self.text_img_gen)[0]\n",
        "            num_proc = min(word_batch['the_input'].shape[0], num_left)\n",
        "            decoded_res = decode_batch(self.test_func, word_batch['the_input'][0:num_proc])\n",
        "            for j in range(num_proc):\n",
        "                edit_dist = editdistance.eval(decoded_res[j], word_batch['source_str'][j])\n",
        "                mean_ed += float(edit_dist)\n",
        "                mean_norm_ed += float(edit_dist) / len(word_batch['source_str'][j])\n",
        "            num_left -= num_proc\n",
        "        mean_norm_ed = mean_norm_ed / num\n",
        "        mean_ed = mean_ed / num\n",
        "        print('\\nOut of %d samples:  Mean edit distance: %.3f Mean normalized edit distance: %0.3f'\n",
        "              % (num, mean_ed, mean_norm_ed))\n",
        "\n",
        "    def on_epoch_end(self, epoch, logs={}):\n",
        "        self.model.save_weights(os.path.join(self.output_dir, 'weights%02d.h5' % (epoch)))\n",
        "        self.show_edit_distance(256)\n",
        "        word_batch = next(self.text_img_gen)[0]\n",
        "        res = decode_batch(self.test_func, word_batch['the_input'][0:self.num_display_words])\n",
        "        if word_batch['the_input'][0].shape[0] < 256:\n",
        "            cols = 2\n",
        "        else:\n",
        "            cols = 1\n",
        "        for i in range(self.num_display_words):\n",
        "            plt.subplot(self.num_display_words // cols, cols, i + 1)\n",
        "            if K.image_data_format() == 'channels_first':\n",
        "                the_input = word_batch['the_input'][i, 0, :, :]\n",
        "            else:\n",
        "                the_input = word_batch['the_input'][i, :, :, 0]\n",
        "            plt.imshow(the_input.T, cmap='Greys_r')\n",
        "            plt.xlabel('Truth = \\'%s\\'\\nDecoded = \\'%s\\'' % (word_batch['source_str'][i], res[i]))\n",
        "        fig = pylab.gcf()\n",
        "        fig.set_size_inches(10, 13)\n",
        "        plt.savefig(os.path.join(self.output_dir, 'e%02d.png' % (epoch)))\n",
        "        plt.close()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eaRrE_u-s2p_"
      },
      "source": [
        "def train(run_name, start_epoch, stop_epoch, img_w):\n",
        "    # Input Parameters\n",
        "    img_h = 64\n",
        "    words_per_epoch = 16000\n",
        "    val_split = 0.2\n",
        "    val_words = int(words_per_epoch * (val_split))\n",
        "\n",
        "    # Network parameters\n",
        "    conv_filters = 16\n",
        "    kernel_size = (3, 3)\n",
        "    pool_size = 2\n",
        "    time_dense_size = 32\n",
        "    rnn_size = 512\n",
        "    minibatch_size = 32\n",
        "\n",
        "    if K.image_data_format() == 'channels_first':\n",
        "        input_shape = (1, img_w, img_h)\n",
        "    else:\n",
        "        input_shape = (img_w, img_h, 1)\n",
        "\n",
        "    fdir = os.path.dirname(get_file('wordlists.tgz',\n",
        "                                    origin='http://pasky.or.cz/wordlists.tgz', untar=True))\n",
        "\n",
        "    img_gen = TextImageGenerator(monogram_file=os.path.join(fdir, 'wordlist_mono_clean.txt'),\n",
        "                                 bigram_file=os.path.join(fdir, 'wordlist_bi_clean.txt'),\n",
        "                                 minibatch_size=minibatch_size,\n",
        "                                 img_w=img_w,\n",
        "                                 img_h=img_h,\n",
        "                                 downsample_factor=(pool_size ** 2),\n",
        "                                 val_split=words_per_epoch - val_words\n",
        "                                 )\n",
        "    act = 'relu'\n",
        "    input_data = Input(name='the_input', shape=input_shape, dtype='float32')\n",
        "    inner = Conv2D(conv_filters, kernel_size, padding='same',\n",
        "                   activation=act, kernel_initializer='he_normal',\n",
        "                   name='conv1')(input_data)\n",
        "    inner = MaxPooling2D(pool_size=(pool_size, pool_size), name='max1')(inner)\n",
        "    inner = Conv2D(conv_filters, kernel_size, padding='same',\n",
        "                   activation=act, kernel_initializer='he_normal',\n",
        "                   name='conv2')(inner)\n",
        "    inner = MaxPooling2D(pool_size=(pool_size, pool_size), name='max2')(inner)\n",
        "\n",
        "    conv_to_rnn_dims = (img_w // (pool_size ** 2), (img_h // (pool_size ** 2)) * conv_filters)\n",
        "    inner = Reshape(target_shape=conv_to_rnn_dims, name='reshape')(inner)\n",
        "\n",
        "    # cuts down input size going into RNN:\n",
        "    inner = Dense(time_dense_size, activation=act, name='dense1')(inner)\n",
        "\n",
        "    # --------------------------------\n",
        "    # Two layers of bidirecitonal GRUs\n",
        "    # GRU seems to work as well, if not better than LSTM:\n",
        "    gru_1 = Bidirectional(GRU(rnn_size, return_sequences=True, name='gru1'), merge_mode='sum')(inner)\n",
        "    gru_2 = Bidirectional(GRU(rnn_size, return_sequences=True, name='gru2'), merge_mode='concat')(gru_1)\n",
        "    inner = Dense(img_gen.get_output_size(), name='dense2')(gru_2)\n",
        "    # ---------------------------------\n",
        "\n",
        "    # ---------------------------------\n",
        "    # Two layers of bidirectional GRUs\n",
        "    # GRU seems to work as well, if not better than LSTM:\n",
        "    # gru_1 = GRU(rnn_size, return_sequences=True, name='gru1')(inner)\n",
        "    # gru_1b = GRU(rnn_size, return_sequences=True, go_backwards=True, name='gru1_b')(inner)\n",
        "    # gru1_merged = add([gru_1, gru_1b])\n",
        "    # gru_2 = GRU(rnn_size, return_sequences=True, name='gru2')(gru1_merged)\n",
        "    # gru_2b = GRU(rnn_size, return_sequences=True, go_backwards=True, name='gru2_b')(gru1_merged)\n",
        "\n",
        "    # # transforms RNN output to character activations:\n",
        "    # inner = Dense(img_gen.get_output_size(), kernel_initializer='he_normal',\n",
        "    #               name='dense2')(concatenate([gru_2, gru_2b]))\n",
        "    # --------------------------------- \n",
        "    \n",
        "    y_pred = Activation('softmax', name='softmax')(inner)\n",
        "    # Model(inputs=input_data, outputs=y_pred).summary()\n",
        "\n",
        "    labels = Input(name='the_labels', shape=[img_gen.absolute_max_string_len], dtype='float32')\n",
        "    input_length = Input(name='input_length', shape=[1], dtype='int64')\n",
        "    label_length = Input(name='label_length', shape=[1], dtype='int64')\n",
        "    # Keras doesn't currently support loss funcs with extra parameters\n",
        "    # so CTC loss is implemented in a lambda layer\n",
        "    loss_out = Lambda(ctc_lambda_func, output_shape=(1,), name='ctc')([y_pred, labels, input_length, label_length])\n",
        "\n",
        "    # clipnorm seems to speeds up convergence\n",
        "    sgd = SGD(learning_rate=0.02, decay=1e-6, momentum=0.9, nesterov=True, clipnorm=5)\n",
        "\n",
        "    model = Model(inputs=[input_data, labels, input_length, label_length], outputs=loss_out)\n",
        "\n",
        "    # the loss calc occurs elsewhere, so use a dummy lambda func for the loss\n",
        "    model.compile(loss={'ctc': lambda y_true, y_pred: y_pred}, optimizer=sgd)\n",
        "    if start_epoch > 0:\n",
        "        weight_file = os.path.join(OUTPUT_DIR, os.path.join(run_name, 'weights%02d.h5' % (start_epoch - 1)))\n",
        "        model.load_weights(weight_file)\n",
        "    # captures output of softmax so we can decode the output during visualization\n",
        "    test_func = K.function([input_data], [y_pred])\n",
        "\n",
        "    viz_cb = VizCallback(run_name, test_func, img_gen.next_val())\n",
        "\n",
        "    model.fit(\n",
        "        x=img_gen.next_train(),\n",
        "        steps_per_epoch=(words_per_epoch - val_words) // minibatch_size,\n",
        "        epochs=stop_epoch,\n",
        "        validation_data=img_gen.next_val(),\n",
        "        validation_steps=val_words // minibatch_size,\n",
        "        callbacks=[viz_cb, img_gen],\n",
        "        initial_epoch=start_epoch\n",
        "    )\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rH2NdvPus8dc",
        "outputId": "92d19909-11bf-4735-9bac-5ec4a10c93d2"
      },
      "source": [
        "run_name = datetime.datetime.now().strftime('%Y:%m:%d:%H:%M:%S')\n",
        "train(run_name, 0, 20, 128)\n",
        "# increase to wider images and start at epoch 20. The learned weights are reloaded\n",
        "train(run_name, 20, 25, 512)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 1/20\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.7/dist-packages/keras/engine/functional.py:585: UserWarning: Input dict contained keys ['source_str'] which did not match any model input. They will be ignored by the model.\n",
            "  [n for n in tensors.keys() if n not in ref_input_names])\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "400/400 [==============================] - 73s 164ms/step - loss: 6.9904 - val_loss: 0.0559\n",
            "\n",
            "Out of 256 samples:  Mean edit distance: 0.020 Mean normalized edit distance: 0.006\n",
            "Epoch 2/20\n",
            "400/400 [==============================] - 64s 161ms/step - loss: -0.0206 - val_loss: 0.0219\n",
            "\n",
            "Out of 256 samples:  Mean edit distance: 0.004 Mean normalized edit distance: 0.001\n",
            "Epoch 3/20\n",
            "400/400 [==============================] - 63s 157ms/step - loss: -0.0501 - val_loss: 0.0104\n",
            "\n",
            "Out of 256 samples:  Mean edit distance: 0.008 Mean normalized edit distance: 0.002\n",
            "Epoch 4/20\n",
            "400/400 [==============================] - 63s 159ms/step - loss: 2.0468 - val_loss: 0.0770\n",
            "\n",
            "Out of 256 samples:  Mean edit distance: 0.020 Mean normalized edit distance: 0.005\n",
            "Epoch 5/20\n",
            "400/400 [==============================] - 63s 157ms/step - loss: -7.4291e-04 - val_loss: 0.0667\n",
            "\n",
            "Out of 256 samples:  Mean edit distance: 0.051 Mean normalized edit distance: 0.014\n",
            "Epoch 6/20\n",
            "400/400 [==============================] - 64s 160ms/step - loss: -0.0308 - val_loss: 0.0178\n",
            "\n",
            "Out of 256 samples:  Mean edit distance: 0.008 Mean normalized edit distance: 0.002\n",
            "Epoch 7/20\n",
            "400/400 [==============================] - 64s 160ms/step - loss: 0.0478 - val_loss: 0.0313\n",
            "\n",
            "Out of 256 samples:  Mean edit distance: 0.008 Mean normalized edit distance: 0.002\n",
            "Epoch 8/20\n",
            "400/400 [==============================] - 64s 161ms/step - loss: -0.0386 - val_loss: 0.0279\n",
            "\n",
            "Out of 256 samples:  Mean edit distance: 0.020 Mean normalized edit distance: 0.006\n",
            "Epoch 9/20\n",
            "400/400 [==============================] - 64s 161ms/step - loss: -0.0440 - val_loss: 0.0288\n",
            "\n",
            "Out of 256 samples:  Mean edit distance: 0.004 Mean normalized edit distance: 0.001\n",
            "Epoch 10/20\n",
            "400/400 [==============================] - 67s 168ms/step - loss: 0.1026 - val_loss: 0.1475\n",
            "\n",
            "Out of 256 samples:  Mean edit distance: 0.043 Mean normalized edit distance: 0.013\n",
            "Epoch 11/20\n",
            "400/400 [==============================] - 65s 163ms/step - loss: -3.0416e-04 - val_loss: 0.1001\n",
            "\n",
            "Out of 256 samples:  Mean edit distance: 0.031 Mean normalized edit distance: 0.009\n",
            "Epoch 12/20\n",
            "400/400 [==============================] - 65s 163ms/step - loss: -0.0033 - val_loss: 0.0933\n",
            "\n",
            "Out of 256 samples:  Mean edit distance: 0.031 Mean normalized edit distance: 0.010\n",
            "Epoch 13/20\n",
            "400/400 [==============================] - 64s 161ms/step - loss: -0.0080 - val_loss: 0.0680\n",
            "\n",
            "Out of 256 samples:  Mean edit distance: 0.023 Mean normalized edit distance: 0.007\n",
            "Epoch 14/20\n",
            "400/400 [==============================] - 65s 164ms/step - loss: -0.0308 - val_loss: 0.0777\n",
            "\n",
            "Out of 256 samples:  Mean edit distance: 0.016 Mean normalized edit distance: 0.004\n",
            "Epoch 15/20\n",
            "400/400 [==============================] - 65s 163ms/step - loss: -0.0323 - val_loss: 0.0418\n",
            "\n",
            "Out of 256 samples:  Mean edit distance: 0.023 Mean normalized edit distance: 0.007\n",
            "Epoch 16/20\n",
            "400/400 [==============================] - 64s 160ms/step - loss: -0.0383 - val_loss: 0.0539\n",
            "\n",
            "Out of 256 samples:  Mean edit distance: 0.016 Mean normalized edit distance: 0.004\n",
            "Epoch 17/20\n",
            "400/400 [==============================] - 65s 164ms/step - loss: -0.0380 - val_loss: 0.0299\n",
            "\n",
            "Out of 256 samples:  Mean edit distance: 0.016 Mean normalized edit distance: 0.004\n",
            "Epoch 18/20\n",
            "400/400 [==============================] - 65s 162ms/step - loss: -0.0401 - val_loss: 0.0219\n",
            "\n",
            "Out of 256 samples:  Mean edit distance: 0.004 Mean normalized edit distance: 0.001\n",
            "Epoch 19/20\n",
            "400/400 [==============================] - 65s 162ms/step - loss: -0.0366 - val_loss: 0.0229\n",
            "\n",
            "Out of 256 samples:  Mean edit distance: 0.012 Mean normalized edit distance: 0.004\n",
            "Epoch 20/20\n",
            "400/400 [==============================] - 63s 158ms/step - loss: -0.0456 - val_loss: 0.0191\n",
            "\n",
            "Out of 256 samples:  Mean edit distance: 0.000 Mean normalized edit distance: 0.000\n",
            "Epoch 21/25\n",
            "400/400 [==============================] - 146s 348ms/step - loss: 0.1706 - val_loss: 0.1948\n",
            "\n",
            "Out of 256 samples:  Mean edit distance: 0.035 Mean normalized edit distance: 0.010\n",
            "Epoch 22/25\n",
            "400/400 [==============================] - 140s 351ms/step - loss: 0.3167 - val_loss: 0.1518\n",
            "\n",
            "Out of 256 samples:  Mean edit distance: 0.027 Mean normalized edit distance: 0.005\n",
            "Epoch 23/25\n",
            "400/400 [==============================] - 141s 351ms/step - loss: 0.0804 - val_loss: 0.1123\n",
            "\n",
            "Out of 256 samples:  Mean edit distance: 0.008 Mean normalized edit distance: 0.002\n",
            "Epoch 24/25\n",
            "400/400 [==============================] - 140s 351ms/step - loss: 0.0413 - val_loss: 0.1134\n",
            "\n",
            "Out of 256 samples:  Mean edit distance: 0.020 Mean normalized edit distance: 0.002\n",
            "Epoch 25/25\n",
            "400/400 [==============================] - 140s 349ms/step - loss: 0.0334 - val_loss: 0.1248\n",
            "\n",
            "Out of 256 samples:  Mean edit distance: 0.031 Mean normalized edit distance: 0.004\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "67Xm82u4ioh5",
        "outputId": "d3344973-45e6-40a4-bdd4-63b6e4fedeee"
      },
      "source": [
        "char_to_num = layers.StringLookup(\n",
        "    vocabulary=['a', 'b', 'c', 'd'], mask_token=None\n",
        ")\n",
        "\n",
        "print(char_to_num(tf.strings.unicode_split(\"abd\", input_encoding=\"UTF-8\")))"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "tf.Tensor([1 2 4], shape=(3,), dtype=int64)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "k5DdNfz4EiU5"
      },
      "source": [
        "Current experimetns\n",
        "\n",
        "Epoch | tf2x |glorot uniform|Bidirectional| dense2 default init|\n",
        "-----:|:----:|:------------:|:-----------:|:------------------:|\n",
        "10    | 0.88 | 0.008        | 0.009       | 0.013\n",
        "15    | 1.00 | 0.028        | 0.003       | 0.007\n",
        "20    | 1.00 | 0.005        | 0.001       | 0.000\n",
        "25    | -    | 0.002        | 0.003       | 0.004\n",
        "\n",
        "Previous best\n",
        "\n",
        "Epoch |   TF   | \n",
        "-----:|:------:|\n",
        "    10|  0.027   \n",
        "    15|  0.038   \n",
        "    20|  0.043   \n",
        "    25|  0.014 "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bnzC_Jj2EeWt"
      },
      "source": [
        "- Issue 1 [label issue](https://github.com/keras-team/keras/issues/14298)\n",
        "- Issue 2 [Use bidirectional](https://github.com/keras-team/keras/issues/5624)\n",
        "- Issue 3 [he normalization vs glorot_uniform](https://stackoverflow.com/questions/48641192/xavier-and-he-normal-initialization-difference)\n",
        "- [Captcha OCR](https://keras.io/examples/vision/captcha_ocr/)"
      ]
    }
  ]
}