DocBed: A Multi-Stage OCR Solution for Documents with Complex Layouts

Inspired by a related study (Meier et al. 2017) that performs semantic segmentation approach based on the visual appearance of the document, the FCN semantic segmentation model is re-implemented in this work as a baseline. The fully convolutional neural network includes feature extraction, down-sampling convolution, upscaling deconvolution, refinement convolution, and a classification layer that classifies the pixels. The model is trained from scratch with random weights initialization.

In the last layer of the model, two different class settings are considered for pixel classification: seven-class setting that considers each of the semantic content types (i.e. header, article title, article body, image, ad, table) as separate classes and a two-class setting that only regards foreground and background classes where foreground class includes all of the six content types; these settings are referred to as FCN-7 and FCN-2, respectively. The input images and the ground-truth masks are resized to 512x512 during the pre-processing steps and they are randomly combined into batches of 20 samples to be fed into the neural network. A weighted cross entropy loss is applied to better address and minimize the effects of class imbalance on the model. The FCN model is trained for 100 epochs, which takes almost 6 hours on a single GPU.

Mask R-CNN  (He et al. 2017) is a widely-used state-of-art approach for instance segmentation. Since Mask R-CNN model predicts both bounding boxes and semantic masks, this approach saves us a significant extra effort to obtain bounding boxes for OCR engine. Furthermore, all the rest segmentation methods discussed in this paper require morphological processes and labeling components to convert the semantic mask into bounding boxes. Moreover, the semantic mask from Mask R-CNN could be either polygon, which would also work for layout contains more complex shape that rectangles could not represent. The Detectron2 framework  (Wu et al. 2019) provides an extensive collection of baseline models including Mask R-CNN that makes experimenting with this model a fast and friction-less process.

In this work, we used two different backbone structures in the architecture of Mask R-CNN: the first backbone is a Residual Networks (ResNet) combined with a feature pyramid network (FPN) structure that uses standard convolutional and fully connected heads for feature extraction. The second structure, Dilated-C5 (DC5), is a ResNet conv5 backbone with dilation in conv5 and convolutional and fully connected heads for mask and bounding box prediction. Moreover, for the ResNet component, both ResNet-50 and ResNet-101 structures are used. Therefore, three different backbones are explored with Mask R-CNN, which are R50-FPN, R101-FPN, and R50-DC5. The backbones are initialized with Microsoft Research Asia’s original models trained on COCO (Lin et al. 2014) dataset. The input images are resized such that their shorter sides are 800 pixels. The models are trained with a mini-batch of 2 images per GPU. The model with R101-FPN backbone is trained on 4 GPUs for 800 epochs, and the ones with R50-FPN and R50-DC backbones are trained on 8 GPUs for 1000 and 600 epochs, respectively. With this training setup, each experiment takes around 3-4 days. The learning rate is set to 3e-4 for all of the experiments.

The Segmentation transformer (SETR) model (Zheng et al. 2021) has been recently shown to achieve the state of the art results in semantic segmentation on a number of datasets, including ADE20K and Pascal Context.

What makes SETR different from the previous segmentation models with fully convolution encoder-decoder structures is that it reformulates the task of semantic segmentation from a sequence to sequence point of view. Following this intuition, it considers the semantic segmentation task as translating an input image into a probability density map with the same resolution and employs an architecture similar to language translation models to tackle the problem.

To convert an input image into a sequence, we divide the image into a grid of fixed-size patches. Each patch is then flattened and linearly projected into an embedding space. After adding positional encoding to the embedded vectors, they are fed to a transformer encoder that outputs a sequence of feature maps. Finally, a simple convolution decoder followed by some upsampling operations is applied to the feature maps to convert them to the original image resolution.

In our experiments with the SETR model, multiple different augmentation techniques such as random-cropping, random-resizing and flipping were used to make the training more robust to over-fitting. Moreover, a weighted cross-entropy loss is used for optimization to compensate for the class imbalance in input pixel distributions. The model is trained with a batch size of 2 images for 1000 epochs. The input resolution is set to 512x512 pixels to allow training on a single GPU. The initial learning rate is set to 0.001 and it is decayed with a rate of 0.999 after every epoch. Under the above setting, training this model takes about 5 hours.

A series of heuristic post processing steps are applied to prepare the model’s output for the next steps down the pipeline which will be explained in the next section.

Instead of putting the model focus on the entities that need to be digitized, a new deep neural net based model that followed the idea behind classic top-down approach was proposed and experimented. The Separator Detector (SepDet) model was developed to detect separators (e.g. lines or white spaces that separate different articles or multiple columns of a single article) that exist in the documents. These separators are very useful for layout analysis as they contain important information to identify the logical structure of a newspaper page. The purpose of this model is to separate the text regions from the background and split them into a series of logical blocks using detected separators. Instead of using heuristics to identify the locations of these separators, a FCN model is trained to do such task. As the model name indicates, this model does not classify the text regions to different entity types (e.g. table, ad and articles). For this model, the ground truth only contains 2 classes: separators and non-separators. To accelerate training while maintaining enough details for detecting the separators, the original input images are scaled to 512x512 pixels. As these separators often comprise less than 10 percent of the image area, a weighted cross-entropy loss is implemented for training the model. In our experiments with SepDet model, it is trained using a mini-batch size of 20 images for 100 epochs. The learning rate is set to 5e-4. And it takes 1 hour and 20 minutes to train the model on a single GPU. A visualization of the predicted separators by SepDet for a test image and the corresponding processed ground truth map are shown in Figure 4.

In this post-processing procedure, bounding boxes are computed and refined in the following steps. First step is to extract bounding boxes from a predicted 2D array representing class probability distribution over pixels. We started with applying Ostu’s thresholding method (N. Otsu 1979) to convert the probability map as a grayscale mask into a binary mask to tackle this task. Then we applied the opening morphology operation to eliminate small regions in multiple iterations. Next, we used the flooding algorithm to label all of the connected components distinctively. Finally, a bounding box is fitted to each refined connected component. Figure 5 illustrates the explained process.

The second step is removing gaps between bounding boxes. There are often small gaps of several pixels between boundaries of annotated segments in the ground truth when the text columns are very close. Since several models are trained on down-sized ground-truth, and when scaling the bounding boxes to the original scale, those gaps can be inflated, resulting in missing text characters near the boundaries. To reduce those gaps, we expanded the predicted bounding boxes by applying a clustering algorithm (Figure 6) on bounding box vertices to replace the neighboring vertices with each cluster’s centroid.

This post-processing algorithm serves three primary purposes. The first purpose is to convert the generated maps at the output of a segmentation model to bounding boxes. These bounding boxes are used to crop the input images around the text segments so that they can be digested by the OCR engine. The second purpose is to combine the generated bounding boxes by the previous step, whenever it is possible, to minimize the number of OCR API calls for cost optimization. Finally, the last purpose is to logically order the final bounding boxes so that the final extracted text can be read in the correct human-readable order.

In order to accomplish the first purpose, similar techniques explained in the geometric post-processing approach are employed. To accomplish the second purpose, bounding boxes that are vertically aligned with some tolerance threshold are combined regardless of their class (as they are expected to be in the same column). Finally, to accomplish the third purpose, final combined bounding boxes are ordered from top left to bottom right to follow human reading habits.