from paddleocr import PaddleOCR
import json
from PIL import Image
import gradio as gr
import numpy as np
import cv2


# 获取随机的颜色
def get_random_color():
    c = tuple(np.random.randint(0 ,256 ,3).tolist())
    return c


# 绘制ocr识别结果
def draw_ocr_bbox( image ,boxes ,colors ):
    print(colors)
    box_num = len(boxes)
    for i in range(box_num):
        box = np.reshape(np.array(boxes[i]) ,[-1 ,1 ,2]).astype(np.int64)
        image = cv2.polylines(np.array(image) ,[box] ,True ,colors[i] ,2)
    return image


# torch.hub.download_url_to_file('https://i.imgur.com/aqMBT0i.jpg', 'example.jpg')

def inference( img: Image.Image ,lang ,confidence ):
    ocr = PaddleOCR(use_angle_cls = True ,lang = lang ,use_gpu = False)
    # img_path = img.name
    img2np = np.array(img)
    result = ocr.ocr(img2np ,cls = True)[0]
    # rgb
    image = img.convert('RGB')
    boxes = [line[0] for line in result]
    txts = [line[1][0] for line in result]
    scores = [line[1][1] for line in result]

    # 识别结果
    final_result = [dict(boxes = box ,txt = txt ,score = score ,_c = get_random_color()) for box ,txt ,score in
                    zip(boxes ,txts ,scores)]
    # 过滤 score < 0.5 的
    final_result = [item for item in final_result if item['score'] > confidence]

    im_show = draw_ocr_bbox(image ,[item['boxes'] for item in final_result] ,[item['_c'] for item in final_result])
    im_show = Image.fromarray(im_show)
    data = [[json.dumps(item['boxes']) ,round(item['score'] ,3) ,item['txt']] for item in final_result]
    return im_show ,data


title = 'PaddleOCR'
description = 'Gradio demo for PaddleOCR.'

examples = [
    # ['example_imgs/example.jpg' ,'en' ,0.5] ,
    # ['example_imgs/ch.jpg' ,'ch' ,0.7] ,
    # ['example_imgs/img_12.jpg' ,'en' ,0.7] ,
]

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

if __name__ == '__main__':
    demo = gr.Interface(
        inference ,
        [gr.Image(type = 'pil' ,label = 'Input') ,
         gr.Dropdown(choices = ['ch' ,'en' ,'fr' ,'german' ,'korean' ,'japan'] ,value = 'ch' ,label = 'language') ,
         gr.Slider(0.1 ,1 ,0.5 ,step = 0.1 ,label = 'confidence_threshold')
         ] ,
        # 输出
        [gr.Image(type = 'pil' ,label = 'Output') ,
         gr.Dataframe(headers = ['bbox' ,'score' ,'text'] ,label = 'Result')] ,
        title = title ,
        description = description ,
        examples = examples ,
        css = css ,
    )
    demo.queue(max_size = 10)
    demo.launch(debug = True ,server_name = "0.0.0.0")