Optical character recognition

    This starts off with 4 letter words. For the first 12 epochs, thedifficulty is gradually increased using the TextImageGenerator classwhich is both a generator class for test/train data and a Kerascallback class. After 20 epochs, longer sequences are thrown at itby recompiling the model to handle a wider image and rebuildingthe word list to include two words separated by a space.

    Additional dependencies

    This requires and editdistance packages:

    Then install Python dependencies:

    1. import os
    2. import itertools
    3. import codecs
    4. import re
    5. import datetime
    6. import cairocffi as cairo
    7. import editdistance
    8. import numpy as np
    9. from scipy import ndimage
    10. import pylab
    11. from keras import backend as K
    12. from keras.layers.convolutional import Conv2D, MaxPooling2D
    13. from keras.layers import Input, Dense, Activation
    14. from keras.layers import Reshape, Lambda
    15. from keras.layers.merge import add, concatenate
    16. from keras.models import Model
    17. from keras.layers.recurrent import GRU
    18. from keras.optimizers import SGD
    19. from keras.utils.data_utils import get_file
    20. from keras.preprocessing import image
    21. import keras.callbacks
    24. OUTPUT_DIR = 'image_ocr'
    26. # character classes and matching regex filter
    27. regex = r'^[a-z ]+$'
    28. alphabet = u'abcdefghijklmnopqrstuvwxyz '
    30. np.random.seed(55)
    33. # this creates larger "blotches" of noise which look
    34. # more realistic than just adding gaussian noise
    35. # assumes greyscale with pixels ranging from 0 to 1
    37. def speckle(img):
    38.     severity = np.random.uniform(0, 0.6)
    39.     blur = ndimage.gaussian_filter(np.random.randn(*img.shape) * severity, 1)
    40.     img_speck = (img + blur)
    41.     img_speck[img_speck > 1] = 1
    42.     img_speck[img_speck <= 0] = 0
    43.     return img_speck
    46. # paints the string in a random location the bounding box
    47. # also uses a random font, a slight random rotation,
    48. # and a random amount of speckle noise
    50. def paint_text(text, w, h, rotate=False, ud=False, multi_fonts=False):
    51.     surface = cairo.ImageSurface(cairo.FORMAT_RGB24, w, h)
    52.     with cairo.Context(surface) as context:
    53.         context.set_source_rgb(1, 1, 1)  # White
    54.         context.paint()
    55.         # this font list works in CentOS 7
    56.         if multi_fonts:
    57.             fonts = [
    58.                 'Century Schoolbook', 'Courier', 'STIX',
    59.                 'URW Chancery L', 'FreeMono']
    60.             context.select_font_face(
    61.                 np.random.choice(fonts),
    62.                 cairo.FONT_SLANT_NORMAL,
    63.                 np.random.choice([cairo.FONT_WEIGHT_BOLD, cairo.FONT_WEIGHT_NORMAL]))
    64.         else:
    65.             context.select_font_face('Courier',
    66.                                      cairo.FONT_SLANT_NORMAL,
    67.                                      cairo.FONT_WEIGHT_BOLD)
    68.         context.set_font_size(25)
    69.         box = context.text_extents(text)
    70.         border_w_h = (4, 4)
    71.         if box[2] > (w - 2 * border_w_h[1]) or box[3] > (h - 2 * border_w_h[0]):
    72.             raise IOError(('Could not fit string into image.'
    73.                            'Max char count is too large for given image width.'))
    75.         # teach the RNN translational invariance by
    76.         # fitting text box randomly on canvas, with some room to rotate
    77.         max_shift_x = w - box[2] - border_w_h[0]
    78.         max_shift_y = h - box[3] - border_w_h[1]
    79.         top_left_x = np.random.randint(0, int(max_shift_x))
    80.         if ud:
    81.             top_left_y = np.random.randint(0, int(max_shift_y))
    82.         else:
    83.             top_left_y = h // 2
    84.         context.move_to(top_left_x - int(box[0]), top_left_y - int(box[1]))
    85.         context.set_source_rgb(0, 0, 0)
    86.         context.show_text(text)
    88.     buf = surface.get_data()
    89.     a = np.frombuffer(buf, np.uint8)
    90.     a.shape = (h, w, 4)
    91.     a = a[:, :, 0]  # grab single channel
    92.     a = a.astype(np.float32) / 255
    93.     a = np.expand_dims(a, 0)
    94.     if rotate:
    95.         a = image.random_rotation(a, 3 * (w - top_left_x) / w + 1)
    96.     a = speckle(a)
    98.     return a
    101. def shuffle_mats_or_lists(matrix_list, stop_ind=None):
    102.     ret = []
    103.     assert all([len(i) == len(matrix_list[0]) for i in matrix_list])
    104.     len_val = len(matrix_list[0])
    105.     if stop_ind is None:
    106.         stop_ind = len_val
    107.     assert stop_ind <= len_val
    109.     a = list(range(stop_ind))
    110.     np.random.shuffle(a)
    111.     a += list(range(stop_ind, len_val))
    112.     for mat in matrix_list:
    113.         if isinstance(mat, np.ndarray):
    114.             ret.append(mat[a])
    115.         elif isinstance(mat, list):
    116.             ret.append([mat[i] for i in a])
    117.         else:
    118.             raise TypeError('`shuffle_mats_or_lists` only supports '
    119.                             'numpy.array and list objects.')
    120.     return ret
    123. # Translation of characters to unique integer values
    124. def text_to_labels(text):
    125.     ret = []
    126.     for char in text:
    127.         ret.append(alphabet.find(char))
    128.     return ret
    131. # Reverse translation of numerical classes back to characters
    132. def labels_to_text(labels):
    133.     ret = []
    134.     for c in labels:
    135.         if c == len(alphabet):  # CTC Blank
    136.             ret.append("")
    137.         else:
    138.             ret.append(alphabet[c])
    139.     return "".join(ret)
    142. # only a-z and space..probably not to difficult
    143. # to expand to uppercase and symbols
    145. def is_valid_str(in_str):
    146.     search = re.compile(regex, re.UNICODE).search
    147.     return bool(search(in_str))
    150. # Uses generator functions to supply train/test with
    151. # data. Image renderings and text are created on the fly
    152. # each time with random perturbations
    154. class TextImageGenerator(keras.callbacks.Callback):
    156.     def __init__(self, monogram_file, bigram_file, minibatch_size,
    157.                  img_w, img_h, downsample_factor, val_split,
    158.                  absolute_max_string_len=16):
    160.         self.minibatch_size = minibatch_size
    161.         self.img_w = img_w
    162.         self.img_h = img_h
    163.         self.monogram_file = monogram_file
    164.         self.bigram_file = bigram_file
    165.         self.downsample_factor = downsample_factor
    166.         self.val_split = val_split
    167.         self.blank_label = self.get_output_size() - 1
    168.         self.absolute_max_string_len = absolute_max_string_len
    170.     def get_output_size(self):
    171.         return len(alphabet) + 1
    172.     # num_words can be independent of the epoch size due to the use of generators
    173.     # as max_string_len grows, num_words can grow
    174.     def build_word_list(self, num_words, max_string_len=None, mono_fraction=0.5):
    175.         assert max_string_len <= self.absolute_max_string_len
    176.         assert num_words % self.minibatch_size == 0
    177.         assert (self.val_split * num_words) % self.minibatch_size == 0
    178.         self.num_words = num_words
    179.         self.string_list = [''] * self.num_words
    180.         tmp_string_list = []
    181.         self.max_string_len = max_string_len
    182.         self.Y_data = np.ones([self.num_words, self.absolute_max_string_len]) * -1
    183.         self.X_text = []
    184.         self.Y_len = [0] * self.num_words
    186.         def _is_length_of_word_valid(word):
    187.             return (max_string_len == -1 or
    188.                     max_string_len is None or
    189.                     len(word) <= max_string_len)
    191.         # monogram file is sorted by frequency in english speech
    192.         with codecs.open(self.monogram_file, mode='r', encoding='utf-8') as f:
    193.             for line in f:
    194.                 if len(tmp_string_list) == int(self.num_words * mono_fraction):
    195.                     break
    196.                 word = line.rstrip()
    197.                 if _is_length_of_word_valid(word):
    198.                     tmp_string_list.append(word)
    200.         # bigram file contains common word pairings in english speech
    201.         with codecs.open(self.bigram_file, mode='r', encoding='utf-8') as f:
    202.             lines = f.readlines()
    203.             for line in lines:
    204.                 if len(tmp_string_list) == self.num_words:
    205.                     break
    206.                 columns = line.lower().split()
    207.                 word = columns[0] + ' ' + columns[1]
    208.                 if is_valid_str(word) and _is_length_of_word_valid(word):
    209.                     tmp_string_list.append(word)
    210.         if len(tmp_string_list) != self.num_words:
    211.             raise IOError('Could not pull enough words'
    212.                           'from supplied monogram and bigram files.')
    213.         # interlace to mix up the easy and hard words
    214.         self.string_list[::2] = tmp_string_list[:self.num_words // 2]
    215.         self.string_list[1::2] = tmp_string_list[self.num_words // 2:]
    217.         for i, word in enumerate(self.string_list):
    218.             self.Y_len[i] = len(word)
    219.             self.Y_data[i, 0:len(word)] = text_to_labels(word)
    220.             self.X_text.append(word)
    221.         self.Y_len = np.expand_dims(np.array(self.Y_len), 1)
    223.         self.cur_val_index = self.val_split
    224.         self.cur_train_index = 0
    226.     # each time an image is requested from train/val/test, a new random
    227.     # painting of the text is performed
    228.     def get_batch(self, index, size, train):
    229.         # width and height are backwards from typical Keras convention
    230.         # because width is the time dimension when it gets fed into the RNN
    231.         if K.image_data_format() == 'channels_first':
    232.             X_data = np.ones([size, 1, self.img_w, self.img_h])
    233.         else:
    234.             X_data = np.ones([size, self.img_w, self.img_h, 1])
    236.         labels = np.ones([size, self.absolute_max_string_len])
    237.         input_length = np.zeros([size, 1])
    238.         label_length = np.zeros([size, 1])
    239.         source_str = []
    240.         for i in range(size):
    241.             # Mix in some blank inputs.  This seems to be important for
    242.             # achieving translational invariance
    243.             if train and i > size - 4:
    244.                 if K.image_data_format() == 'channels_first':
    245.                     X_data[i, 0, 0:self.img_w, :] = self.paint_func('')[0, :, :].T
    246.                 else:
    247.                     X_data[i, 0:self.img_w, :, 0] = self.paint_func('',)[0, :, :].T
    248.                 labels[i, 0] = self.blank_label
    249.                 input_length[i] = self.img_w // self.downsample_factor - 2
    250.                 label_length[i] = 1
    251.                 source_str.append('')
    252.             else:
    253.                 if K.image_data_format() == 'channels_first':
    254.                     X_data[i, 0, 0:self.img_w, :] = (
    255.                         self.paint_func(self.X_text[index + i])[0, :, :].T)
    256.                 else:
    257.                     X_data[i, 0:self.img_w, :, 0] = (
    258.                         self.paint_func(self.X_text[index + i])[0, :, :].T)
    259.                 labels[i, :] = self.Y_data[index + i]
    260.                 input_length[i] = self.img_w // self.downsample_factor - 2
    261.                 label_length[i] = self.Y_len[index + i]
    262.                 source_str.append(self.X_text[index + i])
    263.         inputs = {'the_input': X_data,
    264.                   'the_labels': labels,
    265.                   'input_length': input_length,
    266.                   'label_length': label_length,
    267.                   'source_str': source_str  # used for visualization only
    268.                   }
    269.         outputs = {'ctc': np.zeros([size])}  # dummy data for dummy loss function
    270.         return (inputs, outputs)
    272.     def next_train(self):
    273.         while 1:
    274.             ret = self.get_batch(self.cur_train_index,
    275.                                  self.minibatch_size, train=True)
    276.             self.cur_train_index += self.minibatch_size
    277.             if self.cur_train_index >= self.val_split:
    278.                 self.cur_train_index = self.cur_train_index % 32
    279.                 (self.X_text, self.Y_data, self.Y_len) = shuffle_mats_or_lists(
    280.                     [self.X_text, self.Y_data, self.Y_len], self.val_split)
    281.             yield ret
    283.     def next_val(self):
    284.         while 1:
    285.             ret = self.get_batch(self.cur_val_index,
    286.                                  self.minibatch_size, train=False)
    287.             self.cur_val_index += self.minibatch_size
    288.             if self.cur_val_index >= self.num_words:
    289.                 self.cur_val_index = self.val_split + self.cur_val_index % 32
    290.             yield ret
    292.     def on_train_begin(self, logs={}):
    293.         self.build_word_list(16000, 4, 1)
    294.         self.paint_func = lambda text: paint_text(
    295.             text, self.img_w, self.img_h,
    296.             rotate=False, ud=False, multi_fonts=False)
    298.     def on_epoch_begin(self, epoch, logs={}):
    299.         # rebind the paint function to implement curriculum learning
    300.         if 3 <= epoch < 6:
    301.             self.paint_func = lambda text: paint_text(
    302.                 text, self.img_w, self.img_h,
    303.                 rotate=False, ud=True, multi_fonts=False)
    304.         elif 6 <= epoch < 9:
    305.             self.paint_func = lambda text: paint_text(
    306.                 text, self.img_w, self.img_h,
    307.                 rotate=False, ud=True, multi_fonts=True)
    308.         elif epoch >= 9:
    309.             self.paint_func = lambda text: paint_text(
    310.                 text, self.img_w, self.img_h,
    311.                 rotate=True, ud=True, multi_fonts=True)
    312.         if epoch >= 21 and self.max_string_len < 12:
    313.             self.build_word_list(32000, 12, 0.5)
    316. # the actual loss calc occurs here despite it not being
    317. # an internal Keras loss function
    319. def ctc_lambda_func(args):
    320.     y_pred, labels, input_length, label_length = args
    321.     # the 2 is critical here since the first couple outputs of the RNN
    322.     # tend to be garbage:
    323.     y_pred = y_pred[:, 2:, :]
    324.     return K.ctc_batch_cost(labels, y_pred, input_length, label_length)
    327. # For a real OCR application, this should be beam search with a dictionary
    328. # and language model.  For this example, best path is sufficient.
    330. def decode_batch(test_func, word_batch):
    331.     out = test_func([word_batch])[0]
    332.     ret = []
    333.     for j in range(out.shape[0]):
    334.         out_best = list(np.argmax(out[j, 2:], 1))
    335.         out_best = [k for k, g in itertools.groupby(out_best)]
    336.         outstr = labels_to_text(out_best)
    337.         ret.append(outstr)
    338.     return ret
    341. class VizCallback(keras.callbacks.Callback):
    343.     def __init__(self, run_name, test_func, text_img_gen, num_display_words=6):
    344.         self.output_dir = os.path.join(
    345.             OUTPUT_DIR, run_name)
    346.         self.text_img_gen = text_img_gen
    347.         self.num_display_words = num_display_words
    348.         if not os.path.exists(self.output_dir):
    349.             os.makedirs(self.output_dir)
    351.     def show_edit_distance(self, num):
    352.         num_left = num
    353.         mean_norm_ed = 0.0
    354.         mean_ed = 0.0
    355.         while num_left > 0:
    356.             word_batch = next(self.text_img_gen)[0]
    357.             num_proc = min(word_batch['the_input'].shape[0], num_left)
    358.             decoded_res = decode_batch(self.test_func,
    359.                                        word_batch['the_input'][0:num_proc])
    360.             for j in range(num_proc):
    361.                 edit_dist = editdistance.eval(decoded_res[j],
    362.                                               word_batch['source_str'][j])
    363.                 mean_ed += float(edit_dist)
    364.                 mean_norm_ed += float(edit_dist) / len(word_batch['source_str'][j])
    365.             num_left -= num_proc
    366.         mean_norm_ed = mean_norm_ed / num
    367.         mean_ed = mean_ed / num
    368.         print('\nOut of %d samples:  Mean edit distance:'
    369.               '%.3f Mean normalized edit distance: %0.3f'
    370.               % (num, mean_ed, mean_norm_ed))
    372.     def on_epoch_end(self, epoch, logs={}):
    373.         self.model.save_weights(
    374.             os.path.join(self.output_dir, 'weights%02d.h5' % (epoch)))
    375.         self.show_edit_distance(256)
    376.         word_batch = next(self.text_img_gen)[0]
    377.         res = decode_batch(self.test_func,
    378.                            word_batch['the_input'][0:self.num_display_words])
    379.         if word_batch['the_input'][0].shape[0] < 256:
    380.             cols = 2
    381.         else:
    382.             cols = 1
    383.         for i in range(self.num_display_words):
    384.             pylab.subplot(self.num_display_words // cols, cols, i + 1)
    385.             if K.image_data_format() == 'channels_first':
    386.                 the_input = word_batch['the_input'][i, 0, :, :]
    387.             else:
    388.                 the_input = word_batch['the_input'][i, :, :, 0]
    389.             pylab.imshow(the_input.T, cmap='Greys_r')
    390.             pylab.xlabel(
    391.                 'Truth = \'%s\'\nDecoded = \'%s\'' %
    392.                 (word_batch['source_str'][i], res[i]))
    393.         fig = pylab.gcf()
    394.         fig.set_size_inches(10, 13)
    395.         pylab.savefig(os.path.join(self.output_dir, 'e%02d.png' % (epoch)))
    396.         pylab.close()
    399. def train(run_name, start_epoch, stop_epoch, img_w):
    400.     # Input Parameters
    401.     img_h = 64
    402.     words_per_epoch = 16000
    403.     val_split = 0.2
    404.     val_words = int(words_per_epoch * (val_split))
    406.     # Network parameters
    407.     conv_filters = 16
    408.     kernel_size = (3, 3)
    409.     pool_size = 2
    410.     time_dense_size = 32
    411.     rnn_size = 512
    412.     minibatch_size = 32
    414.     if K.image_data_format() == 'channels_first':
    415.         input_shape = (1, img_w, img_h)
    416.     else:
    417.         input_shape = (img_w, img_h, 1)
    419.     fdir = os.path.dirname(
    420.         get_file('wordlists.tgz',
    421.                  origin='http://www.mythic-ai.com/datasets/wordlists.tgz',
    422.                  untar=True))
    424.     img_gen = TextImageGenerator(
    425.         monogram_file=os.path.join(fdir, 'wordlist_mono_clean.txt'),
    426.         bigram_file=os.path.join(fdir, 'wordlist_bi_clean.txt'),
    427.         minibatch_size=minibatch_size,
    428.         img_w=img_w,
    429.         img_h=img_h,
    430.         downsample_factor=(pool_size ** 2),
    431.         val_split=words_per_epoch - val_words)
    432.     act = 'relu'
    433.     input_data = Input(name='the_input', shape=input_shape, dtype='float32')
    434.     inner = Conv2D(conv_filters, kernel_size, padding='same',
    435.                    activation=act, kernel_initializer='he_normal',
    436.                    name='conv1')(input_data)
    437.     inner = MaxPooling2D(pool_size=(pool_size, pool_size), name='max1')(inner)
    438.     inner = Conv2D(conv_filters, kernel_size, padding='same',
    439.                    activation=act, kernel_initializer='he_normal',
    440.                    name='conv2')(inner)
    441.     inner = MaxPooling2D(pool_size=(pool_size, pool_size), name='max2')(inner)
    443.     conv_to_rnn_dims = (img_w // (pool_size ** 2),
    444.                         (img_h // (pool_size ** 2)) * conv_filters)
    445.     inner = Reshape(target_shape=conv_to_rnn_dims, name='reshape')(inner)
    447.     # cuts down input size going into RNN:
    448.     inner = Dense(time_dense_size, activation=act, name='dense1')(inner)
    450.     # Two layers of bidirectional GRUs
    451.     # GRU seems to work as well, if not better than LSTM:
    452.     gru_1 = GRU(rnn_size, return_sequences=True,
    453.                 kernel_initializer='he_normal', name='gru1')(inner)
    454.     gru_1b = GRU(rnn_size, return_sequences=True,
    455.                  go_backwards=True, kernel_initializer='he_normal',
    456.                  name='gru1_b')(inner)
    457.     gru1_merged = add([gru_1, gru_1b])
    458.     gru_2 = GRU(rnn_size, return_sequences=True,
    459.                 kernel_initializer='he_normal', name='gru2')(gru1_merged)
    460.     gru_2b = GRU(rnn_size, return_sequences=True, go_backwards=True,
    461.                  kernel_initializer='he_normal', name='gru2_b')(gru1_merged)
    463.     # transforms RNN output to character activations:
    464.     inner = Dense(img_gen.get_output_size(), kernel_initializer='he_normal',
    465.                   name='dense2')(concatenate([gru_2, gru_2b]))
    466.     y_pred = Activation('softmax', name='softmax')(inner)
    467.     Model(inputs=input_data, outputs=y_pred).summary()
    469.     labels = Input(name='the_labels',
    470.                    shape=[img_gen.absolute_max_string_len], dtype='float32')
    471.     input_length = Input(name='input_length', shape=[1], dtype='int64')
    472.     label_length = Input(name='label_length', shape=[1], dtype='int64')
    473.     # Keras doesn't currently support loss funcs with extra parameters
    474.     # so CTC loss is implemented in a lambda layer
    475.     loss_out = Lambda(
    476.         ctc_lambda_func, output_shape=(1,),
    477.         name='ctc')([y_pred, labels, input_length, label_length])
    479.     # clipnorm seems to speeds up convergence
    480.     sgd = SGD(learning_rate=0.02,
    481.               decay=1e-6,
    482.               momentum=0.9,
    483.               nesterov=True)
    485.     model = Model(inputs=[input_data, labels, input_length, label_length],
    486.                   outputs=loss_out)
    488.     # the loss calc occurs elsewhere, so use a dummy lambda func for the loss
    489.     model.compile(loss={'ctc': lambda y_true, y_pred: y_pred}, optimizer=sgd)
    490.     if start_epoch > 0:
    491.         weight_file = os.path.join(
    492.             OUTPUT_DIR,
    493.             os.path.join(run_name, 'weights%02d.h5' % (start_epoch - 1)))
    494.         model.load_weights(weight_file)
    495.     # captures output of softmax so we can decode the output during visualization
    496.     test_func = K.function([input_data], [y_pred])
    498.     viz_cb = VizCallback(run_name, test_func, img_gen.next_val())
    500.     model.fit_generator(
    501.         generator=img_gen.next_train(),
    502.         steps_per_epoch=(words_per_epoch - val_words) // minibatch_size,
    503.         epochs=stop_epoch,
    504.         validation_data=img_gen.next_val(),
    505.         validation_steps=val_words // minibatch_size,
    506.         callbacks=[viz_cb, img_gen],
    507.         initial_epoch=start_epoch)
    510. if __name__ == '__main__':
    511.     run_name = datetime.datetime.now().strftime('%Y:%m:%d:%H:%M:%S')
    512.     train(run_name, 0, 20, 128)
    513.     # increase to wider images and start at epoch 20.
    514.     # The learned weights are reloaded
    515.     train(run_name, 20, 25, 512)