import cv2
import gradio as gr
import io
import numpy as np

from z_app_factory import get_app

thickness = 3
lineType = 8
font = cv2.FONT_HERSHEY_SIMPLEX

def plot3(fig):
    with io.BytesIO() as buff:
        fig.savefig(buff, format='raw')
        buff.seek(0)
        data = np.frombuffer(buff.getvalue(), dtype=np.uint8)
    w, h = fig.canvas.get_width_height()
    im = data.reshape((int(h), int(w), -1))
    im = im[:, :, :-1]
    # im = im[:, :, ::-1]
    im = im[125:350,220:440,:]
    # cv2.imwrite("res.png", im)
    # 90/0
    return im


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"
    import numpy as np
    img_out = np.zeros((500, 600, 3), dtype=np.uint8)
    if code!=200:
        cv2.putText(img_out, msg, (20, 200), font, 1, (0, 255, 0), 2)
        return img_out

    if code==200 and not lst2d_res:
        cv2.putText(img_out, "no face detected", (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
    landmarks = lst2d_res['landmarks']
    import numpy as np  # 用来处理数据
    import matplotlib
    import matplotlib.pyplot as plt
    matplotlib.use('agg')
    x = np.array([i[0] for i in landmarks])
    y = np.array([i[1] for i in landmarks])
    z = np.array([i[2] for i in landmarks])
    z = z - z.min()
    z = z / z.max()
    fig = plt.figure()
    ax = plt.subplot(projection='3d')  # 创建一个三维的绘图工程
    ax.scatter(x, y, z, c=[int(i) for i in z * 255])  # 绘制数据点 c: 'r'红色，'y'黄色，等颜色

    # ax.set_xlabel('X')  # 设置x坐标轴
    # ax.set_ylabel('Y')  # 设置y坐标轴
    # ax.set_zlabel('Z')  # 设置z坐标轴
    ax.view_init(-90, -90)
    im_plot = plot3(fig)
    plt.close()
    h, w, _ = im_plot.shape
    # image[:h, :w] = im_plot
    return im_plot



title = "Face Keypoint 3d"
description = "demo for Face Keypoint 3d. 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_3d' target='_blank'>Project Documents</a> | <a href='https://www.bilibili.com/video/BV1GD4y1y77z' 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)