#This code is largely based on: https://huggingface.co/spaces/diffusers/controlnet-3d-pose
import gradio as gr
from PIL import Image, ImageFilter, ImageOps
from io import BytesIO
import numpy as np 
import base64
import cv2
import os
import uuid
uid=uuid.uuid4()
canvas_html = '<div style="height:50%"><pose-maker/></div>'
low_threshold = 100
high_threshold = 200
load_js = """
async () => {
  const url = "https://huggingface.co/datasets/mishig/gradio-components/raw/main/mannequinAll.js"
  fetch(url)
    .then(res => res.text())
    .then(text => {
      const script = document.createElement('script');
      script.type = "module"
      script.src = URL.createObjectURL(new Blob([text], { type: 'application/javascript' }));
      document.head.appendChild(script);
    });
}
"""
get_js_image = """
async (canvas, prompt) => {
  const poseMakerEl = document.querySelector("pose-maker");
  const imgBase64 = poseMakerEl.captureScreenshotDepthMap();
  return [imgBase64, prompt]
}
"""
js_change_rotation_axis = """
async (axis) => {
  const poseMakerEl = document.querySelector("pose-maker");
  poseMakerEl.changeRotationAxis(axis);
}
"""
js_pose_template = """
async (pose) => {
  const poseMakerEl = document.querySelector("pose-maker");
  poseMakerEl.setPose(pose);
}
"""
def get_canny_filter(image):
    if not isinstance(image, np.ndarray):
        image = np.array(image) 
    image = cv2.Canny(image, low_threshold, high_threshold)
    image = image[:, :, None]
    image = np.concatenate([image, image, image], axis=2)
    canny_image = Image.fromarray(image)
    return canny_image
def generate_images(canvas,output_type):        
    base64_img = canvas
    image_data = base64.b64decode(base64_img.split(',')[1])
    
    
    left =128
    top = 0
    right = 384
    bottom = 512   

    input_img = Image.open(BytesIO(image_data)).convert(
        'RGB').resize((512,512)).crop((left, top, right, bottom))

    #input_img = input_img.filter(ImageFilter.GaussianBlur(radius=2))
    input_img = ImageOps.invert(get_canny_filter(input_img)) if output_type == "Canny" else input_img
    input_img = ImageOps.invert(input_img) if output_type == "Depth" else input_img
    #input_img = Image.open(BytesIO(base64_img)).resize((512,512)).crop((left, top, right, bottom)) if output_type == "Base" else input_img
    
    #input_img = ImageOps.invert(input_img)
    input_img.save(f"{uid}-pose.png")
    out = os.path.abspath(f"{uid}-pose.png")
    out_url = f'https://omnibus-model-mover.hf.space/file={out}'       
    return f"{uid}-pose.png",f"{uid}-pose.png",out_url
def placeholder_fn(axis):
    pass
with gr.Blocks() as b:
    with gr.Row():
        canvas = gr.HTML(canvas_html, elem_id="canvas_html", visible=True)
        with gr.Column():
            rotation_axis = gr.Radio(choices=["x", "y", "z"], value="x", label="Joint rotation axis")
            pose_template = gr.Radio(choices=["regular", "ballet", "handstand", "split", "kick", "chilling"], value="regular", label="Pose template")
            output_type = gr.Radio(choices=["Depth", "Canny"], value="Depth", label="Image Output")
            run_button = gr.Button("Generate")
            out_text = gr.Textbox(label="Image URL")
            out_file=gr.File()
            out_im = gr.Image(type='filepath',show_share_button=True)
    
    rotation_axis.change(fn=placeholder_fn,
                            inputs=[rotation_axis],
                            outputs=[],
                            queue=False,
                            _js=js_change_rotation_axis)
    pose_template.change(fn=placeholder_fn,
                            inputs=[pose_template],
                            outputs=[],
                            queue=False,
                            _js=js_pose_template)
    run_button.click(fn=generate_images,
                     inputs=[canvas,output_type],
                     outputs=[out_im,out_file,out_text],
                     _js=get_js_image)
    b.load(None,None,None,_js=load_js)

b.launch()