app.py 1.9 KB
Newer Older
Z
zhoujun 已提交
1 2 3 4 5
import gradio as gr
import base64
from io import BytesIO
from PIL import Image

6
import paddleocr
Z
zhoujun 已提交
7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25
from paddleocr import PaddleOCR, draw_ocr

ocr = PaddleOCR(ocr_version='PP-OCRv2', use_angle_cls=True, lang="ch")


def image_to_base64(image):
    # 输入为PIL读取的图片,输出为base64格式
    byte_data = BytesIO()  # 创建一个字节流管道
    image.save(byte_data, format="JPEG")  # 将图片数据存入字节流管道
    byte_data = byte_data.getvalue()  # 从字节流管道中获取二进制
    base64_str = base64.b64encode(byte_data).decode("ascii")  # 二进制转base64
    return base64_str


# UGC: Define the inference fn() for your models
def model_inference(image):
    result = ocr.ocr(image, cls=True)

    # 显示结果
26 27
    if paddleocr.__version__ >= "2.6.0.2":
        result = result[0]
Z
zhoujun 已提交
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
    boxes = [line[0] for line in result]
    txts = [line[1][0] for line in result]
    scores = [line[1][1] for line in result]
    im_show = draw_ocr(image, boxes, txts=None, scores=None)
    im_show = Image.fromarray(im_show)

    res = []
    for i in range(len(boxes)):
        res.append(dict(boxes=boxes[i], txt=txts[i], score=scores[i]))
    json_out = {"base64": image_to_base64(im_show), "result": res}
    return im_show, json_out


def clear_all():
    return None, None, None


with gr.Blocks() as demo:
    gr.Markdown("PP-OCRv2")

    with gr.Column(scale=1, min_width=100):
        img_in = gr.Image(
            value="https://raw.githubusercontent.com/PaddlePaddle/PaddleOCR/dygraph/doc/imgs/11.jpg",
            label="Input")

        with gr.Row():
            btn1 = gr.Button("Clear")
            btn2 = gr.Button("Submit")
        img_out = gr.Image(label="Output").style(height=400)
        json_out = gr.JSON(label="jsonOutput")

    btn2.click(fn=model_inference, inputs=img_in, outputs=[img_out, json_out])
    btn1.click(fn=clear_all, inputs=None, outputs=[img_in, img_out, json_out])
    gr.Button.style(1)

demo.launch()