import os, io
from paddleocr import PaddleOCR, draw_ocr,PPStructure
from PIL import Image, ImageDraw
import gradio as gr


def inference__ppocr(img_path):

    ocr = PaddleOCR(
        rec_char_dict_path='zhtw_common_dict.txt',
        use_gpu=False,
        rec_image_shape="3, 48, 320"
    )
    
    result = ocr.ocr(img_path)
    
    for idx in range(len(result)):
        res = result[idx]
        for line in res:
            print(line)
    
    result = result[0]
    image = Image.open(img_path).convert('RGB')
    boxes = [line[0] for line in result]
    txts = [line[1][0] if line[1] else '' for line in result]  # 確保在無文字時 txts 還是個空字串
    scores = [line[1][1] for line in result]
    im_show_pil = draw_ocr(image, boxes, txts, scores, font_path="./simfang.ttf")

    return im_show_pil, "\n".join(txts)


def inference__ppstructure(img_path):

    ppsutructure = PPStructure(
        rec_char_dict_path='zhtw_common_dict.txt',
        use_gpu=False,
        rec_image_shape="3, 48, 320",
        ser_dict_path='ppocr/utils/dict/kie/clinical_class_list.txt'
    )
    
    result = ppsutructure.__call__(img_path)

    for idx in range(len(result)):
        res = result[idx]
        for line in res:
            print(line)
    
    result = result[0]
    image = Image.open(img_path).convert('RGB')
    boxes = [line[0] for line in result]
    txts = [line[1][0] if line[1] else '' for line in result]  # 確保在無文字時 txts 還是個空字串
    scores = [line[1][1] for line in result]
    im_show_pil = draw_ocr(image, boxes, txts, scores, font_path="./simfang.ttf")

    return im_show_pil, "\n".join(txts)

    
gr.Interface(
    inference__ppstructure,
    [gr.Image(type='filepath', label='圖片上傳')],
    outputs=[
        gr.Image(type="pil", label="識別結果"),
        "text"
    ],
    ).launch(debug=True)