Ideal PodiumBuilding a Complete OCR Engine From Scratch In Python
Vikas Kumar Ojha
MachineLearning@Avalara Ex-Samsung
OCR is one of the most common applications of Machine Learning that we see in our lives. OCR has plenty of applications in the real world from scanning price tags in malls to license plates on roads. Due to advancements in deep learning techniques over the years the OCR has become quite accurate especially for documents. OCR is being used in a lot of places to digitize documents to make them more usable and search-friendly.
Around a year ago when I was in the final year of my college I was trying to make something related to performing OCR on document images. In my case the documents were PDFs. I first converted the PDF into images and tried to perform OCR on the document images using the popular OCR engine Tesseract. But I realized although it works well it still makes mistakes even on quite clear images and tesseract cannot be made to inference with android studio. Also, I wanted to learn more about OCR hence I decided to make my own. For this, I did around a month of research conducting different experiments for detecting the text lines which was essential for doing OCR and finally I choose the segmentation approach explained in my previous article.
In this article, I will use the concept from the previous article to develop a complete OCR pipeline that will give you enough idea to design your OCR system. Hence you are recommended to read that article in order to properly understand the concept.
Overview
If you have read my previous article or you have some background in computer vision/text recognition then you would be knowing that most text recognition algorithms can recognize single words or lines. Hence for performing recognition on the document/book page it is necessary to separate the page into lines and then the lines into individual words and then perform recognition on them or you could directly perform recognition on the line images. In this article for the sake of simplicity, we will first segment our page into lines and then lines into individual words and perform recognition on the word images.
Line segmentation
To understand the theory behind line segmentation and how can we train a UNet model to segment lines in a document image please read my previous article. I have used the same technique to segment lines into words too. In this article, we will just load the line segmentation model that we trained in the previous article and perform line segmentation. The pretrained model weights can be downloaded from this link.
So let’s create a python module SegmentPage which loads the model and has the utility function which segments the page into lines.
#Importing modules.
from keras.layers import *
from keras.models import Model
from keras.optimizers import Adam
import matplotlib.pyplot as plt
import numpy as np
import cv2
def unet(pretrained_weights = None,input_size = (512,512,1)):
#Creates the UNet model same as the model that was trained to segment page into lines.
inputs = Input(input_size)
conv1 = Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(inputs)
conv1 = Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv1)
pool1 = MaxPooling2D(pool_size=(2, 2))(conv1)
conv2 = Conv2D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(pool1)
conv2 = Conv2D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv2)
pool2 = MaxPooling2D(pool_size=(2, 2))(conv2)
conv3 = Conv2D(256, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(pool2)
conv3 = Conv2D(256, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv3)
pool3 = MaxPooling2D(pool_size=(2, 2))(conv3)
conv4 = Conv2D(512, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(pool3)
conv4 = Conv2D(512, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv4)
drop4 = Dropout(0.5)(conv4)
pool4 = MaxPooling2D(pool_size=(2, 2))(drop4)
conv5 = Conv2D(1024, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(pool4)
conv5 = Conv2D(1024, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv5)
drop5 = Dropout(0.5)(conv5)
up6 = Conv2D(512, 2, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(UpSampling2D(size = (2,2))(drop5))
merge6 = concatenate([drop4,up6], axis = 3)
conv6 = Conv2D(512, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(merge6)
conv6 = Conv2D(512, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv6)
up7 = Conv2D(256, 2, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(UpSampling2D(size = (2,2))(conv6))
merge7 = concatenate([conv3,up7], axis = 3)
conv7 = Conv2D(256, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(merge7)
conv7 = Conv2D(256, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv7)
up8 = Conv2D(128, 2, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(UpSampling2D(size = (2,2))(conv7))
merge8 = concatenate([conv2,up8], axis = 3)
conv8 = Conv2D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(merge8)
conv8 = Conv2D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv8)
up9 = Conv2D(64, 2, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(UpSampling2D(size = (2,2))(conv8))
merge9 = concatenate([conv1,up9], axis = 3)
conv9 = Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(merge9)
conv9 = Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv9)
conv9 = Conv2D(2, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv9)
conv10 = Conv2D(1, 1, activation = 'sigmoid')(conv9)
model = Model(inputs,conv10)
model.compile(optimizer = Adam(lr = 1e-4), loss = 'binary_crossentropy', metrics = ['accuracy'])
if(pretrained_weights):
model.load_weights(pretrained_weights)
return model
model=unet()
model.load_weights('./text_seg_model.h5')
line_img_array=[]
def segment_into_lines(filename):
#Loading the image and performing thresholding on it and then resizing.
img=cv2.imread(f'{filename}',0)
ret,img=cv2.threshold(img,150,255,cv2.THRESH_BINARY_INV)
img=cv2.resize(img,(512,512))
#Expanding the dimension to account for the batch dimension.
img= np.expand_dims(img,axis=-1)
#Expanding dimension along channel axis.
img=np.expand_dims(img,axis=0)
#Predict the segmentation mask.
pred=model.predict(img)
#Remove the batch and channel dimension for performing the binarization.
pred=np.squeeze(np.squeeze(pred,axis=0),axis=-1)
#Performing the binarization of the predicted mask for contour detection.
coordinates=[]
img = cv2.normalize(src=pred, dst=None, alpha=0, beta=255, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8UC1)
cv2.threshold(img,0,255,cv2.THRESH_BINARY+cv2.THRESH_OTSU,img)
#Opening the original image to get the original dimension information.
ori_img=cv2.imread(f'{filename}',0)
(H, W) = ori_img.shape[:2]
(newW, newH) = (512, 512)
rW = W / float(newW)
rH = H / float(newH)
#Contour detection and bouding box generation.
contours, hier = cv2.findContours(img, cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_NONE)
for c in contours:
# get the bounding rect
x, y, w, h = cv2.boundingRect(c)
#cv2.rectangle(ori_img, (int(x*rW), int(y*rH)), (int((x+w)*rW),int((y+h)*rH)), (255,0,0), 1)
coordinates.append((int(x*rW),int(y*rH),int((x+w)*rW),int((y+h)*rH)))
#cv2.imwrite("output.jpg",ori_img)
#Cropping the lines from the original image using the bouding boxes generated above.
for i in range(len(coordinates)-1,-1,-1):
coors=coordinates[i]
p_img=ori_img[coors[1]:coors[3],coors[0]:coors[2]].copy()
line_img_array.append(p_img)
return line_img_array
Here we are creating a UNet model same as that we used to train for line segmentation and we load the trained weights. Then we have defined a function called segment_into_lines which accepts the link of the page image. This function reads the file in grayscale format and then it performs the binarization on the image. After binarization, it resizes the image to 512x512 as that is the input size of our line segmentation model. Then it expands the dimension across axis 0 to account for the batch dimension. It also expands the dimension along the channel axis by default because when the image is opened in the grayscale format it doesn’t have a color channel. Then it passed it to the model for predicting the line segmentation mask.
Again it removes the batch dimension and the channel dimension because again we have to perform binarization of the segmentation mask to perform contour detection on it. After performing binarization we perform contour detection on the image using the findCountours function. After having the contours we pass them to OpenCV's boundingRect function which approximates the rectangular boxes for those contours. We also scale these bounding boxes relative to the size of the original image as the original image is bigger and more clear. Finally, we use these bounding boxes to crop line images from the image and return them appended in a list.
Below is the output after running calling this function on a page.
Word Segmentation
Now we have the cropped lines of the page. Further, we need to segment these lines into words that will be recognized by our text recognition algorithm. We could also train the model to recognize the whole line. But for the sake of simplicity, I decided to segment the lines further into words. So let’s train a model which segments these lines into words.
Model Training
To train the model for word segmentation I have to use the same architecture from the above section. For generating the data I used the previous model to segment some random pages into lines and then I created the annotation for them using the pixel annotation tool. Below the image is an example of the line image and its annotation.
But as we remember from the previous section our model accepts 512x512 images hence we need to pad the image to 512x512. Below are an example of images after padding.
Having understood the theory let's dive into the code and train our model. I used google colab to train the model. Since I have already explained the concept in my previous article, I will explain the concepts at a high level only.
Downloading the dataset
%tensorflow_version 1.x
!gdown --id 155wulmctFNkAMYFp770XLCWm1FcYEsfy
!unzip WordSegData2 Using magic function %tensorflow_version 1. x, I am selecting tensorflow version 1.15 on colab since at that time tensorflow 2. x was new. The gdown command downloads the dataset with its id. The unzip command is then used to extract the dataset.
Update: As of now tensorflow 1.x is not available anymore and tf.keras can be used instead. Since Keras is merged into tensorflow except few imports everything should work fine.
Library imports and image path initialization
#Importing required libraries.
import numpy as np
import matplotlib.pyplot as plt
import cv2
import os
from keras.layers import *
from keras.models import Model
from keras.optimizers import Adam
from keras.callbacks import ModelCheckpoint
import random
import math
image_list=os.listdir('/content/Dataset1/img/')
image_list=[filename.split(".")[0]for filename in image_list]
Here we are doing library imports and listing the images in the dataset directory. Since our mask has the image name as the source_image except that it ends with _mask hence we don’t need to list it separately.
Defining the visualization function, pad functions, and batch generator
#Importing required libraries.
import numpy as np
import matplotlib.pyplot as plt
import cv2
import os
from keras.layers import *
from keras.models import Model
from keras.optimizers import Adam
import random
import math
image_list=os.listdir('/content/Dataset1/img/')
image_list=[filename.split(".")[0]for filename in image_list]
def visualize(img,seg_img):
"""
Visualizes image
"""
plt.figure(figsize=(20,20))
plt.subplot(1,2,1)
plt.imshow(img)
plt.title('Image')
plt.subplot(1,2,2)
plt.imshow(seg_img,cmap='gray')
plt.title('Segmented Image')
plt.show()
def roundup(x):
return int(math.ceil(x / 10.0)) * 10
def get_segmented_img(img,n_classes):
"""
Loads in the segmented image and create suitable segmentation label.
"""
seg_labels=np.zeros((512,512,1))
img=cv2.resize(img,(512,512))
img=img[:,:,0]
cl_list=[0,24]
seg_labels[:,:,0]=(img!=0).astype(int)
return seg_labels
def pad_img(img):
old_h,old_w=img.shape[0],img.shape[1]
#Pad the height.
#If height is less than 512 then pad to 512
if old_h<512:
to_pad=np.ones((512-old_h,old_w))*255
img=np.concatenate((img,to_pad))
new_height=512
else:
#If height >512 then pad to nearest 10.
to_pad=np.ones((roundup(old_h)-old_h,old_w))*255
img=np.concatenate((img,to_pad))
new_height=roundup(old_h)
#Pad the width.
if old_w<512:
to_pad=np.ones((new_height,512-old_w))*255
img=np.concatenate((img,to_pad),axis=1)
new_width=512
else:
to_pad=np.ones((new_height,roundup(old_w)-old_w))*255
img=np.concatenate((img,to_pad),axis=1)
new_width=roundup(old_w)-old_w
return img
def pad_seg(img):
old_h,old_w=img.shape[0],img.shape[1]
#Pad the height.
#If height is less than 512 then pad to 512
if old_h<512:
to_pad=np.zeros((512-old_h,old_w))
img=np.concatenate((img,to_pad))
new_height=512
else:
#If height >512 then pad to nearest 10.
to_pad=np.zeros((roundup(old_h)-old_h,old_w))
img=np.concatenate((img,to_pad))
new_height=roundup(old_h)
#Pad the width.
if old_w<512:
to_pad=np.zeros((new_height,512-old_w))
img=np.concatenate((img,to_pad),axis=1)
new_width=512
else:
to_pad=np.zeros((new_height,roundup(old_w)-old_w))
img=np.concatenate((img,to_pad),axis=1)
new_width=roundup(old_w)-old_w
return img
def batch_generator(filelist,n_classes,batch_size):
while True:
X=[]
Y=[]
for i in range(batch_size):
fn=random.choice(filelist)
img=cv2.imread(f'/content/Dataset1/img/{fn}.jpg',0)
img=pad_img(img)
ret,img=cv2.threshold(img,150,255,cv2.THRESH_BINARY_INV)
img=cv2.resize(img,(512,512))
img=np.expand_dims(img,axis=-1)
#img = np.stack((img,)*3, axis=-1)
img=img/255
seg=cv2.imread(f'/content/Dataset1/mask/{fn}_mask.png',0)
seg=pad_seg(seg)
seg=cv2.resize(seg,(512,512))
seg = np.stack((seg,)*3, axis=-1)
seg=get_segmented_img(seg,n_classes)
X.append(img)
Y.append(seg)
yield np.array(X),np.array(Y)
The visualize and get_segmented_img are the same as in my previous article. The pad_seg and pad_img functions perform padding on the segmentation and source image respectively and make them into 512x512 size. The batch_generator function is also almost the same as my previous article except that it pads the image and the segmentation before doing any operations.
Splitting the data into train and test splits
random.shuffle(image_list)
file_train=image_list[0:int(0.75*len(image_list))]
file_test=image_list[int(0.75*len(image_list)):]
Defining the model and callbacks
def unet(pretrained_weights = None,input_size = (512,512,1)):
inputs = Input(input_size)
conv1 = Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(inputs)
conv1 = Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv1)
pool1 = MaxPooling2D(pool_size=(2, 2))(conv1)
conv2 = Conv2D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(pool1)
conv2 = Conv2D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv2)
pool2 = MaxPooling2D(pool_size=(2, 2))(conv2)
conv3 = Conv2D(256, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(pool2)
conv3 = Conv2D(256, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv3)
pool3 = MaxPooling2D(pool_size=(2, 2))(conv3)
conv4 = Conv2D(512, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(pool3)
conv4 = Conv2D(512, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv4)
drop4 = Dropout(0.5)(conv4)
pool4 = MaxPooling2D(pool_size=(2, 2))(drop4)
conv5 = Conv2D(1024, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(pool4)
conv5 = Conv2D(1024, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv5)
drop5 = Dropout(0.5)(conv5)
up6 = Conv2D(512, 2, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(UpSampling2D(size = (2,2))(drop5))
merge6 = concatenate([drop4,up6], axis = 3)
conv6 = Conv2D(512, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(merge6)
conv6 = Conv2D(512, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv6)
up7 = Conv2D(256, 2, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(UpSampling2D(size = (2,2))(conv6))
merge7 = concatenate([conv3,up7], axis = 3)
conv7 = Conv2D(256, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(merge7)
conv7 = Conv2D(256, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv7)
up8 = Conv2D(128, 2, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(UpSampling2D(size = (2,2))(conv7))
merge8 = concatenate([conv2,up8], axis = 3)
conv8 = Conv2D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(merge8)
conv8 = Conv2D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv8)
up9 = Conv2D(64, 2, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(UpSampling2D(size = (2,2))(conv8))
merge9 = concatenate([conv1,up9], axis = 3)
conv9 = Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(merge9)
conv9 = Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv9)
conv9 = Conv2D(2, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv9)
conv10 = Conv2D(1, 1, activation = 'sigmoid')(conv9)
model = Model(inputs,conv10)
model.compile(optimizer = Adam(lr = 1e-4), loss = 'binary_crossentropy', metrics = ['accuracy'])
#model.summary()
if(pretrained_weights):
model.load_weights(pretrained_weights)
return model
model=unet()
mc = ModelCheckpoint('weights{epoch:08d}.h5',
save_weights_only=True, period=1)
Training the model
model.fit_generator(batch_generator(file_train,2,2),epochs=5,steps_per_epoch=1000,validation_data=batch_generator(file_test,2,2),
validation_steps=400,callbacks=[mc],shuffle=1)
Finally, we trained the model with batch size 2 for around 5 epochs with callbacks we defined above.
Testing the model
#Prediction.
file_test='/content/Dataset1/img/lineA1.jpg'
img=cv2.imread(f'{file_test}',0)
img=pad_img(img)
ret,img=cv2.threshold(img,150,255,cv2.THRESH_BINARY_INV)
img=cv2.resize(img,(512,512))
img=np.expand_dims(img,axis=-1)
#img = np.stack((img,)*3, axis=-1)
img=img/255
img=np.expand_dims(img,axis=0)
pred=model.predict(img)
pred=np.squeeze(np.squeeze(pred,axis=0),axis=-1)
plt.imshow(pred,cmap='gray')
plt.imsave('test_img_mask.JPG',pred)
img = cv2.imread('/content/test_img_mask.JPG',0)
cv2.threshold(img,0,255,cv2.THRESH_BINARY+cv2.THRESH_OTSU,img)
ori_img=cv2.imread(f'{file_test}',0)
ori_img=pad_img(ori_img)
(H, W) = ori_img.shape[:2]
(newW, newH) = (512, 512)
rW = W / float(newW)
rH = H / float(newH)
ori_img_copy=np.stack((ori_img,)*3, axis=-1)
contours, hier = cv2.findContours(img, cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_NONE)
for c in contours:
# get the bounding rect
x, y, w, h = cv2.boundingRect(c)
# draw a white rectangle to visualize the bounding rect
cv2.rectangle(ori_img_copy, (int(x*rW), int(y*rH)), (int((x+w)*rW),int((y+h)*rH)), (255,0,0), 1)
#coordinates.append([x,y,(x+w),(y+h)])
#cv2.drawContours(img, contours, -1, (255, 255, 0), 1)
cv2.imwrite("output.png",ori_img_copy)
Here we have done prediction using the same technique as in the line segmentation section.
Following was the output:
Word Recognition
Now the final piece which is missing from our puzzle is the text recognizer. For recognizing the cropped word images we will be training a CRNN model which performs the text recognition. The CRNN model is a neural network that consists of convolutional layers which are followed by RNN(LSTM) layers. In the end, we have a CTC layer that predicts the text given the image. The convolutional layers extract the visual information while the recurrent layers capture the temporal information of the text in the image. The CTC layer does the prediction of the characters. To understand more about the working of the CRNN model please read this excellent article.
Training the CRNN model
Here we will be training a CRNN model which takes a 128x32 image as input. Since the aim of this experiment was just to figure out the OCR training pipeline and the recognition was to be performed on clear PDF images I decided to use synthetic data instead of downloading or procuring real-world data. For generating the synthetic data I wrote my own data generator which can be found here.
Cloning the data generator and generating images
!git clone https://github.com/VikasOjha666/Data-generator-for-CRNN.git
%cd Data-generator-for-CRNN
!mkdir images/
!python generate_data.py --n_samples 300000 Here we are cloning my data generator repository and changing our directory into it. Then we are making a folder called images. Finally, we generate around 300000 images. One thing that should be noted is that the data generator I have created contains the font files similar to the text I want to recognize I used certain font identification websites to list out fonts similar to the page I wanted to identify. Every text recognition engine is biased towards the font it has seen during the training and that's the reason we use a lot of different font images to train the model in real life.
Importing necessary libraries
import fnmatch
import cv2
import numpy as np
import string
import time
from keras.preprocessing.sequence import pad_sequences
from keras.layers import Dense, LSTM, Reshape, BatchNormalization, Input, Conv2D, MaxPool2D, Lambda, Bidirectional
from keras.models import Model
from keras.activations import relu, sigmoid, softmax
import keras.backend as K
from keras.utils import to_categorical
from keras.callbacks import ModelCheckpoint
import os
import tensorflow as tf
import random
import matplotlib.pyplot as plt
from PIL import Image
from sklearn.utils import shuffle
Defining character list and some utility functions
char_list = string.ascii_letters+string.digits
def encode_to_labels(txt):
# encoding each output word into digits
dig_lst = []
for index, char in enumerate(txt):
try:
dig_lst.append(char_list.index(char))
except:
print(char)
return dig_lst
def find_dominant_color(image):
#Resizing parameters
width, height = 150,150
image = image.resize((width, height),resample = 0)
#Get colors from image object
pixels = image.getcolors(width * height)
#Sort them by count number(first element of tuple)
sorted_pixels = sorted(pixels, key=lambda t: t[0])
#Get the most frequent color
dominant_color = sorted_pixels[-1][1]
return dominant_color
def preprocess_img(img, imgSize):
"put img into target img of size imgSize and normalize gray-values"
# In case of black images with no text just use black image instead.
if img is None:
img = np.zeros([imgSize[1], imgSize[0]])
print("Image None!")
# create target image and copy sample image into it
(wt, ht) = imgSize
(h, w) = img.shape
fx = w / wt
fy = h / ht
f = max(fx, fy)
newSize = (max(min(wt, int(w / f)), 1),
max(min(ht, int(h / f)), 1)) # scale according to f (result at least 1 and at most wt or ht)
img = cv2.resize(img, newSize, interpolation=cv2.INTER_CUBIC) # INTER_CUBIC interpolation best approximate the pixels image
# see this https://stackoverflow.com/a/57503843/7338066
most_freq_pixel=find_dominant_color(Image.fromarray(img))
target = np.ones([ht, wt]) * most_freq_pixel
target[0:newSize[1], 0:newSize[0]] = img
img = target
return img
Here we are defining the characters list that we want to recognize using our text recognition model. We are using a string module to get the list of characters. We have defined two utility functions find_dominant_color and preprocess_img. The find_dominant_color function is used to find the background color in the image. Since in any image the number of pixels of background will always be more than the text color pixels hence the pixels having the greatest occurrence in the image are the pixels related to the background and hence the color of these pixels is the background color. We have created this function so that we can pad our input image to 128x32 and avoid resizing as resizing have effects on accuracy. The preprocess_img function performs the actual resizing operation. The encode_to_labels function converts the text to indexes as neural networks processes numbers and not text.
Loading the generated images and preprocessing them
training_img = []
training_txt = []
train_input_length = []
train_label_length = []
orig_txt = []
#lists for validation dataset
valid_img = []
valid_txt = []
valid_input_length = []
valid_label_length = []
valid_orig_txt = []
max_label_len = 0
annot=open('/content/Data-generator-for-CRNN/annotation.txt','r').readlines()
imagenames=[]
txts=[]
for cnt in annot:
filename,txt=cnt.split(',')[0],cnt.split(',')[1].split('\n')[0]
imagenames.append(filename)
txts.append(txt)
c = list(zip(imagenames, txts))
random.shuffle(c)
imagenames, txts = zip(*c)
for i in range(len(imagenames)):
img = cv2.imread('/content/Data-generator-for-CRNN/images/'+imagenames[i],0)
img=preprocess_img(img,(128,32))
img=np.expand_dims(img,axis=-1)
img = img/255.
txt = txts[i]
# compute maximum length of the text
if len(txt) > max_label_len:
max_label_len = len(txt)
# split the 150000 data into validation and training dataset as 10% and 90% respectively
if i%10 == 0:
valid_orig_txt.append(txt)
valid_label_length.append(len(txt))
valid_input_length.append(31)
valid_img.append(img)
valid_txt.append(encode_to_labels(txt))
else:
orig_txt.append(txt)
train_label_length.append(len(txt))
train_input_length.append(31)
training_img.append(img)
training_txt.append(encode_to_labels(txt))
# break the loop if total data is 150000
if i == 150000:
flag = 1
break
i+=1
Here we are loading the images and processing them. We are also processing the text and encoding them. While preprocessing we are resizing each image to 128x32. After preprocessing we are appending images to training_img and valid_img lists respectively. We are also appending the preprocessed text to training_txt and valid_txt respectively. Additionally, we are appending 31 to training_label_length and valid_input_length because the maximum character our model can recognize is 32 as RNN has 32 timesteps in our model. If we want to recognize more we can increase the dimension of the image. Since Google colab provides around 12 GB RAM we only keep around 150000 examples.
Padding each sequence to the max length of the text in the dataset
#pad each output label to maximum text length
train_padded_txt = pad_sequences(training_txt, maxlen=max_label_len, padding='post', value = len(char_list))
valid_padded_txt = pad_sequences(valid_txt, maxlen=max_label_len, padding='post', value = len(char_list))
Here we are padding the encoded texts to the max length of the text in the dataset to make the length of each sequence in the batch equal. Post padding method is being used here and the value which is being used is len(char_list).
Define the CRNN model
#CRNN model
inputs = Input(shape=(32,128,1))
# convolution layer with kernel size (3,3)
conv_1 = Conv2D(64, (3,3), activation = 'relu', padding='same')(inputs)
# poolig layer with kernel size (2,2)
pool_1 = MaxPool2D(pool_size=(2, 2), strides=2)(conv_1)
conv_2 = Conv2D(128, (3,3), activation = 'relu', padding='same')(pool_1)
pool_2 = MaxPool2D(pool_size=(2, 2), strides=2)(conv_2)
conv_3 = Conv2D(256, (3,3), activation = 'relu', padding='same')(pool_2)
conv_4 = Conv2D(256, (3,3), activation = 'relu', padding='same')(conv_3)
# poolig layer with kernel size (2,1)
pool_4 = MaxPool2D(pool_size=(2, 1))(conv_4)
conv_5 = Conv2D(512, (3,3), activation = 'relu', padding='same')(pool_4)
# Batch normalization layer
batch_norm_5 = BatchNormalization()(conv_5)
conv_6 = Conv2D(512, (3,3), activation = 'relu', padding='same')(batch_norm_5)
batch_norm_6 = BatchNormalization()(conv_6)
pool_6 = MaxPool2D(pool_size=(2, 1))(batch_norm_6)
conv_7 = Conv2D(512, (2,2), activation = 'relu')(pool_6)
squeezed = Lambda(lambda x: K.squeeze(x, 1))(conv_7)
# bidirectional LSTM layers with units=128
blstm_1 = Bidirectional(LSTM(128, return_sequences=True, dropout = 0.2))(squeezed)
blstm_2 = Bidirectional(LSTM(128, return_sequences=True, dropout = 0.2))(blstm_1)
outputs = Dense(len(char_list)+1, activation = 'softmax')(blstm_2)
# model to be used at test time
act_model = Model(inputs, outputs)
Here we are defining our CRNN model to recognize the model. Here you might have observed that the model doesn’t have any CTC layer here. That's because we only use the CTC layer to do the training. For inference, we would be using this model. So let's define the CTC layer next and train this model using the CTC layer. This model is called act_model here in the code.
Defining CTC based model
labels = Input(name='the_labels', shape=[max_label_len], dtype='float32')
input_length = Input(name='input_length', shape=[1], dtype='int64')
label_length = Input(name='label_length', shape=[1], dtype='int64')
def ctc_lambda_func(args):
#Defining the CTC loss.
y_pred, labels, input_length, label_length = args
return K.ctc_batch_cost(labels, y_pred, input_length, label_length)
#CTC layer declaration using lambda.
loss_out = Lambda(ctc_lambda_func, output_shape=(1,), name='ctc')([outputs, labels, input_length, label_length])
#Including the CTC layer to train the model.
model = Model(inputs=[inputs, labels, input_length, label_length], outputs=loss_out) Here we are declaring the CTC loss function and CTC layer using lambda layers of Keras. Finally, we combine the CTC layer with our model.
Compiling the model and defining model checkpointing callback
model.compile(loss={'ctc': lambda y_true, y_pred: y_pred}, optimizer = 'adam')
filepath="/content/best_model.hdf5"
checkpoint = ModelCheckpoint(filepath=filepath, monitor='val_loss', verbose=1, save_best_only=True, mode='auto')
callbacks_list = [checkpoint]
Converting the lists to NumPy arrays
training_img = np.array(training_img)
train_input_length = np.array(train_input_length)
train_label_length = np.array(train_label_length)
valid_img = np.array(valid_img)
valid_input_length = np.array(valid_input_length)
valid_label_length = np.array(valid_label_length) Here we are converting the lists which contain images and labels to NumPy arrays.
Training the model
batch_size = 256
epochs = 15
model.fit(x=[training_img, train_padded_txt, train_input_length, train_label_length],
y=np.zeros(len(training_img)),
batch_size=batch_size, epochs = epochs,
validation_data = ([valid_img, valid_padded_txt, valid_input_length, valid_label_length],
[np.zeros(len(valid_img))]), verbose = 1, callbacks = callbacks_list)
Loading the trained weights and making a prediction
# load the saved best model weights
act_model.load_weights('best_model.hdf5')
# predict outputs on validation images
prediction = act_model.predict(valid_img[10:20])
# use CTC decoder
out = K.get_value(K.ctc_decode(prediction, input_length=np.ones(prediction.shape[0])*prediction.shape[1],
greedy=True)[0][0])
# see the results
i = 10
for x in out:
print("original_text = ", valid_orig_txt[i])
print("predicted text = ", end = '')
for p in x:
if int(p) != -1:
print(char_list[int(p)], end = '')
print('\n')
i+=1 Here we are loading the model and making predictions. Here it can be observed that the weights are being loaded into the act_model because that is the model which is used for prediction. The CTC layer is used for training and doesn’t have weights.
Creating the OCR pipeline
Now we have built all the building blocks for OCR. Let's use those models and finally build our OCR.
Defining utility function for word segmentation
We have already defined the function segment_into_lines in SegmentPage.py which segments the page into lines. Now we define another file SegmentLine.py having a function segment_into_words which segments the lines into words. This function takes the line image and the line index and it returns an array containing cropped words and the same index. We have used the same code from the prediction section of word segmentation.
import numpy as np
import cv2
import os
from keras.layers import *
from keras.models import Model
from keras.optimizers import Adam
import matplotlib.pyplot as plt
from spellchecker import SpellChecker
import math
from PIL import Image
def find_dominant_color(image):
#Resizing parameters
width, height = 150,150
image = image.resize((width, height),resample = 0)
#Get colors from image object
pixels = image.getcolors(width * height)
#Sort them by count number(first element of tuple)
sorted_pixels = sorted(pixels, key=lambda t: t[0])
#Get the most frequent color
dominant_color = sorted_pixels[-1][1]
return dominant_color
def preprocess_img(img, imgSize):
"put img into target img of size imgSize, transpose for TF and normalize gray-values"
# there are damaged files in IAM dataset - just use black image instead
if img is None:
img = np.zeros([imgSize[1], imgSize[0]])
print("Image None!")
# create target image and copy sample image into it
(wt, ht) = imgSize
(h, w) = img.shape
fx = w / wt
fy = h / ht
f = max(fx, fy)
newSize = (max(min(wt, int(w / f)), 1),
max(min(ht, int(h / f)), 1)) # scale according to f (result at least 1 and at most wt or ht)
img = cv2.resize(img, newSize, interpolation=cv2.INTER_CUBIC) # INTER_CUBIC interpolation best approximate the pixels image
# see this https://stackoverflow.com/a/57503843/7338066
most_freq_pixel=find_dominant_color(Image.fromarray(img))
target = np.ones([ht, wt]) * most_freq_pixel
target[0:newSize[1], 0:newSize[0]] = img
img = target
return img
def pad_img(img):
old_h,old_w=img.shape[0],img.shape[1]
#Pad the height.
#If height is less than 512 then pad to 512
if old_h<512:
to_pad=np.ones((512-old_h,old_w))*255
img=np.concatenate((img,to_pad))
new_height=512
else:
#If height >512 then pad to nearest 10.
to_pad=np.ones((roundup(old_h)-old_h,old_w))*255
img=np.concatenate((img,to_pad))
new_height=roundup(old_h)
#Pad the width.
if old_w<512:
to_pad=np.ones((new_height,512-old_w))*255
img=np.concatenate((img,to_pad),axis=1)
new_width=512
else:
to_pad=np.ones((new_height,roundup(old_w)-old_w))*255
img=np.concatenate((img,to_pad),axis=1)
new_width=roundup(old_w)-old_w
return img
def roundup(x):
return int(math.ceil(x / 10.0)) * 10
def unet(pretrained_weights = None,input_size = (512,512,1)):
inputs = Input(input_size)
conv1 = Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(inputs)
conv1 = Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv1)
pool1 = MaxPooling2D(pool_size=(2, 2))(conv1)
conv2 = Conv2D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(pool1)
conv2 = Conv2D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv2)
pool2 = MaxPooling2D(pool_size=(2, 2))(conv2)
conv3 = Conv2D(256, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(pool2)
conv3 = Conv2D(256, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv3)
pool3 = MaxPooling2D(pool_size=(2, 2))(conv3)
conv4 = Conv2D(512, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(pool3)
conv4 = Conv2D(512, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv4)
drop4 = Dropout(0.5)(conv4)
pool4 = MaxPooling2D(pool_size=(2, 2))(drop4)
conv5 = Conv2D(1024, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(pool4)
conv5 = Conv2D(1024, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv5)
drop5 = Dropout(0.5)(conv5)
up6 = Conv2D(512, 2, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(UpSampling2D(size = (2,2))(drop5))
merge6 = concatenate([drop4,up6], axis = 3)
conv6 = Conv2D(512, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(merge6)
conv6 = Conv2D(512, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv6)
up7 = Conv2D(256, 2, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(UpSampling2D(size = (2,2))(conv6))
merge7 = concatenate([conv3,up7], axis = 3)
conv7 = Conv2D(256, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(merge7)
conv7 = Conv2D(256, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv7)
up8 = Conv2D(128, 2, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(UpSampling2D(size = (2,2))(conv7))
merge8 = concatenate([conv2,up8], axis = 3)
conv8 = Conv2D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(merge8)
conv8 = Conv2D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv8)
up9 = Conv2D(64, 2, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(UpSampling2D(size = (2,2))(conv8))
merge9 = concatenate([conv1,up9], axis = 3)
conv9 = Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(merge9)
conv9 = Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv9)
conv9 = Conv2D(2, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv9)
conv10 = Conv2D(1, 1, activation = 'sigmoid')(conv9)
model = Model(inputs,conv10)
model.compile(optimizer = Adam(lr = 1e-4), loss = 'binary_crossentropy', metrics = ['accuracy'])
if(pretrained_weights):
model.load_weights(pretrained_weights)
return model
model=unet()
model.load_weights('./word_seg_model.h5')
def sort_word(wordlist):
wordlist.sort(key=lambda x:x[0])
return wordlist
def segment_into_words(line_img,idx):
"""This function takes in the line image and line index returns word images and the reference
of line they belong to."""
img=pad_img(line_img)
ori_img=img.copy()
#ori_img=np.stack((ori_img,)*3, axis=-1)
ret,img=cv2.threshold(img,150,255,cv2.THRESH_BINARY_INV)
img=cv2.resize(img,(512,512))
img=np.expand_dims(img,axis=-1)
img=img/255
img=np.expand_dims(img,axis=0)
seg_pred=model.predict(img)
seg_pred=np.squeeze(np.squeeze(seg_pred,axis=0),axis=-1)
seg_pred=cv2.normalize(src=seg_pred, dst=None, alpha=0, beta=255, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8UC1)
cv2.threshold(seg_pred,0,255,cv2.THRESH_BINARY+cv2.THRESH_OTSU,seg_pred)
contours, hier = cv2.findContours(seg_pred, cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_NONE)
(H, W) = ori_img.shape[:2]
(newW, newH) = (512, 512)
rW = W / float(newW)
rH = H / float(newH)
coordinates=[]
for c in contours:
# get the bounding rect
x, y, w, h = cv2.boundingRect(c)
# draw a white rectangle to visualize the bounding rect
# cv2.rectangle(ori_img, (int(x*rW), int(y*rH)), (int((x+w)*rW),int((y+h)*rH)), (255,0,0), 1)
coordinates.append((int(x*rW),int(y*rH),int((x+w)*rW),int((y+h)*rH)))
coordinates=sort_word(coordinates) #Sorting according to x-coordinates.
word_counter=0
word_array=[]
line_indicator=[]
for (x1,y1,x2,y2) in coordinates:
word_img=ori_img[y1:y2,x1:x2]
word_img=preprocess_img(word_img,(128,32))
word_img=np.expand_dims(word_img,axis=-1)
word_array.append(word_img)
line_indicator.append(idx)
return line_indicator,word_array
Creating utility function to recognize word images
import cv2
import numpy as np
import string
import os
#os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
from keras.preprocessing.sequence import pad_sequences
from keras.layers import Dense, LSTM, Reshape, BatchNormalization, Input, Conv2D, MaxPool2D, Lambda, Bidirectional,Dropout
from keras.models import Model
import keras.backend as K
from spellchecker import SpellChecker
import tensorflow as tf
from PIL import Image
spell = SpellChecker()
char_list = string.ascii_letters+string.digits
inputs = Input(shape=(32,128,1))
conv_1 = Conv2D(64, (3,3), activation = 'relu', padding='same')(inputs)
pool_1 = MaxPool2D(pool_size=(2, 2), strides=2)(conv_1)
conv_2 = Conv2D(128, (3,3), activation = 'relu', padding='same')(pool_1)
pool_2 = MaxPool2D(pool_size=(2, 2), strides=2)(conv_2)
conv_3 = Conv2D(256, (3,3), activation = 'relu', padding='same')(pool_2)
conv_4 = Conv2D(256, (3,3), activation = 'relu', padding='same')(conv_3)
pool_4 = MaxPool2D(pool_size=(2, 1))(conv_4)
conv_5 = Conv2D(512, (3,3), activation = 'relu', padding='same')(pool_4)
batch_norm_5 = BatchNormalization()(conv_5)
conv_6 = Conv2D(512, (3,3), activation = 'relu', padding='same')(batch_norm_5)
batch_norm_6 = BatchNormalization()(conv_6)
pool_6 = MaxPool2D(pool_size=(2, 1))(batch_norm_6)
conv_7 = Conv2D(512, (2,2), activation = 'relu')(pool_6)
squeezed = Lambda(lambda x: K.squeeze(x, 1))(conv_7)
blstm_1 = Bidirectional(LSTM(128, return_sequences=True, dropout = 0.2))(squeezed)
blstm_2 = Bidirectional(LSTM(128, return_sequences=True, dropout = 0.2))(blstm_1)
outputs = Dense(len(char_list)+1, activation = 'softmax')(blstm_2)
act_model = Model(inputs, outputs)
act_model.load_weights('best_model.hdf5')
def recognize_words(line_indicator,word_array,n_lines):
file=open('recognized_texts.txt','w')
line_rec=[]
for listidx in range(n_lines):
line_rec.append([])
predictions=act_model.predict(word_array)
out = K.get_value(K.ctc_decode(predictions, input_length=np.ones(predictions.shape[0])*predictions.shape[1],
greedy=True)[0][0])
lw_idx=0
for wordidxs in out:
word=[]
for char in wordidxs:
if int(char)!=-1:
word.append(char_list[int(char)])
word=spell.correction(''.join(word))
word=''.join(word)
line_rec[line_indicator[lw_idx]].append(word)
lw_idx+=1
for listidx in range(n_lines):
line=' '.join(line_rec[listidx])
print(line)
file.writelines(line+'\n')
file.close()
The file RecognizeWord.py defines the utility function recognize_words that recognizes the text in the lists of word images obtained after segmentation. It takes three parameters line_indicator which contains the array of the line number which is helpful in knowing which line number we are recognizing. The n_lines contain the number of lines we have passed to the function. The parameter word_array is a 2-dimensional array that is actually a list of list of word images(that is each element is a list of words). This function does the prediction using the CRNN model and converts logits into texts using greedy search. This function also utilizes a python library called spellchecker to correct spelling mistakes during recognition. The function saves the recognition to a file recognized_texts.txt.
Let's perform recognition on the page
import numpy as np
from SegmentPage import segment_into_lines
from SegmentLine import segment_into_words
from RecognizeWord import recognize_words
#Open image and segment into lines
line_img_array=segment_into_lines('./test_image.jpg')
#Creating lists to store the line indexes,words list.
full_index_indicator=[]
all_words_list=[]
#Variable to count the total no of lines in page.
len_line_arr=0
#Segment the lines into words and store as arrays.
for idx,im in enumerate(line_img_array):
line_indicator,word_array=segment_into_words(im,idx)
for k in range(len(word_array)):
full_index_indicator.append(line_indicator[k])
all_words_list.append(word_array[k])
len_line_arr+=1
all_words_list=np.array(all_words_list)
#Perform the recognition on list of list of words.
recognize_words(full_index_indicator,all_words_list,len_line_arr) The above file reads test_img.jpg and performs recognition on it.
Below are the outputs
Left is recognition in python console and right is the test image. Cool right. There is a lot of improvement that can be done which is covered in the section below.
Improvements
* The CRNN model can be trained to perform recognition on line level instead of word level.
* Detections of the words can be done once instead of first segmenting into lines and then into words. This can be achieved by simply doing annotation on word level in the page but that would be time-consuming.
* Light text detection models like Differential Binarization could be used. This model is much lighter and faster and is designed explicitly for text recognition. A lot of OCR engines like PaddleOCR, MMOCR, etc uses this algorithm.
* Real-world data with a lot of variations can be used to train detection and recognition models.
* A lot of font variations should be used to make the model robust. Data cleaning can be performed to improve accuracy.
* The model can be trained only with a large amount of data with batching for a long period of time with better hyperparameters.
* As mentioned earlier we have used greed search to decode our prediction. Hence Beam search can be used instead to improve the accuracy.
References and further reading
Congrats if you have reached here you have learned an important concept in computer vision. If you like my article please hit the clap button.
Link to the git repo and notebooks:-
The repo readme also contains the link to the pretrained models.
Comments
No comments yet. Be the first to comment!