import cv2
import gradio as gr
import numpy as np

from z_app_factory import get_app

thickness = 3
lineType = 8
font = cv2.FONT_HERSHEY_SIMPLEX

def inference(image):
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    code,lst2d_res = get_app(image)
    msg = "aaa"
    if code == 401:
        msg = "Not RGB three channel picture"
    elif code == 402:
        msg = "Pixels less than 32 * 32"
    elif code == 403:
        msg = "Pixels greater than 4096 * 4096"
    elif code == 404:
        msg = "Files greater than 5MB"
    elif code == 405:
        msg = "System error, please contact\n the server for troubleshooting"

    if code!=200:
        img_out = np.zeros((500, 600,3),dtype=np.uint8)
        cv2.putText(img_out, msg, (20, 200), font, 1, (0, 255, 0), 2)
        return img_out

    # img_out = np.zeros((500, 600, 3),dtype=np.uint8)
    # cv2.putText(img_out, "ease contact the server for trouble", (20, 100), font, 0.8, (0, 255, 0), 2)
    # return img_out
    # img_out = np.zeros((500, 600, 3), dtype=np.uint8)
    # cv2.putText(img_out, msg, (20, 20), font, 1, (0, 255, 0), 2)
    # return img_out
    for face in lst2d_res:

        bbox = [int(i) for i in face["bbox"]]
        score = face['score']
        point_color = (0, int(255 * score), 0) # BGR
        x1, y1 = bbox[:2]
        x2, y2 = bbox[2:]
        cv2.putText(image, str(score)[:4], (x1, y1 - 10), font, 0.8, (0, 255, 0), 2)
        cv2.line(image, (x1, y1), (x2, y1), point_color, thickness, lineType)
        cv2.line(image, (x2, y1), (x2, y2), point_color, thickness, lineType)
        cv2.line(image, (x1, y1), (x1, y2), point_color, thickness, lineType)
        cv2.line(image, (x1, y2), (x2, y2), point_color, thickness, lineType)

        for kp in face["kps"]:
            x, y = [int(i) for i in kp]
            cv2.circle(image, (x, y), 2, (2, 30, 200), 2)

        landmarks = face["landmarks"]
        lst = [[int(i) for i in kp] for kp in landmarks]
        for i, kp in enumerate(lst):
            x, y = kp
            cv2.circle(image, (x, y), 2, (200, 200, 20), 2)
    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    return image



title = "Face Keypoint"
description = "demo for Face Keypoint. To use it, simply upload your image, or click one of the examples to load them. Read more at the links below."
article = "<p style='text-align: center'><a href='https://www.yuque.com/itmorn/ability/face_keypoint' target='_blank'>Project Documents</a> | <a href='https://www.bilibili.com/video/BV1DN4y1P7j3' target='_blank'>Video Demo</a></p>"

gr.Interface(
    inference,
    [gr.inputs.Image(label="Input")],
    gr.outputs.Image(type="pil", label="Output"),
    title=title,
    description=description,
    article=article,
    examples=[
              ["imgs/face1.jpg"],
              ["imgs/face2.jpg"],
              ["imgs/cc.png"],
              ["imgs/2313.png"]
    ]).launch(debug=True)