ketanmore/surya-ocr-arabic-layout

Script For Layout Model Of Surya OCR.

Please refer to surya-layout-fien-tuneCrossEntropyLoss.ipynb (Differnt Repository) file to understand fine-tuning process.

Setup Instructions

Clone the Surya OCR GitHub Repository

git clone https://github.com/vikp/surya.git
cd surya

Switch to v0.4.14

git checkout f7c6c04

Install Dependencies

You can install the required dependencies using the following command:

pip install -r requirements.txt

Image Pre-processing

For image pre-processing we can directly import a function and image processor from surya ocr github repository.

from surya.input.processing import prepare_image_detection

from surya.model.detection.segformer import load_processor

from PIL import Image
image = Image.open("path/to/image")
images = [prepare_image_detection(img=image, processor=load_processor())]

images = torch.stack(images, dim=0).to(model.dtype).to(torch.float32).to(device)

When an image is passed through prepare_image_detection function, the image is converted into the shape : [1 ,12 , 300 , 300 ] (This is the format(shape) that is expected by Surya Layout Model)

1 (Batch Size): The first dimension is the batch size, which is 1 in this case. It means that this tensor contains data for one example (or image) at a time.

12 (Number of Channels or Classes): The second dimension represents the number of channels or classes.

300 (Height): The third dimension corresponds to the height of the image or spatial dimensions of the data. In this case, the height is 300 pixels.

300 (Width): The fourth dimension corresponds to the width of the image or spatial dimensions. The width is also 300 pixels.

Loading Model

from surya.model.detection.segformer import load_model
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = load_model("ketanmore/surya-ocr-arabic-layout") 

output = model(images).to(device).to(torch.float32)

Post-Processing

Now, since we have output let's understand it. The surya layout model is based of Segformer architecture, this means that surya model produces a segmentation mask. Following are the labels that are known to surya :

print(model.config.id2label) # This snippet can be used to print classes(labels) used for segmentaion

{0: 'Blank',
 1: 'Caption',
 2: 'Footnote',
 3: 'Formula',
 4: 'List-item',
 5: 'Page-footer',
 6: 'Page-header',
 7: 'Picture',
 8: 'Section-header',
 9: 'Table',
 10: 'Text',
 11: 'Title'}

The surya layout model classifies each pixesl into any of the above value resulting in formation of a mask.

Ploting the mask

import torch
import matplotlib.colors as mcolors
import matplotlib.pyplot as plt

probabilities = torch.softmax(outputs.logits, dim=1)
mask = torch.argmax(probabilities, dim=1).squeeze(0) 
mask = mask.cpu().numpy()

class_labels = {
    0: 'Blank',
    1: 'Caption',
    2: 'Footnote',
    3: 'Formula',
    4: 'List-item',
    5: 'Page-footer',
    6: 'Page-header',
    7: 'Picture',
    8: 'Section-header',
    9: 'Table',
    10: 'Text',
    11: 'Title'
}

# Create a color map, one unique color for each class
colors = plt.cm.get_cmap('tab20', len(class_labels))

def plot_segmentation_mask(mask, class_labels, colors):
    fig, ax = plt.subplots(figsize=(10, 10))
    # Display the image
    ax.imshow(mask, cmap=colors, interpolation='nearest')

    # Add class labels
    unique_classes = np.unique(mask)
    for cls in unique_classes:
        # Find the first occurrence of this class in the mask
        indices = np.argwhere(mask == cls)
        if indices.size > 0:
            representative_index = indices[len(indices)//2]  # Pick middle index as representative
            ax.text(representative_index[1], representative_index[0], class_labels[cls],
                    verticalalignment='center', horizontalalignment='center',
                    color='white', fontsize=12, weight='bold')

    plt.axis('off')
    plt.show()

# Call the function with the mask, labels, and colors
plot_segmentation_mask(mask, class_labels, colors)

Ploting Bounding Boxes on the image

To perform this task, we can use some function from surya OCR github repository :

Converting Logits into bounding boxes.

from surya.layout import parallel_get_regions

def logits_to_bboxes(logits,image) :
    correct_shape = (300, 300)  
    logits_temp = F.interpolate(logits, size=correct_shape, mode='bilinear', align_corners=False)
    logits_temp = logits_temp.cpu().detach().numpy().astype(np.float32)

    heatmap_count = logits_temp.shape[1]
    heatmaps = [logits_temp[i][k] for i in range(logits_temp.shape[0]) for k in range(heatmap_count)]
    regions = parallel_get_regions(heatmaps=heatmaps, orig_size=image.size, id2label=model.config.id2label)

    final_bboxes = []
    for i in regions.bboxes :
        final_bboxes.append(i.bbox)
    return final_bboxes


bb = logits_to_bboxes(outputs.logits,image)
print(bb)

# Sample Output
[[41.0, 293.0, 345.0, 517.0],
 [39.0, 141.0, 343.0, 320.0],
 [332.0, 136.0, 528.0, 405.0],
 [332.0, 367.0, 528.0, 714.0],
 [126.0, 62.0, 437.0, 95.0],
 [218.0, 106.0, 343.0, 138.0]]

Now we have bounding boxes that can be used for ploting on images. For this we can directly use the function present in the surya OCR github repository :

from surya.postprocessing.heatmap import draw_bboxes_on_image
from PIL import Image

img = draw_bboxes_on_image(bboxes,Image.open(image_path))

Add-ons

You can directly refer to Compare.ipynb notebook present in the repo to directly start comparing orignal and fine-tuned model.