from models.model import Model as AutoLink
import gradio as gr
import PIL
import torch
import os
import imageio
import numpy as np

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


autolink = AutoLink.load_from_checkpoint(os.path.join("checkpoints", "celeba_wild_k32_m0.8_b16_t0.00075_sklr512", "model.ckpt"))
autolink.to(device)


def predict_image(image_in: PIL.Image.Image) -> PIL.Image.Image:
    if image_in == None:
        raise gr.Error("Please upload a video or image.")
    edge_map = autolink(image_in)
    return edge_map


def predict_video(video_in: str) -> str:
    if video_in == None:
        raise gr.Error("Please upload a video or image.")
    video_out = video_in[:-4] + '_out.mp4'
    video_in = imageio.get_reader(video_in)
    writer = imageio.get_writer(video_out, mode='I', fps=video_in.get_meta_data()['fps'])
    for image_in in video_in:
        image_in = PIL.Image.fromarray(image_in)
        edge_map = autolink(image_in)
        writer.append_data(np.array(edge_map))
    writer.close()
    return video_out


with gr.Blocks() as blocks:
    gr.Markdown("""
    # AutoLink
    ## Self-supervised Learning of Human Skeletons and Object Outlines by Linking Keypoints
    ## This demo is specifically for self-supervised facial landmark detection
    #### Note that there is no face detection in this demo, so please make sure that the face is center-around in the image.
    * [Paper](https://arxiv.org/abs/2205.10636)
    * [Project Page](https://xingzhehe.github.io/autolink/)
    * [GitHub](https://github.com/xingzhehe/AutoLink-Self-supervised-Learning-of-Human-Skeletons-and-Object-Outlines-by-Linking-Keypoints)
""")
    
    with gr.Tab("Image"):
        with gr.Row():
            with gr.Column():
                image_in = gr.Image(source="upload",  type="pil", visible=True)
            with gr.Column():
                image_out = gr.Image()
        run_btn = gr.Button("Run")
        run_btn.click(fn=predict_image, inputs=[image_in], outputs=[image_out])
        gr.Examples(fn=predict_image, examples=[["assets/jackie_chan.jpg", None]],
                    inputs=[image_in], outputs=[image_out],
                    cache_examples=True)
        
    with gr.Tab("Video") as tab:
        with gr.Row():
            with gr.Column():
                video_in = gr.Video(source="upload", type="mp4")
            with gr.Column():
                video_out = gr.Video()
        run_btn = gr.Button("Run")
        run_btn.click(fn=predict_video, inputs=[video_in], outputs=[video_out])
        gr.Examples(fn=predict_video, examples=[["assets/00344.mp4"],],
                    inputs=[video_in], outputs=[video_out],
                    cache_examples=True)

blocks.launch()