Upload 4 files

app.py

@@ -0,0 +1,40 @@
+import pytesseract
+import torch
+import gradio as gr
+from transformers import LayoutLMForSequenceClassification
+from preprocess import apply_ocr,encode_example
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+pytesseract.pytesseract.tesseract_cmd = r"C:\\Program Files\\Tesseract-OCR\\tesseract.exe"
+model = LayoutLMForSequenceClassification.from_pretrained("models/document_model")
+model.to(device)
+classes=['questionnaire', 'memo', 'budget', 'file_folder', 'specification', 'invoice', 'resume',
+                               'advertisement', 'news_article', 'email', 'scientific_publication', 'presentation',
+                               'letter', 'form', 'handwritten', 'scientific_report']
+
+
+def predict(image):
+    example = apply_ocr(image)
+    encoded_example = encode_example(example)
+    input_ids = torch.tensor(encoded_example['input_ids']).unsqueeze(0)
+    bbox = torch.tensor(encoded_example['bbox']).unsqueeze(0)
+    attention_mask = torch.tensor(encoded_example['attention_mask']).unsqueeze(0)
+    token_type_ids = torch.tensor(encoded_example['token_type_ids']).unsqueeze(0)
+    model.eval()
+    outputs=model(input_ids=input_ids, bbox=bbox, attention_mask=attention_mask, token_type_ids=token_type_ids)
+    classification_results = torch.softmax(outputs.logits, dim=1).tolist()[0]
+    max_prob_index = classification_results.index(max(classification_results))
+    predicted_class = classes[max_prob_index]
+    return predicted_class
+
+
+
+title="Document Image Classification"
+
+demo = gr.Interface(
+    fn=predict,
+    inputs=gr.inputs.Image(type="pil"),
+    outputs=gr.outputs.Textbox(label="Predicted Class"),
+    title=title,
+)
+demo.launch(share=True)

models/document_model/config.json

@@ -0,0 +1,63 @@
+{
+  "_name_or_path": "microsoft/layoutlm-base-uncased",
+  "architectures": [
+    "LayoutLMForSequenceClassification"
+  ],
+  "attention_probs_dropout_prob": 0.1,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 768,
+  "id2label": {
+    "0": "LABEL_0",
+    "1": "LABEL_1",
+    "2": "LABEL_2",
+    "3": "LABEL_3",
+    "4": "LABEL_4",
+    "5": "LABEL_5",
+    "6": "LABEL_6",
+    "7": "LABEL_7",
+    "8": "LABEL_8",
+    "9": "LABEL_9",
+    "10": "LABEL_10",
+    "11": "LABEL_11",
+    "12": "LABEL_12",
+    "13": "LABEL_13",
+    "14": "LABEL_14",
+    "15": "LABEL_15"
+  },
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "label2id": {
+    "LABEL_0": 0,
+    "LABEL_1": 1,
+    "LABEL_10": 10,
+    "LABEL_11": 11,
+    "LABEL_12": 12,
+    "LABEL_13": 13,
+    "LABEL_14": 14,
+    "LABEL_15": 15,
+    "LABEL_2": 2,
+    "LABEL_3": 3,
+    "LABEL_4": 4,
+    "LABEL_5": 5,
+    "LABEL_6": 6,
+    "LABEL_7": 7,
+    "LABEL_8": 8,
+    "LABEL_9": 9
+  },
+  "layer_norm_eps": 1e-12,
+  "max_2d_position_embeddings": 1024,
+  "max_position_embeddings": 512,
+  "model_type": "layoutlm",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "output_past": true,
+  "pad_token_id": 0,
+  "position_embedding_type": "absolute",
+  "problem_type": "single_label_classification",
+  "torch_dtype": "float32",
+  "transformers_version": "4.30.2",
+  "type_vocab_size": 2,
+  "use_cache": true,
+  "vocab_size": 30522
+}

preprocess.py

@@ -0,0 +1,83 @@
+import pytesseract
+from PIL import Image
+import numpy as np
+from transformers import LayoutLMTokenizer
+
+
+pytesseract.pytesseract.tesseract_cmd = r"C:\\Program Files\\Tesseract-OCR\\tesseract.exe"
+tokenizer = LayoutLMTokenizer.from_pretrained("microsoft/layoutlm-base-uncased")
+def normalize_box(box, width, height):
+     return [
+         int(1000 * (box[0] / width)),
+         int(1000 * (box[1] / height)),
+         int(1000 * (box[2] / width)),
+         int(1000 * (box[3] / height)),
+     ]
+
+def apply_ocr(image):
+        # get the image
+        # image = Image.open(example['image_path'])
+
+        width, height = image.size
+        example={}
+        # apply ocr to the image
+        ocr_df = pytesseract.image_to_data(image, output_type='data.frame')
+        float_cols = ocr_df.select_dtypes('float').columns
+        ocr_df = ocr_df.dropna().reset_index(drop=True)
+        ocr_df[float_cols] = ocr_df[float_cols].round(0).astype(int)
+        ocr_df = ocr_df.replace(r'^\s*$', np.nan, regex=True)
+        ocr_df = ocr_df.dropna().reset_index(drop=True)
+
+        # get the words and actual (unnormalized) bounding boxes
+        #words = [word for word in ocr_df.text if str(word) != 'nan'])
+        words = list(ocr_df.text)
+        words = [str(w) for w in words]
+        coordinates = ocr_df[['left', 'top', 'width', 'height']]
+        actual_boxes = []
+        for idx, row in coordinates.iterrows():
+            x, y, w, h = tuple(row) # the row comes in (left, top, width, height) format
+            actual_box = [x, y, x+w, y+h] # we turn it into (left, top, left+width, top+height) to get the actual box
+            actual_boxes.append(actual_box)
+
+        # normalize the bounding boxes
+        boxes = []
+        for box in actual_boxes:
+            boxes.append(normalize_box(box, width, height))
+
+        # add as extra columns
+        assert len(words) == len(boxes)
+        example['words'] = words
+        example['bbox'] = boxes
+        return example
+def encode_example(example, max_seq_length=512, pad_token_box=[0, 0, 0, 0]):
+  words = example['words']
+  normalized_word_boxes = example['bbox']
+
+  assert len(words) == len(normalized_word_boxes)
+
+  token_boxes = []
+  for word, box in zip(words, normalized_word_boxes):
+      word_tokens = tokenizer.tokenize(word)
+      token_boxes.extend([box] * len(word_tokens))
+
+  # Truncation of token_boxes
+  special_tokens_count = 2
+  if len(token_boxes) > max_seq_length - special_tokens_count:
+      token_boxes = token_boxes[: (max_seq_length - special_tokens_count)]
+
+  # add bounding boxes of cls + sep tokens
+  token_boxes = [[0, 0, 0, 0]] + token_boxes + [[1000, 1000, 1000, 1000]]
+
+  encoding = tokenizer(' '.join(words), padding='max_length', truncation=True)
+  # Padding of token_boxes up the bounding boxes to the sequence length.
+  input_ids = tokenizer(' '.join(words), truncation=True)["input_ids"]
+  padding_length = max_seq_length - len(input_ids)
+  token_boxes += [pad_token_box] * padding_length
+  encoding['bbox'] = token_boxes
+
+  assert len(encoding['input_ids']) == max_seq_length
+  assert len(encoding['attention_mask']) == max_seq_length
+  assert len(encoding['token_type_ids']) == max_seq_length
+  assert len(encoding['bbox']) == max_seq_length
+
+  return encoding

requirements.txt

@@ -0,0 +1,82 @@
+aiofiles==23.1.0
+aiohttp==3.8.4
+aiosignal==1.3.1
+altair==5.0.1
+annotated-types==0.5.0
+anyio==3.7.1
+async-timeout==4.0.2
+attrs==23.1.0
+certifi==2023.5.7
+charset-normalizer==3.2.0
+click==8.1.6
+colorama==0.4.6
+contourpy==1.1.0
+cycler==0.11.0
+datasets==2.13.1
+dill==0.3.6
+exceptiongroup==1.1.2
+fastapi==0.100.0
+ffmpy==0.3.1
+filelock==3.12.2
+fonttools==4.41.0
+frozenlist==1.4.0
+fsspec==2023.6.0
+gradio==3.37.0
+gradio_client==0.2.10
+h11==0.14.0
+httpcore==0.17.3
+httpx==0.24.1
+huggingface-hub==0.16.4
+idna==3.4
+Jinja2==3.1.2
+jsonschema==4.18.4
+jsonschema-specifications==2023.7.1
+kiwisolver==1.4.4
+linkify-it-py==2.0.2
+markdown-it-py==2.2.0
+MarkupSafe==2.1.3
+matplotlib==3.7.2
+mdit-py-plugins==0.3.3
+mdurl==0.1.2
+mpmath==1.3.0
+multidict==6.0.4
+multiprocess==0.70.14
+networkx==3.1
+numpy==1.25.1
+orjson==3.9.2
+packaging==23.1
+pandas==2.0.3
+Pillow==10.0.0
+pyarrow==12.0.1
+pydantic==2.0.3
+pydantic_core==2.3.0
+pydub==0.25.1
+pyparsing==3.0.9
+pytesseract==0.3.10
+python-dateutil==2.8.2
+python-multipart==0.0.6
+pytz==2023.3
+PyYAML==6.0.1
+referencing==0.30.0
+regex==2023.6.3
+requests==2.31.0
+rpds-py==0.9.2
+safetensors==0.3.1
+semantic-version==2.10.0
+six==1.16.0
+sniffio==1.3.0
+starlette==0.27.0
+sympy==1.12
+tokenizers==0.13.3
+toolz==0.12.0
+torch==2.0.1
+tqdm==4.65.0
+transformers==4.31.0
+typing_extensions==4.7.1
+tzdata==2023.3
+uc-micro-py==1.0.2
+urllib3==2.0.3
+uvicorn==0.23.1
+websockets==11.0.3
+xxhash==3.2.0
+yarl==1.9.2