Ideal PodiumDetecting Text-lines in a Document Image Using Deep Learning
Vikas Kumar Ojha
MachineLearning@Avalara Ex-Samsung
OCR is one of the most important and popular problems that computer vision has tried to solve for a long period. A huge amount of research has been done in the field of OCR as a result of which a large number of approach has been discovered. Several of them are classical computer vision approach while others are deep learning. Generally, deep learning beat the classical techniques techniques in terms of speed while deep learning techniques beat the classical approaches in terms of accuracy.
A lot of text recognition algorithms have evolved due to a huge amount of research but most of the text recognition algorithms suffer in terms of accuracy while recognizing long paragraphs.Hence in this article I propose an approach to detect text lines in a document image which helps us to break the image into various images each containing a single line of text.
Image Segmentation
Before understanding the approach it is a strict requirement for the user to have solid understanding of the image segmentation.
Image segmentation is a pixel-wise mask for each of the objects present in a image.
Below is a example of image segmentation.
In the above figure, we can see that each object has been given separate colored masks. The color of the mask has no relationship with the object.
Now the obvious question which is to be answered here is that, “how is that mask achieved?” , which can be answered through another definition of image segmentation which is as follows:
Image segmentation is the pixel-wise classification of the image.
From the above definition, you might have got a slight idea of how this mask is obtained. In the case of image segmentation instead of predicting a label from the image, we predict the class for each of the pixels in the image, and the predicted classes are assigned some random fixed color for the purpose of display.
Proposed method
Here I have used the same concept of image segmentation to detect the text lines in a page. Here the text is considered as one class and the background as the other class. The area within the image which contains text has to be classified as 1 while the background area has 0. The below figure will help you better understand it.
In the above figure we can see that the text regions have been given gray colored mask while the background region is black. That is what we intend to predict.
After we have the mask as shown above in the figure,we find the contours of the mask. The contour indexes can be used to get the bounding boxes coordinates for the lines which would help to crop the text regions from the page. Suitable text recognition algorithms can be applied on them thereafter.
Dataset
The dataset which I have used here has been created on my own using the pixel annotation tool. I took screenshot of many pages from various documents and annotated them. This dataset contains around 300 images which took a lot of time for annotation. You can annotate more images or use similar datasets if you can find them on the internet. As you will increase the dataset size and quality you will observe drastic performance improvement. With data of only 300 annotated pages the model was doing quite a decent job but to avoid failure cases you should include more images. The dataset can be downloaded from here.
The Model
I have used the CNN UNet for doing image segmentation. I decided to use UNet because of its structure. In UNet, the image is first downsampled and then upsampled to its original size. In this process, the most important detail of the image is preserved which helps in improving the segmentation. I have also tried using Fully Convolutional Networks(FCNs) but they didn’t perform well on my dataset hence I trained the model from scratch.
Code Implementation
Importing necessary libraries
#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
Listing images to use later in batch generator.
My data is lying in directory /content/PageSegData/.Page images are in PageImg directory inside the PageSegData folder while the segmentation mask is in PageSeg.
image_list=os.listdir('/content/PageSegData/PageImg/')
image_list=[filename.split(".")[0]for filename in image_list]
Some helping functions
Visualize function
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()
This function helps in visualizing the image segmentation mask and image side by side using matplotlib subplots.
get_segmented_function
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 As I discussed in the dataset section above, that I have used Pixel Annotation Tool for creating the dataset. When we use pixel annotation tool it gives different color or number to different classes when we annotate.In our case the no of classes were two i.e 0 and 24. It could have been anything as long as each class has seperate number. But since there are two classes I decided to remap 0 to 0 and 24 to 1 that is what being done in this function.Also this function converts the image from 3 channel (RGB)to single channel (gray).
The preprocess_img function
def preprocess_img(img):
img=cv2.resize(img,(512,512))
return imgThis function resizes the image to 512x512 which is the shape the UNet model expects.
The batch_generator function
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/PageSegData/PageImg/{fn}.JPG',0)
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
seg=cv2.imread(f'/content/PageSegData/PageSeg/{fn}_mask.png',1)
seg=get_segmented_img(seg,n_classes)
X.append(img)
Y.append(seg)
yield np.array(X),np.array(Y) This is a simple batch generator function which reads the image from the image directory. There is one additional preprocessing we are doing here which is inverse image binarization. You can even try avoiding this step and check the result. After performing image binarization, the image is normalized and appended to an array which stores the image batch. The segmentation mask is pre-processed using the get_segmented function.
Lets’s define model and view its summary.
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 = FCN(n_classes=2,
# input_height=320,
# input_width=320)
model=unet()
model.summary()
Model: "model_1"
__________________________________________________________________________________________________
Layer (type) Output Shape Param # Connected to
==================================================================================================
input_1 (InputLayer) (None, 512, 512, 1) 0
__________________________________________________________________________________________________
conv2d_1 (Conv2D) (None, 512, 512, 64) 640 input_1[0][0]
__________________________________________________________________________________________________
conv2d_2 (Conv2D) (None, 512, 512, 64) 36928 conv2d_1[0][0]
__________________________________________________________________________________________________
max_pooling2d_1 (MaxPooling2D) (None, 256, 256, 64) 0 conv2d_2[0][0]
__________________________________________________________________________________________________
conv2d_3 (Conv2D) (None, 256, 256, 128 73856 max_pooling2d_1[0][0]
__________________________________________________________________________________________________
conv2d_4 (Conv2D) (None, 256, 256, 128 147584 conv2d_3[0][0]
__________________________________________________________________________________________________
max_pooling2d_2 (MaxPooling2D) (None, 128, 128, 128 0 conv2d_4[0][0]
__________________________________________________________________________________________________
conv2d_5 (Conv2D) (None, 128, 128, 256 295168 max_pooling2d_2[0][0]
__________________________________________________________________________________________________
conv2d_6 (Conv2D) (None, 128, 128, 256 590080 conv2d_5[0][0]
__________________________________________________________________________________________________
max_pooling2d_3 (MaxPooling2D) (None, 64, 64, 256) 0 conv2d_6[0][0]
__________________________________________________________________________________________________
conv2d_7 (Conv2D) (None, 64, 64, 512) 1180160 max_pooling2d_3[0][0]
__________________________________________________________________________________________________
conv2d_8 (Conv2D) (None, 64, 64, 512) 2359808 conv2d_7[0][0]
__________________________________________________________________________________________________
dropout_1 (Dropout) (None, 64, 64, 512) 0 conv2d_8[0][0]
__________________________________________________________________________________________________
max_pooling2d_4 (MaxPooling2D) (None, 32, 32, 512) 0 dropout_1[0][0]
____________________________________________________________________
conv2d_9 (Conv2D) (None, 32, 32, 1024) 4719616 max_pooling2d_4[0][0]
__________________________________________________________________________________________________
conv2d_10 (Conv2D) (None, 32, 32, 1024) 9438208 conv2d_9[0][0]
__________________________________________________________________________________________________
dropout_2 (Dropout) (None, 32, 32, 1024) 0 conv2d_10[0][0]
__________________________________________________________________________________________________
up_sampling2d_1 (UpSampling2D) (None, 64, 64, 1024) 0 dropout_2[0][0]
__________________________________________________________________________________________________
conv2d_11 (Conv2D) (None, 64, 64, 512) 2097664 up_sampling2d_1[0][0]
__________________________________________________________________________________________________
concatenate_1 (Concatenate) (None, 64, 64, 1024) 0 dropout_1[0][0]
conv2d_11[0][0]
__________________________________________________________________________________________________
conv2d_12 (Conv2D) (None, 64, 64, 512) 4719104 concatenate_1[0][0]
__________________________________________________________________________________________________
conv2d_13 (Conv2D) (None, 64, 64, 512) 2359808 conv2d_12[0][0]
__________________________________________________________________________________________________
up_sampling2d_2 (UpSampling2D) (None, 128, 128, 512 0 conv2d_13[0][0]
__________________________________________________________________________________________________
conv2d_14 (Conv2D) (None, 128, 128, 256 524544 up_sampling2d_2[0][0]
__________________________________________________________________________________________________
concatenate_2 (Concatenate) (None, 128, 128, 512 0 conv2d_6[0][0]
conv2d_14[0][0]
__________________________________________________________________________________________________
conv2d_15 (Conv2D) (None, 128, 128, 256 1179904 concatenate_2[0][0]
__________________________________________________________________________________________________
conv2d_16 (Conv2D) (None, 128, 128, 256 590080 conv2d_15[0][0]
__________________________________________________________________________________________________
up_sampling2d_3 (UpSampling2D) (None, 256, 256, 256 0 conv2d_16[0][0]
__________________________________________________________________________________________________
conv2d_17 (Conv2D) (None, 256, 256, 128 131200 up_sampling2d_3[0][0]
__________________________________________________________________________________________________
concatenate_3 (Concatenate) (None, 256, 256, 256 0 conv2d_4[0][0]
conv2d_17[0][0]
__________________________________________________________________________________________________
conv2d_18 (Conv2D) (None, 256, 256, 128 295040 concatenate_3[0][0]
__________________________________________________________________________________________________
conv2d_19 (Conv2D) (None, 256, 256, 128 147584 conv2d_18[0][0]
__________________________________________________________________________________________________
up_sampling2d_4 (UpSampling2D) (None, 512, 512, 128 0 conv2d_19[0][0]
__________________________________________________________________________________________________
conv2d_20 (Conv2D) (None, 512, 512, 64) 32832 up_sampling2d_4[0][0]
__________________________________________________________________________________________________
concatenate_4 (Concatenate) (None, 512, 512, 128 0 conv2d_2[0][0]
conv2d_20[0][0]
__________________________________________________________________________________________________
conv2d_21 (Conv2D) (None, 512, 512, 64) 73792 concatenate_4[0][0]
__________________________________________________________________________________________________
conv2d_22 (Conv2D) (None, 512, 512, 64) 36928 conv2d_21[0][0]
__________________________________________________________________________________________________
conv2d_23 (Conv2D) (None, 512, 512, 2) 1154 conv2d_22[0][0]
__________________________________________________________________________________________________
conv2d_24 (Conv2D) (None, 512, 512, 1) 3 conv2d_23[0][0]
==================================================================================================
Total params: 31,031,685
Trainable params: 31,031,685
Non-trainable params: 0
____________________________________________________________________
Let’s split the dataset into training and test set. I decided to use 70% data for training and 30% for testing.
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)):]
Let’s train the model.
We use model check-pointing to train the model.The batch size is 2.
from keras.callbacks import ModelCheckpoint
mc = ModelCheckpoint('weights{epoch:08d}.h5',
save_weights_only=True, period=1)
model.fit_generator(batch_generator(file_train,2,2),epochs=3,steps_per_epoch=1000,validation_data=batch_generator(file_test,2,2),
validation_steps=400,callbacks=[mc],shuffle=1)
Let’s test the trained model on a new page image.
img=cv2.imread(f'/content/test_img.JPG',0)
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.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)
Now let’s detect the contour around these segmentation mask and draw bounding boxes around the original image using these contour coordinates.
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('/content/test_img.JPG')
ori_img=cv2.resize(ori_img,(512,512))
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, (x, y), (x+w,y+h), 255, 1)
coordinates.append([x,y,(x+w),(y+h)])
#cv2.drawContours(img, contours, -1, (255, 255, 0), 1)
cv2.imwrite("output.png",ori_img)
OpenCV function findContours is used to detect contours around these segmentation and the function boundingRect is used to estimate the bounding box from the contours.
The output image:
Further Work
The cropped lines images could be later be passed to a CRNN model which can do text recognition and hence thats how you can build your own OCR.
Congrats you have just learned to solve one of the fundamental problem in recognizing the whole page. If you have doubts regrading anything feel free to leave it down in the comment section and clap the article if you like it.
Link to jupyter notebook:-
Comments
No comments yet. Be the first to comment!