title ocr model for reading captchas aknainhttps twitter comaknain date created 20200614 last modified 20200626 description how to implement an ocr model using cnns rnns and ctc loss accelerator gpu introduction this example demonstrates a simple ocr model built with the functional api apart from combining cnn and rnn it also illustrates how you can instantiate a new layer and use it as an endpoint layer for implementing ctc loss for a detailed guide to layer subclassing please check out this pagehttps keras ioguidesmakingnewlayersandmodelsviasubclassing in the developer guides setup load the data captcha imageshttps www kaggle comfournierpcaptchaversion2images let s download the data shell curl lo https github comaakashkumarnaincaptchacrackerrawmastercaptchaimagesv2 zip unzip qq captchaimagesv2 zip the dataset contains 1040 captcha files as png images the label for each sample is a string the name of the file minus the file extension we will map each character in the string to an integer for training the model similary we will need to map the predictions of the model back to strings for this purpose we will maintain two dictionaries mapping characters to integers and integers to characters respectively path to the data directory get list of all the images batch size for training and validation desired image dimensions factor by which the image is going to be downsampled by the convolutional blocks we will be using two convolution blocks and each block will have a pooling layer which downsample the features by a factor of 2 hence total downsampling factor would be 4 maximum length of any captcha in the dataset preprocessing mapping characters to integers mapping integers back to original characters 1 get the total size of the dataset 2 make an indices array and shuffle it if required 3 get the size of training samples 4 split data into training and validation sets splitting data into training and validation sets 1 read image 2 decode and convert to grayscale 3 convert to float32 in 0 1 range 4 resize to the desired size 5 transpose the image because we want the time dimension to correspond to the width of the image 6 map the characters in label to numbers 7 return a dict as our model is expecting two inputs create dataset objects visualize the data model compute the trainingtime loss value and add it to the layer using self addloss at test time just return the computed predictions inputs to the model first conv block second conv block we have used two max pool with pool size and strides 2 hence downsampled feature maps are 4x smaller the number of filters in the last layer is 64 reshape accordingly before passing the output to the rnn part of the model rnns output layer add ctc layer for calculating ctc loss at each step define the model optimizer compile the model and return get the model training add early stopping train the model inference you can use the trained model hosted on hugging face hubhttps huggingface cokerasioocrforcaptcha and try the demo on hugging face spaceshttps huggingface cospaceskerasioocrforcaptcha get the prediction model by extracting layers till the output layer a utility function to decode the output of the network use greedy search for complex tasks you can use beam search iterate over the results and get back the text let s check results on some validation samples introduction this example demonstrates a simple ocr model built with the functional api apart from combining cnn and rnn it also illustrates how you can instantiate a new layer and use it as an endpoint layer for implementing ctc loss for a detailed guide to layer subclassing please check out this page https keras io guides making_new_layers_and_models_via_subclassing in the developer guides setup load the data captcha images https www kaggle com fournierp captcha version 2 images let s download the data shell curl lo https github com aakashkumarnain captchacracker raw master captcha_images_v2 zip unzip qq captcha_images_v2 zip the dataset contains 1040 captcha files as png images the label for each sample is a string the name of the file minus the file extension we will map each character in the string to an integer for training the model similary we will need to map the predictions of the model back to strings for this purpose we will maintain two dictionaries mapping characters to integers and integers to characters respectively path to the data directory get list of all the images batch size for training and validation desired image dimensions factor by which the image is going to be downsampled by the convolutional blocks we will be using two convolution blocks and each block will have a pooling layer which downsample the features by a factor of 2 hence total downsampling factor would be 4 maximum length of any captcha in the dataset preprocessing mapping characters to integers mapping integers back to original characters 1 get the total size of the dataset 2 make an indices array and shuffle it if required 3 get the size of training samples 4 split data into training and validation sets splitting data into training and validation sets 1 read image 2 decode and convert to grayscale 3 convert to float32 in 0 1 range 4 resize to the desired size 5 transpose the image because we want the time dimension to correspond to the width of the image 6 map the characters in label to numbers 7 return a dict as our model is expecting two inputs create dataset objects visualize the data model compute the training time loss value and add it to the layer using self add_loss at test time just return the computed predictions inputs to the model first conv block second conv block we have used two max pool with pool size and strides 2 hence downsampled feature maps are 4x smaller the number of filters in the last layer is 64 reshape accordingly before passing the output to the rnn part of the model rnns output layer add ctc layer for calculating ctc loss at each step define the model optimizer compile the model and return get the model training add early stopping train the model inference you can use the trained model hosted on hugging face hub https huggingface co keras io ocr for captcha and try the demo on hugging face spaces https huggingface co spaces keras io ocr for captcha get the prediction model by extracting layers till the output layer a utility function to decode the output of the network use greedy search for complex tasks you can use beam search iterate over the results and get back the text let s check results on some validation samples
import os
import numpy as np
import matplotlib.pyplot as plt

from pathlib import Path
from collections import Counter

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers












data_dir = Path("./captcha_images_v2/")


images = sorted(list(map(str, list(data_dir.glob("*.png")))))
labels = [img.split(os.path.sep)[-1].split(".png")[0] for img in images]
characters = set(char for label in labels for char in label)
characters = sorted(list(characters))

print("Number of images found: ", len(images))
print("Number of labels found: ", len(labels))
print("Number of unique characters: ", len(characters))
print("Characters present: ", characters)


batch_size = 16


img_width = 200
img_height = 50






downsample_factor = 4


max_length = max([len(label) for label in labels])






char_to_num = layers.StringLookup(vocabulary=list(characters), mask_token=None)


num_to_char = layers.StringLookup(
    vocabulary=char_to_num.get_vocabulary(), mask_token=None, invert=True
)


def split_data(images, labels, train_size=0.9, shuffle=True):
    
    size = len(images)
    
    indices = np.arange(size)
    if shuffle:
        np.random.shuffle(indices)
    
    train_samples = int(size * train_size)
    
    x_train, y_train = (
        images[indices[:train_samples]],
        labels[indices[:train_samples]],
    )
    x_valid, y_valid = (
        images[indices[train_samples:]],
        labels[indices[train_samples:]],
    )
    return x_train, x_valid, y_train, y_valid



x_train, x_valid, y_train, y_valid = split_data(
    np.array(images), np.array(labels)
)


def encode_single_sample(img_path, label):
    
    img = tf.io.read_file(img_path)
    
    img = tf.io.decode_png(img, channels=1)
    
    img = tf.image.convert_image_dtype(img, tf.float32)
    
    img = tf.image.resize(img, [img_height, img_width])
    
    
    img = tf.transpose(img, perm=[1, 0, 2])
    
    label = char_to_num(tf.strings.unicode_split(label, input_encoding="UTF-8"))
    
    return {"image": img, "label": label}





train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
train_dataset = (
    train_dataset.map(encode_single_sample, num_parallel_calls=tf.data.AUTOTUNE)
    .batch(batch_size)
    .prefetch(buffer_size=tf.data.AUTOTUNE)
)

validation_dataset = tf.data.Dataset.from_tensor_slices((x_valid, y_valid))
validation_dataset = (
    validation_dataset.map(
        encode_single_sample, num_parallel_calls=tf.data.AUTOTUNE
    )
    .batch(batch_size)
    .prefetch(buffer_size=tf.data.AUTOTUNE)
)




_, ax = plt.subplots(4, 4, figsize=(10, 5))
for batch in train_dataset.take(1):
    images = batch["image"]
    labels = batch["label"]
    for i in range(16):
        img = (images[i] * 255).numpy().astype("uint8")
        label = (
            tf.strings.reduce_join(num_to_char(labels[i]))
            .numpy()
            .decode("utf-8")
        )
        ax[i // 4, i % 4].imshow(img[:, :, 0].T, cmap="gray")
        ax[i // 4, i % 4].set_title(label)
        ax[i // 4, i % 4].axis("off")
plt.show()




class CTCLayer(layers.Layer):
    def __init__(self, name=None):
        super().__init__(name=name)
        self.loss_fn = keras.backend.ctc_batch_cost

    def call(self, y_true, y_pred):
        
        
        batch_len = tf.cast(tf.shape(y_true)[0], dtype="int64")
        input_length = tf.cast(tf.shape(y_pred)[1], dtype="int64")
        label_length = tf.cast(tf.shape(y_true)[1], dtype="int64")

        input_length = input_length * tf.ones(
            shape=(batch_len, 1), dtype="int64"
        )
        label_length = label_length * tf.ones(
            shape=(batch_len, 1), dtype="int64"
        )

        loss = self.loss_fn(y_true, y_pred, input_length, label_length)
        self.add_loss(loss)

        
        return y_pred


def build_model():
    
    input_img = layers.Input(
        shape=(img_width, img_height, 1), name="image", dtype="float32"
    )
    labels = layers.Input(name="label", shape=(None,), dtype="float32")

    
    x = layers.Conv2D(
        32,
        (3, 3),
        activation="relu",
        kernel_initializer="he_normal",
        padding="same",
        name="Conv1",
    )(input_img)
    x = layers.MaxPooling2D((2, 2), name="pool1")(x)

    
    x = layers.Conv2D(
        64,
        (3, 3),
        activation="relu",
        kernel_initializer="he_normal",
        padding="same",
        name="Conv2",
    )(x)
    x = layers.MaxPooling2D((2, 2), name="pool2")(x)

    
    
    
    
    new_shape = ((img_width // 4), (img_height // 4) * 64)
    x = layers.Reshape(target_shape=new_shape, name="reshape")(x)
    x = layers.Dense(64, activation="relu", name="dense1")(x)
    x = layers.Dropout(0.2)(x)

    
    x = layers.Bidirectional(
        layers.LSTM(128, return_sequences=True, dropout=0.25)
    )(x)
    x = layers.Bidirectional(
        layers.LSTM(64, return_sequences=True, dropout=0.25)
    )(x)

    
    x = layers.Dense(
        len(char_to_num.get_vocabulary()) + 1,
        activation="softmax",
        name="dense2",
    )(x)

    
    output = CTCLayer(name="ctc_loss")(labels, x)

    
    model = keras.models.Model(
        inputs=[input_img, labels], outputs=output, name="ocr_model_v1"
    )
    
    opt = keras.optimizers.Adam()
    
    model.compile(optimizer=opt)
    return model



model = build_model()
model.summary()




epochs = 1
early_stopping_patience = 10

early_stopping = keras.callbacks.EarlyStopping(
    monitor="val_loss",
    patience=early_stopping_patience,
    restore_best_weights=True,
)


history = model.fit(
    train_dataset,
    validation_data=validation_dataset,
    epochs=epochs,
    callbacks=[early_stopping],
)






prediction_model = keras.models.Model(
    model.get_layer(name="image").input, model.get_layer(name="dense2").output
)
prediction_model.summary()



def decode_batch_predictions(pred):
    input_len = np.ones(pred.shape[0]) * pred.shape[1]
    
    results = keras.backend.ctc_decode(
        pred, input_length=input_len, greedy=True
    )[0][0][:, :max_length]
    
    output_text = []
    for res in results:
        res = tf.strings.reduce_join(num_to_char(res)).numpy().decode("utf-8")
        output_text.append(res)
    return output_text



for batch in validation_dataset.take(1):
    batch_images = batch["image"]
    batch_labels = batch["label"]

    preds = prediction_model.predict(batch_images)
    pred_texts = decode_batch_predictions(preds)

    orig_texts = []
    for label in batch_labels:
        label = (
            tf.strings.reduce_join(num_to_char(label)).numpy().decode("utf-8")
        )
        orig_texts.append(label)

    _, ax = plt.subplots(4, 4, figsize=(15, 5))
    for i in range(len(pred_texts)):
        img = (batch_images[i, :, :, 0] * 255).numpy().astype(np.uint8)
        img = img.T
        title = f"Prediction: {pred_texts[i]}"
        ax[i // 4, i % 4].imshow(img, cmap="gray")
        ax[i // 4, i % 4].set_title(title)
        ax[i // 4, i % 4].axis("off")
plt.show()