vtomoasv's picture
Update app.py
d90fed6
import numpy as np
from io import BytesIO
import matplotlib.font_manager as mfm
from matplotlib import mathtext
import pandas as pd
import gradio as gr
from cnocr import CnOcr
from PIL import Image
from pix2tex.cli import LatexOCR
from pix2tex.utils.utils import post_process
from paddleocr import PPStructure#,draw_structure_result,save_structure_res
def table_text(img):
table_engine = PPStructure(layout=False, show_log=True)
result = table_engine(img)
im_show = pd.read_html(result[0]['res']['html'])[0]
#print(im_show)
return im_show
def pix_text(x):
x = Image.fromarray(x)
model = LatexOCR()
out = model(x)
out = '$'+out+'$'
#out = re.compile(r'\left').sub('', out)
#out = re.compile(r'\right').sub('', out)
return out
def merge_strings(lst):
indices = []
for num,char in enumerate(lst):
if char[-1] == '。' or char[-1] == '!' or char[-1] == '.':
indices.append(num+1)
groups = [' '.join(lst[i:j]) for i, j in zip([0]+indices, indices+[None])]
return groups
def ocr_image(x,det_model,rec_model):
text = []
ocr = CnOcr(det_model_name=det_model,rec_model_name=rec_model)
out = ocr.ocr(x)
out = list(filter(None, [i['text'] for i in out]))
out = merge_strings(out)
text = ' '.join([j for j in out])
return text
with gr.Blocks() as demo:
gr.Markdown("Product identification")
with gr.Tab("Text product extraction"):
with gr.Row():
det_model_text = gr.Dropdown(["db_shufflenet_v2","db_shufflenet_v2_tiny","db_mobilenet_v3","db_resnet34","db_resnet18","ch_PP-OCRv3_det",
"ch_PP-OCRv2_det","en_PP-OCRv3_det"],value='ch_PP-OCRv3_det' ,interactive=True,
label="detection models")
rec_model_text = gr.Dropdown(["densenet_lite_114-fc","densenet_lite_134-fc","densenet_lite_136-fc","densenet_lite_134-gru",
"densenet_lite_136-gru",
"ch_PP-OCRv3","ch_ppocr_mobile_v2.0","en_PP-OCRv3","en_number_mobile_v2.0","chinese_cht_PP-OCRv3"],
value = "densenet_lite_136-gru",label="recognition model")
with gr.Row():
with gr.Column(min_width=300):
image_input = gr.Image(label='product to be recognized')
with gr.Column(scale=1, min_width=300):
image_output = gr.Textbox(lines=8,label='Identification results').style(show_copy_button=True)
with gr.Row():
clear_button = gr.ClearButton(value='clean')
image_button = gr.Button("recognize")
with gr.Tab("Prodcut table extracto"):
with gr.Row():
table_input = gr.Image(label='Prodcut')
#table_output = gr.File()
table_output = gr.Dataframe(datatype=["str", "number"],row_count=5,
col_count=5,label='Result'
)
with gr.Row():
clear_button3 = gr.ClearButton(value='Clear')
table_button = gr.Button("Extract")
image_button.click(ocr_image, inputs=[image_input,det_model_text,rec_model_text], outputs=image_output)
clear_button.add([image_input,image_output])
table_button.click(table_text, inputs=table_input, outputs=table_output)
clear_button3.add([table_input,table_output])
demo.launch()