app.py

# -*- coding: utf-8 -*-
"""Gradio with DocFormer

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1_XBurG-8jYF4eJJK5VoCJ2Y1v6RV9iAW
"""

## Requirements.txt
import os
os.system('pip install pyyaml==5.1')
## install PyTesseract
os.system('pip install -q pytesseract')
os.environ["TOKENIZERS_PARALLELISM"] = "false"

## Importing the functions from the DocFormer Repo
from dataset import create_features
from modeling import DocFormerEncoder,ResNetFeatureExtractor,DocFormerEmbeddings,LanguageFeatureExtractor
from transformers import BertTokenizerFast
from utils import DocFormer

## Hyperparameters
import torch

seed = 42
target_size = (500, 384)
max_len = 128

## Setting some hyperparameters

device = 'cuda' if torch.cuda.is_available() else 'cpu'

config = {
  "coordinate_size": 96,              ## (768/8), 8 for each of the 8 coordinates of x, y
  "hidden_dropout_prob": 0.1,
  "hidden_size": 768,
  "image_feature_pool_shape": [7, 7, 256],
  "intermediate_ff_size_factor": 4,
  "max_2d_position_embeddings": 1024,
  "max_position_embeddings": 128,
  "max_relative_positions": 8,
  "num_attention_heads": 12,
  "num_hidden_layers": 12,
  "pad_token_id": 0,
  "shape_size": 96,
  "vocab_size": 30522,
  "layer_norm_eps": 1e-12,
}

## Defining the tokenizer
tokenizer = BertTokenizerFast.from_pretrained("bert-base-uncased")



docformer = DocFormer(config)

# path_to_weights = 'drive/MyDrive/docformer_rvl_checkpoint/docformer_v1.ckpt'

url = 'https://www.kaggleusercontent.com/kf/97691030/eyJhbGciOiJkaXIiLCJlbmMiOiJBMTI4Q0JDLUhTMjU2In0..lvzr1OmuVtg5OU9A8XpvPA.uWaGx0E9rlQBfqEvsGQljo0TLy_hdWvA_tUo5ngWoOgT5zHNICZrid2nvC56lsFxAaFcRzCpQdIgWbDEo8bfyyUVQE3d-g-yAUMJZgy8fzAdzQvf4FHk2c-LqHcIkXj8F_vtKc3QqsvOxBSu8hHB7pd4jGhzPeVkYkLLbrBU2zlT9YboO3RkuPp_6eQEqSRmUPtm4wOa0TA5J3sIXbBMaEeZV8xYGAzph9i9ke900XAQZEYQFScGblT5-FKwKDlZQmasQRezsxYBFima_Q9sbSLtFA0JH-4lUEmpw0HExoZiPiw79IbDVlF4y58prZpHivlc3YMuF94d4GWtH9WQ9gpddDnCmuCuVJswUuKSlAYc_bifGj7wS5S3xXddWjtVMTWuJN-2e6hLcGil7oooWlYv281qMNm6jIOHoyxpkLdlCViTlkqOQcgoO2yeM_QCpewdiWGSFanJgvA8kGGIUj79JqzgSmSqhUvGdTtsHIPIjDOTzKrRBCaC6isJVLifjhTGPKV-E9SHIz7ndyuVzIhU2C0KowXRQSduOT7quvzZjYlTozeeYZGx8me2qALxr81F1jwtljUkZw7m3GTPi-9fOIFixcZw4uxSPYQTlFiSw6lzh0XecuupglGTIvCI18tJuQDf6f_-bJH97g4OwablQUVxQQLnLkxlRfFSQkci_uqooQ5RsqK4dmsOq4i9.WwuNA6C-_f8grPIzmL7SYw/models/epoch=0-step=753.ckpt'

docformer.load_from_checkpoint(url)

id2label = ['scientific_report',
 'resume',
 'memo',
 'file_folder',
 'specification',
 'news_article',
 'letter',
 'form',
 'budget',
 'handwritten',
 'email',
 'invoice',
 'presentation',
 'scientific_publication',
 'questionnaire',
 'advertisement']

import gradio as gr

## Taken from LayoutLMV2 space

image = gr.inputs.Image(type="pil")
label = gr.outputs.Label(num_top_classes=5)
examples = [['00093726.png'], ['00866042.png']]
title = "Interactive demo: DocFormer for Image Classification"
description = "Demo for classifying document images with DocFormer model. To use it, \
simply upload an image or use the example images below and click 'submit' to let the model predict the 5 most probable Document classes. \
Results will show up in a few seconds."

def classify_image(image):

  image.save('sample_img.png')
  final_encoding = create_features(
            './sample_img.png',
            tokenizer,
            add_batch_dim=True,
            target_size=target_size,
            max_seq_length=max_len,
            path_to_save=None,
            save_to_disk=False,
            apply_mask_for_mlm=False,
            extras_for_debugging=False,
            use_ocr = True
    )

  keys_to_reshape = ['x_features', 'y_features', 'resized_and_aligned_bounding_boxes']
  for key in keys_to_reshape:
      final_encoding[key] = final_encoding[key][:, :max_len]

  from torchvision import transforms
  # ## Normalization to these mean and std (I have seen some tutorials used this, and also in image reconstruction, so used it)
  transform = transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))

  final_encoding['resized_scaled_img'] = transform(final_encoding['resized_scaled_img'])
  output = docformer.forward(final_encoding)
  output = output[0].softmax(axis = -1)
  
  final_pred = {}
  for i, score in enumerate(output):
      score = output[i]
      final_pred[id2label[i]] = score.detach().cpu().tolist()
      
  return final_pred

gr.Interface(fn=classify_image, inputs=image, outputs=label, title=title, description=description, examples=examples, enable_queue=True).launch(debug=True)