lamiaaEl commited on
Commit
968b8b2
1 Parent(s): 5537571

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +109 -0
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
+