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


# 設定 Hugging Face Hub 的 Access Token
os.environ["HF_TOKEN"] = "TWOCR"
 
def inference(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)
    
title = "<p style='text-align: center'><a href='https://www.twman.org/AI/CV' target='_blank'>繁體中文醫療診斷書和收據OCR：PaddleOCR</a></p>"

description = """
<p style='text-align: center'><a href="https://blog.twman.org/2023/07/wsl.html" target='_blank'>用PaddleOCR的PPOCRLabel來微調醫療診斷書和收據</a></p><br>
<p style='text-align: center'><a href="https://github.com/Deep-Learning-101" target='_blank'>https://github.com/Deep-Learning-101</a></p><br>
<p style='text-align: center'><a href="https://github.com/Deep-Learning-101/Computer-Vision-Paper" target='_blank'>https://github.com/Deep-Learning-101/Computer-Vision-Paper</a></p><br>
"""


css = ".output_image, .input_image {height: 40rem !important; width: 100% !important;}"

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