Spaces:
Sleeping
Sleeping
Create app.py
Browse files
app.py
ADDED
@@ -0,0 +1,109 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import streamlit as st
|
2 |
+
from transformers import LayoutLMv3ForTokenClassification, LayoutLMv3Processor
|
3 |
+
from PIL import Image
|
4 |
+
import torch
|
5 |
+
import easyocr
|
6 |
+
from PIL import Image
|
7 |
+
import re
|
8 |
+
|
9 |
+
# Load the model and processor from Hugging Face
|
10 |
+
model_name = "capitaletech/LayoutLMv3-v1" # Replace with your model repository name
|
11 |
+
model = LayoutLMv3ForTokenClassification.from_pretrained(model_name)
|
12 |
+
processor = LayoutLMv3Processor.from_pretrained(model_name)
|
13 |
+
|
14 |
+
st.title("LayoutLMv3 Text Extraction")
|
15 |
+
st.write("Upload an image to get text predictions using the fine-tuned LayoutLMv3 model.")
|
16 |
+
|
17 |
+
uploaded_file = st.file_uploader("Choose an image...", type="png")
|
18 |
+
|
19 |
+
if uploaded_file is not None:
|
20 |
+
image = Image.open(uploaded_file)
|
21 |
+
st.image(image, caption='Uploaded Image.', use_column_width=True)
|
22 |
+
st.write("")
|
23 |
+
st.write("Classifying...")
|
24 |
+
|
25 |
+
# Process the image
|
26 |
+
words = uploaded_file["tokens"]
|
27 |
+
boxes = uploaded_file["bboxes"]
|
28 |
+
word_labels = uploaded_file["ner_tags"]
|
29 |
+
|
30 |
+
encoding = processor(image, words, boxes=boxes, word_labels=word_labels, return_tensors="pt")
|
31 |
+
|
32 |
+
with torch.no_grad():
|
33 |
+
outputs = model(**encoding)
|
34 |
+
|
35 |
+
logits = outputs.logits
|
36 |
+
predictions = logits.argmax(-1).squeeze().cpu.tolist()
|
37 |
+
labels = encoding['labels'].squeeze().tolist()
|
38 |
+
|
39 |
+
|
40 |
+
# Set up the EasyOCR reader for multiple languages
|
41 |
+
languages = ["ru", "rs_cyrillic", "be", "bg", "uk", "mn", "en"]
|
42 |
+
reader = easyocr.Reader(languages)
|
43 |
+
|
44 |
+
# Load the image
|
45 |
+
image_path = example["img_path"]
|
46 |
+
image = Image.open(image_path)
|
47 |
+
|
48 |
+
# Perform text detection
|
49 |
+
ocr_results = reader.readtext(image_path, detail=1)
|
50 |
+
|
51 |
+
# Extract text and bounding boxes, filtering non-alphabetic characters from text
|
52 |
+
words = []
|
53 |
+
boxes = []
|
54 |
+
|
55 |
+
# Define a regular expression pattern for non-alphabetic characters
|
56 |
+
non_alphabet_pattern = re.compile(r'[^a-zA-Z]+')
|
57 |
+
|
58 |
+
for result in ocr_results:
|
59 |
+
bbox, text, _ = result
|
60 |
+
filtered_text = re.sub(non_alphabet_pattern, '', text)
|
61 |
+
if filtered_text: # Only append if there are alphabetic characters left
|
62 |
+
words.append(filtered_text)
|
63 |
+
boxes.append([
|
64 |
+
bbox[0][0], bbox[0][1],
|
65 |
+
bbox[2][0], bbox[2][1]
|
66 |
+
])
|
67 |
+
|
68 |
+
words = words[1:]
|
69 |
+
|
70 |
+
|
71 |
+
def unnormalize_box(bbox, width, height):
|
72 |
+
return [
|
73 |
+
width * (bbox[0] / 1000),
|
74 |
+
height * (bbox[1] / 1000),
|
75 |
+
width * (bbox[2] / 1000),
|
76 |
+
height * (bbox[3] / 1000),
|
77 |
+
]
|
78 |
+
|
79 |
+
token_boxes = encoding["bbox"].squeeze().tolist()
|
80 |
+
width, height = image.size
|
81 |
+
|
82 |
+
true_predictions = [model.config.id2label[pred] for pred, label in zip(predictions, labels) if label != - 100]
|
83 |
+
true_labels = [model.config.id2label[label] for prediction, label in zip(predictions, labels) if label != -100]
|
84 |
+
true_boxes = [unnormalize_box(box, width, height) for box, label in zip(token_boxes, labels) if label != -100]
|
85 |
+
true_tokens = words
|
86 |
+
|
87 |
+
# Associate languages with their levels
|
88 |
+
languages_with_levels = {}
|
89 |
+
current_language = None
|
90 |
+
|
91 |
+
j=0
|
92 |
+
|
93 |
+
for i in range(0, len(true_labels)):
|
94 |
+
if true_labels[i] == 'language':
|
95 |
+
current_language = words[j]
|
96 |
+
j= j+1
|
97 |
+
languages_with_levels[current_language] = true_labels[i+1]
|
98 |
+
|
99 |
+
print(languages_with_levels)
|
100 |
+
|
101 |
+
input_ids = encoding["input_ids"]
|
102 |
+
bbox = encoding["bbox"]
|
103 |
+
attention_mask = encoding["attention_mask"]
|
104 |
+
|
105 |
+
st.write("Predicted labels:")
|
106 |
+
# Print languages with their levels
|
107 |
+
for language, level in languages_with_levels.items():
|
108 |
+
st.write(f"{language}: {level}")
|
109 |
+
|