#This code is from: https://huggingface.co/spaces/diffusers/controlnet-3d-pose
import gradio as gr
from PIL import Image, ImageFilter
from io import BytesIO
import base64 
canvas_html = '<div style="height:50%"><pose-maker/></div>'

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 generate_images(canvas):        
    base64_img = canvas
    image_data = base64.b64decode(base64_img.split(',')[1])
    input_img = Image.open(BytesIO(image_data)).convert(
        'RGB').resize((512, 512))
    input_img = input_img.filter(ImageFilter.GaussianBlur(radius=2))
    input_img = get_canny_filter(input_img)
    return input_image
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")
            run_button = gr.Button("Generate")
    
    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],
                     outputs=[gr.Image()],
                     _js=get_js_image)
    b.load(None,None,None,_js=load_js)

b.launch()