from mmocr.ocr import MMOCR
import gradio as gr
import cv2
import sys
import torch

device = 'gpu' if torch.cuda.is_available() else 'cpu'

ocr = MMOCR(det_config='model/det/config.py',
            det_ckpt='model/det/model.pth',
            recog_config='model/recog/config.py',
            recog_ckpt='model/det/model.pth',
            device='cpu')

def get_rec(points):
    xs = []
    ys = []
    for ix, iv in enumerate(points):
        if ix % 2:
            ys.append(iv)
        else:
            xs.append(iv)
    return (min(xs), min(ys)), (max(xs), max(ys))
        
    
def predict(image_input, score_threshold):
    draw_img = image_input.copy()
    print('image shape', image_input.shape, file=sys.stderr)

    try:
        output = ocr.readtext(image_input)
    except:
        output = dict(det_polygons=[], det_scores=[])

    polygons = output['det_polygons']
    scores = output['det_scores']
    for polygon, score in zip(polygons, scores):
        if score < score_threshold:
            continue
        p0, p1 = get_rec([int(i) for i in polygon])
        draw_img = cv2.rectangle(draw_img, p0, p1, (255,0,0), 2)

    return draw_img, output

def run():
    demo = gr.Interface(
        fn=predict,
        inputs=[gr.components.Image(), gr.Slider(0, 1, 0.8)],
        outputs=[gr.components.Image(), gr.JSON()],
    )

    demo.launch(server_name="0.0.0.0", server_port=7860)


if __name__ == "__main__":
    run()