WrittenRec / app.py
JustKiddo's picture
Update app.py
243d9ad verified
raw
history blame
1.1 kB
import pandas as pd
import PIL
from PIL import Image
from PIL import ImageDraw
import gradio as gr
import torch
import easyocr
processor = AutoProcessor.from_pretrained("Sharka/CIVQA_LayoutLMv2_EasyOCR")
model = AutoModelForDocumentQuestionAnswering.from_pretrained("Sharka/CIVQA_LayoutLMv2_EasyOCR")
def draw_boxes(image, bounds, color='red', width=2):
draw = ImageDraw.Draw(image)
for bound in bounds:
p0, p1, p2, p3 = bound[0]
draw.line([*p0, *p1, *p2, *p3, *p0], fill=color, width=width)
return image
def inference(img, lang):
reader = easyocr.Reader(lang)
bounds = reader.readtext(img.name)
im = PIL.Image.open(img.name)
draw_boxes(im, bounds)
im.save('result.jpg')
return ['result.jpg', pd.DataFrame(bounds).iloc[: , 1:]]
gr.Interface(
inference,
[gr.inputs.Image(type='file', label='Input'),gr.inputs.CheckboxGroup(choices, type="value", default=['en', 'vi'], label='language')],
[gr.outputs.Image(type='file', label='Output'), gr.outputs.Dataframe(headers=['text', 'confidence'])],
enable_queue=True
).launch(debug=True)