Spaces:
Running
Running
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() |