EasyOCR-VietOCR / app.py
hantech's picture
Update app.py
6eaa19c verified
import pandas as pd
import PIL
from PIL import Image
from PIL import ImageDraw
import gradio as gr
import torch
import easyocr
import omegaconf
import cv2
import json
from vietocr.vietocr.tool.predictor import Predictor
from vietocr.vietocr.tool.config import Cfg
# Configure of VietOCR
config = Cfg.load_config_from_name('vgg_transformer')
# config = Cfg.load_config_from_file('vietocr/config.yml')
# config['weights'] = '/Users/bmd1905/Desktop/pretrain_ocr/vi00_vi01_transformer.pth'
config['cnn']['pretrained'] = True
config['predictor']['beamsearch'] = True
config['device'] = 'cpu' # mps
recognitor = Predictor(config)
#model_name = "microsoft/xdoc-base-squad2.0"
#nlp = pipeline('question-answering', model=model_name)
#from transformers import AutoModelForQuestionAnswering, AutoTokenizer, pipeline
#model_name = "timpal0l/mdeberta-v3-base-squad2"
#model = AutoModelForQuestionAnswering.from_pretrained(model_name)
#tokenizer = AutoTokenizer.from_pretrained(model_name)
#nlp = pipeline('question-answering', model=model, tokenizer=tokenizer)
'''
def query(doc, labels):
questions = labels.split(", ")
result={}
for question in questions:
QA_input = {
'question': question+'?',
'context': doc
}
res= nlp(QA_input)
print(res)
value = res['answer']
result[question]=value
return result
'''
def draw_boxes(image, bounds, color='yellow', 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(filepath, lang, labels):
img = cv2.imread(filepath)
width, height, _ = img.shape
if width>height:
height, width, _ = img.shape
reader = easyocr.Reader(lang)
bounds = reader.readtext(filepath)
new_bounds=[]
texts=''
for (bbox, text, prob) in bounds:
(tl, tr, br, bl) = bbox
tl = (int(tl[0]), int(tl[1]))
tr = (int(tr[0]), int(tr[1]))
br = (int(br[0]), int(br[1]))
bl = (int(bl[0]), int(bl[1]))
min_x = min(tl[0], tr[0], br[0], bl[0])
min_x = max(0, min_x)
max_x = max(tl[0], tr[0], br[0], bl[0])
max_x = min(width-1, max_x)
min_y = min(tl[1], tr[1], br[1], bl[1])
min_y = max(0, min_y)
max_y = max(tl[1], tr[1], br[1], bl[1])
max_y = min(height-1, max_y)
# crop the region of interest (ROI)
cropped_image = img[min_y:max_y,min_x:max_x] # crop the image
cropped_image = Image.fromarray(cropped_image)
out = recognitor.predict(cropped_image)
texts = texts + '\t' + out
#result = query(texts, labels)
#jsonText = json.dumps(result)
jsonText='{"result":"ok"'
im = PIL.Image.open(filepath)
draw_boxes(im, bounds)
im.save('result.jpg')
return ['result.jpg', texts, jsonText]
title = 'EasyOCR'
description = 'Gradio demo for EasyOCR. EasyOCR demo supports 80+ languages.To use it, simply upload your image and choose a language from the dropdown menu, or click one of the examples to load them. Read more at the links below.'
article = "<p style='text-align: center'><a href='https://www.jaided.ai/easyocr/'>Ready-to-use OCR with 80+ supported languages and all popular writing scripts including Latin, Chinese, Arabic, Devanagari, Cyrillic and etc.</a> | <a href='https://github.com/JaidedAI/EasyOCR'>Github Repo</a></p>"
css = ".output_image, .input_image {height: 40rem !important; width: 100% !important;}"
choices = [
"vi"
]
gr.Interface(
inference,
[gr.Image(type='filepath', label='Input'),gr.CheckboxGroup(choices, type="value", default=['vi'], label='language'), gr.Textbox(label='Labels',default='Người nα»™p, TiαΊΏp nhαΊ­n hα»“ sΖ‘ của')],
[gr.Image(type='pil', label='Output'), gr.Textbox(label='Text'), gr.JSON(label='document')],
title=title,
description=description,
article=article,
css=css,
examples=[['giaytiepnhan.jpg',['vi'],'Người nα»™p, TiαΊΏp nhαΊ­n hα»“ sΖ‘ của'],['giaytiepnhan2.jpg',['vi'],'MΓ£ sα»‘ hα»“ sΖ‘, Địa chỉ']],
enable_queue=True
).launch(debug=True)