import gradio as gr
import json as js
import util
from fileservice import app
from fastapi.staticfiles import StaticFiles
from pose import infer, draw
from abgr import remove_bg, split_image

def image_changed(image):
  if image == None:
    return "estimation", {}

  print("make mask")
  cvimage = util.pil2cv(image)
  mask, fg, bg = split_image(cvimage[..., ::-1])

  print("pose not found")
  pose_result, _ = infer(cvimage)

  candidate, subset = util.convert_to_openpose(pose_result)
  candidateJson = util.candidate_to_json_string(candidate)
  subsetJson = util.subset_to_json_string(subset)
  jsonText = f'{{"candidate":{candidateJson}, "subset":{subsetJson}, "width":{image.width}, "height":{image.height}}}'
  return f'{image.width}px x {image.height}px, {len(subset)} indivisual(s)', jsonText, mask, fg, bg

with gr.Blocks(css="""button { min-width: 80px; }""") as demo:
  with gr.Row():
    with gr.Column(scale=1):
      source = gr.Image(type="pil")
      info = gr.Markdown("""info""")
      gr.Examples(
        examples=["static/sample1.png", "static/sample2.png", "static/sample3.png"],
        inputs=source,
      )      
      btn = gr.Button("Import")  
      with gr.Accordion(label="Parts", open=False):
        mask = gr.Image()
        frontImage = gr.Image(image_mode="RGBA")
        backImage = gr.Image(image_mode="RGBA")
        gr.Markdown("""
#### Reference
Pose estimation: [MMPose](https://github.com/open-mmlab/mmpose)

Image splitting: [anime-segmentation](https://github.com/SkyTNT/anime-segmentation/)
""")   
      with gr.Accordion(label="Json", open=False):
        json = gr.JSON(label="Json")
    with gr.Column(scale=3):
      gr.HTML('<canvas id="canvas" width="512" height="512"></canvas>')
  with gr.Row():
    gr.HTML('<canvas id="canvas2" width="512" height="512" style="display:none"></canvas>')
    gr.HTML('<canvas id="canvas3" width="512" height="512" style="display:none"></canvas>')
    gr.HTML('<canvas id="canvas4" width="512" height="512" style="display:none"></canvas>')
    gr.HTML('<canvas id="microscope" width="50" height="50" style="display:none"></canvas>')

  source.change(
    fn = image_changed,
    inputs = [source],
    outputs = [info, json, mask, frontImage, backImage])

  btn.click(
    fn = None,
    inputs = [frontImage, backImage, json],
    outputs = [],
    _js="(frontImage, backImage, json) => { initializeEditor(); importPose(json); importPicture(frontImage); importBackground(backImage); return []; }")

  demo.load(fn=None, inputs=[], outputs=[], _js="() => { initializeEditor(); importPose(); return []; }")

print("mount")
app.mount("/static", StaticFiles(directory="static"), name="static")
app.mount("/js", StaticFiles(directory="js"), name="js")
gr.mount_gradio_app(app, demo, path="/")