Spaces:
Sleeping
Sleeping
File size: 3,431 Bytes
52c57e0 9c725ef 52c57e0 f1eda89 52c57e0 4093517 d6f6a75 52c57e0 a3a5528 0453cbe 3f72c9d a3a5528 0326abd a3a5528 9b8edeb 808fea7 96db0e7 9b8edeb f88cfb3 670eda8 a0b35de 670eda8 f88cfb3 a3a5528 7dfc92a a3a5528 3ba1c73 a3a5528 4d31291 0f12594 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 |
import torch
import numpy as np
from transformers import LayoutLMv3TokenizerFast, LayoutLMv3Processor, LayoutLMv3ForTokenClassification
from PIL import Image, ImageDraw, ImageFont
from utils import OCR, unnormalize_box
tokenizer = LayoutLMv3TokenizerFast.from_pretrained("mp-02/layoutlmv3-base-cord-sroie", apply_ocr=False)
processor = LayoutLMv3Processor.from_pretrained("mp-02/layoutlmv3-base-cord-sroie", apply_ocr=False)
model = LayoutLMv3ForTokenClassification.from_pretrained("mp-02/layoutlmv3-base-cord-sroie")
id2label = model.config.id2label
label2id = model.config.label2id
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)
import json
def prediction(image):
boxes, words = OCR(image)
# Preprocessa l'immagine e il testo con il processore di LayoutLMv3
encoding = processor(image, words, boxes=boxes, return_offsets_mapping=True, return_tensors="pt", truncation=True)
offset_mapping = encoding.pop('offset_mapping')
# Esegui l'inferenza con il modello fine-tuned
outputs = model(**encoding)
predictions = outputs.logits.argmax(-1).squeeze().tolist()
token_boxes = encoding.bbox.squeeze().tolist()
probabilities = torch.softmax(outputs.logits, dim=-1)
confidence_scores = probabilities.max(-1).values.squeeze().tolist()
inp_ids = encoding.input_ids.squeeze().tolist()
inp_words = [tokenizer.decode(i) for i in inp_ids]
width, height = image.size
is_subword = np.array(offset_mapping.squeeze().tolist())[:, 0] != 0
true_predictions = [id2label[pred] for idx, pred in enumerate(predictions) if not is_subword[idx]]
true_boxes = [unnormalize_box(box, width, height) for idx, box in enumerate(token_boxes) if not is_subword[idx]]
true_confidence_scores = [confidence_scores[idx] for idx, conf in enumerate(confidence_scores) if not is_subword[idx]]
true_words = []
for id, i in enumerate(inp_words):
if not is_subword[id]:
true_words.append(i)
else:
true_words[-1] = true_words[-1]+i
true_predictions = true_predictions[1:-1]
true_boxes = true_boxes[1:-1]
true_words = true_words[1:-1]
true_confidence_scores = true_confidence_scores[1:-1]
d = {}
for id, i in enumerate(true_predictions):
#rimuovo i prefissi
if i != "O":
i = i[2:]
if i not in d.keys():
d[i] = true_words[id]
else:
d[i] = d[i] + ", " + true_words[id]
d = {k: v.strip() for (k, v) in d.items()}
if "O" in d: d.pop("O")
if("MENU.NM" in d and "MENU.PRICE" in d):
if(len(d["MENU.NM"].split(", ")) == len(d["MENU.PRICE"].split(", "))):
menu_names = [name.strip() for name in d["MENU.NM"].split(', ')]
menu_prices = [price.strip() for price in d["MENU.PRICE"].split(', ')]
menu_combined = [{"ITEM": name, "PRICE": price} for name, price in zip(menu_names, menu_prices)]
d.pop("MENU.NM")
d.pop("MENU.PRICE")
d["MENU"] = menu_combined
draw = ImageDraw.Draw(image, "RGBA")
font = ImageFont.load_default()
for prediction, box, confidence in zip(true_predictions, true_boxes, true_confidence_scores):
draw.rectangle(box)
draw.text((box[0]+10, box[1]-10), text=str(prediction)+ ", "+ str(confidence), font=font, fill="black", font_size="15")
return d
|