54-captcha-ocr.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Pb87gBUvRqu_"
      },
      "source": [
        "# OCR model for reading Captchas\n",
        "\n",
        "**Author:** [A_K_Nain](https://twitter.com/A_K_Nain)<br>\n",
        "**Date created:** 2020/06/14<br>\n",
        "**Last modified:** 2020/06/26<br>\n",
        "**Description:** How to implement an OCR model using CNNs, RNNs and CTC loss."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RmORFbjSRqvC"
      },
      "source": [
        "## Introduction\n",
        "\n",
        "This example demonstrates a simple OCR model built with the Functional API. Apart from\n",
        "combining CNN and RNN, it also illustrates how you can instantiate a new layer\n",
        "and use it as an \"Endpoint layer\" for implementing CTC loss. For a detailed\n",
        "guide to layer subclassing, please check out\n",
        "[this page](https://keras.io/guides/making_new_layers_and_models_via_subclassing/)\n",
        "in the developer guides."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "J-bKprA0RqvD"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fFYUPaM5RqvE"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "from pathlib import Path\n",
        "from collections import Counter\n",
        "\n",
        "import tensorflow as tf\n",
        "from tensorflow import keras\n",
        "from tensorflow.keras import layers\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7GfzJJLmRqvF"
      },
      "source": [
        "## Load the data: [Captcha Images](https://www.kaggle.com/fournierp/captcha-version-2-images)\n",
        "Let's download the data."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HotMP3n0RqvG"
      },
      "outputs": [],
      "source": [
        "!curl -LO https://github.com/AakashKumarNain/CaptchaCracker/raw/master/captcha_images_v2.zip\n",
        "!unzip -qq captcha_images_v2.zip"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3MbiZE0dRqvH"
      },
      "source": [
        "The dataset contains 1040 captcha files as `png` images. The label for each sample is a string,\n",
        "the name of the file (minus the file extension).\n",
        "We will map each character in the string to an integer for training the model. Similary,\n",
        "we will need to map the predictions of the model back to strings. For this purpose\n",
        "we will maintain two dictionaries, mapping characters to integers, and integers to characters,\n",
        "respectively."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5_-Bh379RqvI"
      },
      "outputs": [],
      "source": [
        "\n",
        "# Path to the data directory\n",
        "data_dir = Path(\"./captcha_images_v2/\")\n",
        "\n",
        "# Get list of all the images\n",
        "images = sorted(list(map(str, list(data_dir.glob(\"*.png\")))))\n",
        "labels = [img.split(os.path.sep)[-1].split(\".png\")[0] for img in images]\n",
        "characters = set(char for label in labels for char in label)\n",
        "characters = sorted(list(characters))\n",
        "\n",
        "print(\"Number of images found: \", len(images))\n",
        "print(\"Number of labels found: \", len(labels))\n",
        "print(\"Number of unique characters: \", len(characters))\n",
        "print(\"Characters present: \", characters)\n",
        "\n",
        "# Batch size for training and validation\n",
        "batch_size = 16\n",
        "\n",
        "# Desired image dimensions\n",
        "img_width = 200\n",
        "img_height = 50\n",
        "\n",
        "# Factor by which the image is going to be downsampled\n",
        "# by the convolutional blocks. We will be using two\n",
        "# convolution blocks and each block will have\n",
        "# a pooling layer which downsample the features by a factor of 2.\n",
        "# Hence total downsampling factor would be 4.\n",
        "downsample_factor = 4\n",
        "\n",
        "# Maximum length of any captcha in the dataset\n",
        "max_length = max([len(label) for label in labels])\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1lWNMBlWRqvI"
      },
      "source": [
        "## Preprocessing"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fteC_U-bRqvK"
      },
      "outputs": [],
      "source": [
        "\n",
        "# Mapping characters to integers\n",
        "char_to_num = layers.StringLookup(\n",
        "    vocabulary=list(characters), mask_token=None\n",
        ")\n",
        "\n",
        "# Mapping integers back to original characters\n",
        "num_to_char = layers.StringLookup(\n",
        "    vocabulary=char_to_num.get_vocabulary(), mask_token=None, invert=True\n",
        ")\n",
        "\n",
        "\n",
        "def split_data(images, labels, train_size=0.9, shuffle=True):\n",
        "    # 1. Get the total size of the dataset\n",
        "    size = len(images)\n",
        "    # 2. Make an indices array and shuffle it, if required\n",
        "    indices = np.arange(size)\n",
        "    if shuffle:\n",
        "        np.random.shuffle(indices)\n",
        "    # 3. Get the size of training samples\n",
        "    train_samples = int(size * train_size)\n",
        "    # 4. Split data into training and validation sets\n",
        "    x_train, y_train = images[indices[:train_samples]], labels[indices[:train_samples]]\n",
        "    x_valid, y_valid = images[indices[train_samples:]], labels[indices[train_samples:]]\n",
        "    return x_train, x_valid, y_train, y_valid\n",
        "\n",
        "\n",
        "# Splitting data into training and validation sets\n",
        "x_train, x_valid, y_train, y_valid = split_data(np.array(images), np.array(labels))\n",
        "\n",
        "\n",
        "def encode_single_sample(img_path, label):\n",
        "    # 1. Read image\n",
        "    img = tf.io.read_file(img_path)\n",
        "    # 2. Decode and convert to grayscale\n",
        "    img = tf.io.decode_png(img, channels=1)\n",
        "    # 3. Convert to float32 in [0, 1] range\n",
        "    img = tf.image.convert_image_dtype(img, tf.float32)\n",
        "    # 4. Resize to the desired size\n",
        "    img = tf.image.resize(img, [img_height, img_width])\n",
        "    # 5. Transpose the image because we want the time\n",
        "    # dimension to correspond to the width of the image.\n",
        "    img = tf.transpose(img, perm=[1, 0, 2])\n",
        "    # 6. Map the characters in label to numbers\n",
        "    label = char_to_num(tf.strings.unicode_split(label, input_encoding=\"UTF-8\"))\n",
        "    # 7. Return a dict as our model is expecting two inputs\n",
        "    return {\"image\": img, \"label\": label}\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ykyjmxGkRqvL"
      },
      "source": [
        "## Create `Dataset` objects"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WIuLF0-vRqvM"
      },
      "outputs": [],
      "source": [
        "\n",
        "train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))\n",
        "train_dataset = (\n",
        "    train_dataset.map(\n",
        "        encode_single_sample, num_parallel_calls=tf.data.AUTOTUNE\n",
        "    )\n",
        "    .batch(batch_size)\n",
        "    .prefetch(buffer_size=tf.data.AUTOTUNE)\n",
        ")\n",
        "\n",
        "validation_dataset = tf.data.Dataset.from_tensor_slices((x_valid, y_valid))\n",
        "validation_dataset = (\n",
        "    validation_dataset.map(\n",
        "        encode_single_sample, num_parallel_calls=tf.data.AUTOTUNE\n",
        "    )\n",
        "    .batch(batch_size)\n",
        "    .prefetch(buffer_size=tf.data.AUTOTUNE)\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yV-0EewtRqvN"
      },
      "source": [
        "## Visualize the data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QqZEsbV-RqvO"
      },
      "outputs": [],
      "source": [
        "\n",
        "_, ax = plt.subplots(4, 4, figsize=(10, 5))\n",
        "for batch in train_dataset.take(1):\n",
        "    images = batch[\"image\"]\n",
        "    labels = batch[\"label\"]\n",
        "    for i in range(16):\n",
        "        img = (images[i] * 255).numpy().astype(\"uint8\")\n",
        "        label = tf.strings.reduce_join(num_to_char(labels[i])).numpy().decode(\"utf-8\")\n",
        "        ax[i // 4, i % 4].imshow(img[:, :, 0].T, cmap=\"gray\")\n",
        "        ax[i // 4, i % 4].set_title(label)\n",
        "        ax[i // 4, i % 4].axis(\"off\")\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MnJert5dRqvP"
      },
      "source": [
        "## Model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_FQ8fEktRqvP"
      },
      "outputs": [],
      "source": [
        "\n",
        "class CTCLayer(layers.Layer):\n",
        "    def __init__(self, name=None):\n",
        "        super().__init__(name=name)\n",
        "        self.loss_fn = keras.backend.ctc_batch_cost\n",
        "\n",
        "    def call(self, y_true, y_pred):\n",
        "        # Compute the training-time loss value and add it\n",
        "        # to the layer using `self.add_loss()`.\n",
        "        batch_len = tf.cast(tf.shape(y_true)[0], dtype=\"int64\")\n",
        "        input_length = tf.cast(tf.shape(y_pred)[1], dtype=\"int64\")\n",
        "        label_length = tf.cast(tf.shape(y_true)[1], dtype=\"int64\")\n",
        "\n",
        "        input_length = input_length * tf.ones(shape=(batch_len, 1), dtype=\"int64\")\n",
        "        label_length = label_length * tf.ones(shape=(batch_len, 1), dtype=\"int64\")\n",
        "\n",
        "        loss = self.loss_fn(y_true, y_pred, input_length, label_length)\n",
        "        self.add_loss(loss)\n",
        "\n",
        "        # At test time, just return the computed predictions\n",
        "        return y_pred\n",
        "\n",
        "\n",
        "def build_model():\n",
        "    # Inputs to the model\n",
        "    input_img = layers.Input(\n",
        "        shape=(img_width, img_height, 1), name=\"image\", dtype=\"float32\"\n",
        "    )\n",
        "    labels = layers.Input(name=\"label\", shape=(None,), dtype=\"float32\")\n",
        "\n",
        "    # First conv block\n",
        "    x = layers.Conv2D(\n",
        "        32,\n",
        "        (3, 3),\n",
        "        activation=\"relu\",\n",
        "        kernel_initializer=\"he_normal\",\n",
        "        padding=\"same\",\n",
        "        name=\"Conv1\",\n",
        "    )(input_img)\n",
        "    x = layers.MaxPooling2D((2, 2), name=\"pool1\")(x)\n",
        "\n",
        "    # Second conv block\n",
        "    x = layers.Conv2D(\n",
        "        64,\n",
        "        (3, 3),\n",
        "        activation=\"relu\",\n",
        "        kernel_initializer=\"he_normal\",\n",
        "        padding=\"same\",\n",
        "        name=\"Conv2\",\n",
        "    )(x)\n",
        "    x = layers.MaxPooling2D((2, 2), name=\"pool2\")(x)\n",
        "\n",
        "    # We have used two max pool with pool size and strides 2.\n",
        "    # Hence, downsampled feature maps are 4x smaller. The number of\n",
        "    # filters in the last layer is 64. Reshape accordingly before\n",
        "    # passing the output to the RNN part of the model\n",
        "    new_shape = ((img_width // 4), (img_height // 4) * 64)\n",
        "    x = layers.Reshape(target_shape=new_shape, name=\"reshape\")(x)\n",
        "    x = layers.Dense(64, activation=\"relu\", name=\"dense1\")(x)\n",
        "    x = layers.Dropout(0.2)(x)\n",
        "\n",
        "    # RNNs\n",
        "    x = layers.Bidirectional(layers.LSTM(128, return_sequences=True, dropout=0.25))(x)\n",
        "    x = layers.Bidirectional(layers.LSTM(64, return_sequences=True, dropout=0.25))(x)\n",
        "\n",
        "    # Output layer\n",
        "    x = layers.Dense(\n",
        "        len(char_to_num.get_vocabulary()) + 1, activation=\"softmax\", name=\"dense2\"\n",
        "    )(x)\n",
        "\n",
        "    # Add CTC layer for calculating CTC loss at each step\n",
        "    output = CTCLayer(name=\"ctc_loss\")(labels, x)\n",
        "\n",
        "    # Define the model\n",
        "    model = keras.models.Model(\n",
        "        inputs=[input_img, labels], outputs=output, name=\"ocr_model_v1\"\n",
        "    )\n",
        "    # Optimizer\n",
        "    opt = keras.optimizers.Adam()\n",
        "    # Compile the model and return\n",
        "    model.compile(optimizer=opt)\n",
        "    return model\n",
        "\n",
        "\n",
        "# Get the model\n",
        "model = build_model()\n",
        "model.summary()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Yy2UhI__RqvQ"
      },
      "source": [
        "## Training"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cM-7KEVaRqvR"
      },
      "outputs": [],
      "source": [
        "\n",
        "epochs = 100\n",
        "early_stopping_patience = 10\n",
        "# Add early stopping\n",
        "early_stopping = keras.callbacks.EarlyStopping(\n",
        "    monitor=\"val_loss\", patience=early_stopping_patience, restore_best_weights=True\n",
        ")\n",
        "\n",
        "# Train the model\n",
        "history = model.fit(\n",
        "    train_dataset,\n",
        "    validation_data=validation_dataset,\n",
        "    epochs=epochs,\n",
        "    callbacks=[early_stopping],\n",
        ")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MVq9hZgLRqvS"
      },
      "source": [
        "## Inference\n",
        "\n",
        "You can use the trained model hosted on [Hugging Face Hub](https://huggingface.co/keras-io/ocr-for-captcha)\n",
        "and try the demo on [Hugging Face Spaces](https://huggingface.co/spaces/keras-io/ocr-for-captcha)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dyalNqr4RqvS"
      },
      "outputs": [],
      "source": [
        "\n",
        "# Get the prediction model by extracting layers till the output layer\n",
        "prediction_model = keras.models.Model(\n",
        "    model.get_layer(name=\"image\").input, model.get_layer(name=\"dense2\").output\n",
        ")\n",
        "prediction_model.summary()\n",
        "\n",
        "# A utility function to decode the output of the network\n",
        "def decode_batch_predictions(pred):\n",
        "    input_len = np.ones(pred.shape[0]) * pred.shape[1]\n",
        "    # Use greedy search. For complex tasks, you can use beam search\n",
        "    results = keras.backend.ctc_decode(pred, input_length=input_len, greedy=True)[0][0][\n",
        "        :, :max_length\n",
        "    ]\n",
        "    # Iterate over the results and get back the text\n",
        "    output_text = []\n",
        "    for res in results:\n",
        "        res = tf.strings.reduce_join(num_to_char(res)).numpy().decode(\"utf-8\")\n",
        "        output_text.append(res)\n",
        "    return output_text\n",
        "\n",
        "\n",
        "#  Let's check results on some validation samples\n",
        "for batch in validation_dataset.take(1):\n",
        "    batch_images = batch[\"image\"]\n",
        "    batch_labels = batch[\"label\"]\n",
        "\n",
        "    preds = prediction_model.predict(batch_images)\n",
        "    pred_texts = decode_batch_predictions(preds)\n",
        "\n",
        "    orig_texts = []\n",
        "    for label in batch_labels:\n",
        "        label = tf.strings.reduce_join(num_to_char(label)).numpy().decode(\"utf-8\")\n",
        "        orig_texts.append(label)\n",
        "\n",
        "    _, ax = plt.subplots(4, 4, figsize=(15, 5))\n",
        "    for i in range(len(pred_texts)):\n",
        "        img = (batch_images[i, :, :, 0] * 255).numpy().astype(np.uint8)\n",
        "        img = img.T\n",
        "        title = f\"Prediction: {pred_texts[i]}\"\n",
        "        ax[i // 4, i % 4].imshow(img, cmap=\"gray\")\n",
        "        ax[i // 4, i % 4].set_title(title)\n",
        "        ax[i // 4, i % 4].axis(\"off\")\n",
        "plt.show()"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "captcha_ocr",
      "provenance": [],
      "toc_visible": true
    },
    "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.0"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}