Update app.py

app.py

@@ -262,7 +262,7 @@ def pano_depth_to_world_points(depth):
 def rgb2gray(rgb):
     return np.dot(rgb[...,:3], [0.333, 0.333, 0.333])
 
-def get_mesh(image, depth, blur_data, loadall):
+def get_mesh(image, depth, blur_data, loadall, dcolor):
     global frame_selected
     global mesh
     global mesh_n
@@ -328,7 +328,7 @@ def get_mesh(image, depth, blur_data, loadall):
     glb_file = tempfile.NamedTemporaryFile(suffix='.glb', delete=False)
     scene.export(glb_file.name)
     print('file - ok')
-    return glb_file.name, glb_file.name, ",".join(mesh_n)
+    return gr.Model3D(value=glb_file.name, clear_color=json.loads(dcolor)), glb_file.name, ",".join(mesh_n)
 
 def blur_image(image, depth, blur_data):
     blur_a = blur_data.split()
@@ -413,9 +413,9 @@ def optimize(v, d):
             cv2.imwrite(depths[k], depth)
             
     if d == False:
-      return frames
+      return frames, dcolor
     else:
-      return depths
+      return depths, dcolor
 
 def bincount(a):
     a2D = a.reshape(-1,a.shape[-1])
@@ -522,12 +522,13 @@ def draw_mask(l, t, v, d, evt: gr.EventData):
 
 
 load_model="""
-async(f, c, o)=>{
+async(f, c, o, b)=>{
   if (document.getElementById("iframe3D")) {
     document.getElementById("iframe3D").src = "index.htm?file=" + f.path;
     document.getElementById("iframe3D").onload = function() {
       document.getElementById("iframe3D").contentDocument.getElementById("coords").value = c;
       document.getElementById("iframe3D").contentDocument.getElementById("order").value = o;
+      document.getElementById("iframe3D").contentDocument.getElementById("bgcolor").value = b;
     }
     toggleDisplay("model");
   }
@@ -629,7 +630,7 @@ with gr.Blocks(css=css, js=js) as demo:
             output_switch = gr.Checkbox(label="Show depths")
             optimize_switch = gr.Checkbox(label="Optimize")
             output_switch.input(fn=switch_rows, inputs=[output_switch], outputs=[output_frame])
-            optimize_switch.input(fn=optimize, inputs=[optimize_switch, output_switch], outputs=[output_frame])
+            optimize_switch.input(fn=optimize, inputs=[optimize_switch, output_switch], outputs=[output_frame, bgcolor])
             output_mask = gr.ImageEditor(layers=False, sources=('upload', 'clipboard'), show_download_button=True, type="numpy", interactive=True, transforms=(None,), eraser=gr.Eraser(), brush=gr.Brush(default_size=0, colors=['black', '#505050', '#a0a0a0', 'white']), elem_id="image_edit")
             with gr.Row():
               selector = gr.HTML(value="""
@@ -933,9 +934,10 @@ with gr.Blocks(css=css, js=js) as demo:
                 ]"""
                 coords = gr.Textbox(elem_id="coords", value=example_coords, label="Coordinates", interactive=False)
                 mesh_order = gr.Textbox(elem_id="order", value="", label="Order", interactive=False)
+                bgcolor = gr.Textbox(elem_id="bgcolor", value="[127, 127, 127, 255]", label="Background color", interactive=False)
                 
             result_file = gr.File(label="3D file (.glb)")
-            result_file.change(fn=None, inputs=[result_file, coords, mesh_order], outputs=None, js=load_model)
+            result_file.change(fn=None, inputs=[result_file, coords, mesh_order, bgcolor], outputs=None, js=load_model)
             html = gr.HTML(value="""<label for='zoom'>Zoom</label><input id='zoom' type='range' style='width:256px;height:1em;' value='0.8' min='0.157' max='1.57' step='0.001' oninput='
               if (!BABYLON.Engine.LastCreatedScene.activeCamera.metadata) {
                 var evt = document.createEvent(\"Event\");
@@ -1091,7 +1093,7 @@ with gr.Blocks(css=css, js=js) as demo:
         return output_video_path + (json.dumps(locations),)
 
     submit.click(on_submit, inputs=[input_video, model_type, coords], outputs=[processed_video, processed_zip, output_frame, output_mask, coords])
-    render.click(partial(get_mesh), inputs=[output_frame, output_mask, blur_in, load_all], outputs=[result, result_file, mesh_order])
+    render.click(partial(get_mesh), inputs=[output_frame, output_mask, blur_in, load_all, bgcolor], outputs=[result, result_file, mesh_order])
 
     example_files = os.listdir('examples')
     example_files.sort()