Update app.py

app.py

@@ -3,21 +3,61 @@ import numpy as np
 import cv2
 from PIL import Image
 import torch
+import base64
+import requests 
+from io import BytesIO
 from region_control import MultiDiffusion, get_views, preprocess_mask
 from sketch_helper import get_high_freq_colors, color_quantization, create_binary_matrix
 MAX_COLORS = 12
 
 sd = MultiDiffusion("cuda", "2.0")
 
-def process_sketch(image, binary_matrixes):
-  high_freq_colors, image = get_high_freq_colors(image)
-  how_many_colors = len(high_freq_colors)
-  im2arr = np.array(image) # im2arr.shape: height x width x channel
-  im2arr = color_quantization(im2arr, high_freq_colors)
+canvas_html = "<div id='canvas-root'></div>"
+load_js = """
+async () => {
+const url = "https://huggingface.co/datasets/radames/gradio-components/raw/main/sketch-canvas.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_colors = """
+async (canvasData) => {
+  const canvasEl = document.getElementById("canvas-root");
+  return [canvasEl._data]
+}
+"""
+
+set_canvas_size ="""
+async (aspect) => {
+  if(aspect ==='square'){
+    _updateCanvas(512,512)
+  }
+  if(aspect ==='horizontal'){
+    _updateCanvas(768,512)
+  }
+  if(aspect ==='vertical'){
+    _updateCanvas(512,768)
+  }
+}
+"""
+
+def process_sketch(canvas_data, binary_matrixes):
+  base64_img = canvas_data['image']
+  image_data = base64.b64decode(base64_img.split(',')[1])
+  image = Image.open(BytesIO(image_data))
+  im2arr = np.array(image)
   
+  colors = [tuple(int(color.lstrip('#')[i:i+2], 16) for i in (0, 2, 4)) for color in canvas_data['colors']] 
   colors_fixed = []
-  for color in high_freq_colors:
-    r, g, b = color[1]
+  for color in colors:
+    r, g, b = color
     if any(c != 255 for c in (r, g, b)):
       binary_matrix = create_binary_matrix(im2arr, (r,g,b))
       binary_matrixes.append(binary_matrix)
@@ -27,7 +67,7 @@ def process_sketch(image, binary_matrixes):
   for n in range(MAX_COLORS):
     visibilities.append(gr.update(visible=False))
     colors.append(gr.update(value=f'<div class="color-bg-item" style="background-color: black"></div>'))
-  for n in range(how_many_colors-1):
+  for n in range(len(colors)-1):
     visibilities[n] = gr.update(visible=True)
     colors[n] = colors_fixed[n]
   return [gr.update(visible=True), binary_matrixes, *visibilities, *colors]
@@ -65,7 +105,10 @@ with gr.Blocks(css=css) as demo:
   with gr.Row():
     with gr.Box(elem_id="main-image"):
       #with gr.Row():
-      image = gr.Image(interactive=True, tool="color-sketch", source="canvas", type="pil", shape=(512,512), brush_radius=45)
+      canvas_data = gr.JSON(value={}, visible=False)
+      canvas = gr.HTML(canvas_html)
+      #image = gr.Image(interactive=True, tool="color-sketch", source="canvas", type="pil", shape=(512,512), brush_radius=45)
+      
       #image_horizontal = gr.Image(interactive=True, tool="color-sketch", source="canvas", type="pil", shape=(768,512), visible=False, brush_radius=45)
       #image_vertical = gr.Image(interactive=True, tool="color-sketch", source="canvas", type="pil", shape=(512, 768), visible=False, brush_radius=45)
       #with gr.Row():
@@ -90,6 +133,7 @@ with gr.Blocks(css=css) as demo:
   ''')
   #css_height = gr.HTML("<style>#main-image{width: 512px} .fixed-height{height: 512px !important}</style>")
   #aspect.change(update_css, inputs=aspect, outputs=[image, image_horizontal, image_vertical])
-  button_run.click(process_sketch, inputs=[image, binary_matrixes], outputs=[post_sketch, binary_matrixes, *color_row, *colors])
+  button_run.click(process_sketch, inputs=[canvas_data, binary_matrixes], outputs=[post_sketch, binary_matrixes, *color_row, *colors])
   final_run_btn.click(process_generation, inputs=[binary_matrixes, general_prompt, *prompts], outputs=out_image)
+  demo.load(None, None, None, _js=load_js)
 demo.launch(debug=True)