dumb down demo

app.py

@@ -12,12 +12,13 @@ import torch
 import tempfile
 import gradio
 import shutil
+import math
 
 HERE_PATH = path.normpath(path.dirname(__file__))  # noqa
 MASt3R_REPO_PATH = path.normpath(path.join(HERE_PATH, './mast3r'))  # noqa
 sys.path.insert(0, MASt3R_REPO_PATH)  # noqa
 
-from mast3r.demo import get_reconstructed_scene, get_3D_model_from_scene, set_scenegraph_options
+from mast3r.demo import get_reconstructed_scene
 from mast3r.model import AsymmetricMASt3R
 from mast3r.utils.misc import hash_md5
 
@@ -39,133 +40,74 @@ silent = True
 gradio_delete_cache = 7200
 
 
-@spaces.GPU()
-def local_get_reconstructed_scene(current_scene_state,
-                                  filelist, optim_level, lr1, niter1, lr2, niter2, min_conf_thr, matching_conf_thr,
-                                  as_pointcloud, mask_sky, clean_depth, transparent_cams, cam_size, scenegraph_type, winsize,
-                                  win_cyclic, refid, TSDF_thresh, shared_intrinsics, **kw):
-    return get_reconstructed_scene(tmpdirname, gradio_delete_cache, model, device, silent, image_size, current_scene_state,
-                                   filelist, optim_level, lr1, niter1, lr2, niter2, min_conf_thr, matching_conf_thr,
-                                   as_pointcloud, mask_sky, clean_depth, transparent_cams, cam_size, scenegraph_type, winsize,
-                                   win_cyclic, refid, TSDF_thresh, shared_intrinsics, **kw)
-
-
-@spaces.GPU()
-def local_get_3D_model_from_scene(scene_state, min_conf_thr=2, as_pointcloud=False, mask_sky=False,
-                                  clean_depth=False, transparent_cams=False, cam_size=0.05, TSDF_thresh=0):
-    return get_3D_model_from_scene(silent, scene_state, min_conf_thr, as_pointcloud, mask_sky,
-                                   clean_depth, transparent_cams, cam_size, TSDF_thresh)
+class FileState:
+    def __init__(self, outfile_name=None):
+        self.outfile_name = outfile_name
 
+    def __del__(self):
+        if self.outfile_name is not None and os.path.isfile(self.outfile_name):
+            os.remove(self.outfile_name)
+        self.outfile_name = None
 
-recon_fun = local_get_reconstructed_scene
-model_from_scene_fun = local_get_3D_model_from_scene
 
-
-def get_context(delete_cache):
-    css = """.gradio-container {margin: 0 !important; min-width: 100%};"""
-    title = "MASt3R Demo"
-    if delete_cache:
-        return gradio.Blocks(css=css, title=title, delete_cache=(delete_cache, delete_cache))
+@spaces.GPU()
+def local_get_reconstructed_scene(filelist, min_conf_thr, matching_conf_thr,
+                                  as_pointcloud, cam_size,
+                                  shared_intrinsics, **kw):
+    lr1 = 0.07
+    niter1 = 500
+    lr2 = 0.014
+    niter2 = 200
+    optim_level = 'refine'
+    mask_sky, clean_depth, transparent_cams = False, True, False
+    if len(filelist) < 5:
+        scenegraph_type = 'complete'
+        winsize = 1
     else:
-        return gradio.Blocks(css=css, title="MASt3R Demo")  # for compatibility with older versions
-
-
-with get_context(gradio_delete_cache) as demo:
-    # scene state is save so that you can change conf_thr, cam_size... without rerunning the inference
-    scene = gradio.State(None)
+        scenegraph_type = 'logwin'
+        half_size = math.ceil((len(filelist) - 1) / 2)
+        max_winsize = max(1, math.ceil(math.log(half_size, 2)))
+        winsize = min(5, max_winsize)
+    refid = 0
+    win_cyclic = False
+    scene_state, outfile = get_reconstructed_scene(tmpdirname, gradio_delete_cache, model, device, silent, image_size, None,
+                                                   filelist, optim_level, lr1, niter1, lr2, niter2, min_conf_thr, matching_conf_thr,
+                                                   as_pointcloud, mask_sky, clean_depth, transparent_cams, cam_size, scenegraph_type, winsize,
+                                                   win_cyclic, refid, TSDF_thresh=0, shared_intrinsics=shared_intrinsics, **kw)
+    filestate = FileState(scene_state.outfile_name)
+    scene_state.outfile_name = None
+    del scene_state
+    return filestate, outfile
+
+
+css = """.gradio-container {margin: 0 !important; min-width: 100%};"""
+title = "MASt3R Demo"
+with gradio.Blocks(css=css, title=title, delete_cache=(gradio_delete_cache, gradio_delete_cache)) as demo:
+    filestate = gradio.State(None)
     gradio.HTML('<h2 style="text-align: center;">MASt3R Demo</h2>')
     with gradio.Column():
         inputfiles = gradio.File(file_count="multiple")
         with gradio.Row():
-            with gradio.Column():
-                with gradio.Row():
-                    lr1 = gradio.Slider(label="Coarse LR", value=0.07, minimum=0.01, maximum=0.2, step=0.01)
-                    niter1 = gradio.Number(value=500, precision=0, minimum=0, maximum=10_000,
-                                           label="num_iterations", info="For coarse alignment!")
-                    lr2 = gradio.Slider(label="Fine LR", value=0.014, minimum=0.005, maximum=0.05, step=0.001)
-                    niter2 = gradio.Number(value=200, precision=0, minimum=0, maximum=100_000,
-                                           label="num_iterations", info="For refinement!")
-                    optim_level = gradio.Dropdown(["coarse", "refine", "refine+depth"],
-                                                  value='refine', label="OptLevel",
-                                                  info="Optimization level")
-                with gradio.Row():
-                    matching_conf_thr = gradio.Slider(label="Matching Confidence Thr", value=5.,
-                                                      minimum=0., maximum=30., step=0.1,
-                                                      info="Before Fallback to Regr3D!")
-                    shared_intrinsics = gradio.Checkbox(value=False, label="Shared intrinsics",
-                                                        info="Only optimize one set of intrinsics for all views")
-                    scenegraph_type = gradio.Dropdown([("complete: all possible image pairs", "complete"),
-                                                       ("swin: sliding window", "swin"),
-                                                       ("logwin: sliding window with long range", "logwin"),
-                                                       ("oneref: match one image with all", "oneref")],
-                                                      value='complete', label="Scenegraph",
-                                                      info="Define how to make pairs",
-                                                      interactive=True)
-                    with gradio.Column(visible=False) as win_col:
-                        winsize = gradio.Slider(label="Scene Graph: Window Size", value=1,
-                                                minimum=1, maximum=1, step=1)
-                        win_cyclic = gradio.Checkbox(value=False, label="Cyclic sequence")
-                    refid = gradio.Slider(label="Scene Graph: Id", value=0,
-                                          minimum=0, maximum=0, step=1, visible=False)
-        run_btn = gradio.Button("Run")
-
-        with gradio.Row():
+            matching_conf_thr = gradio.Slider(label="Matching Confidence Thr", value=5.,
+                                              minimum=0., maximum=30., step=0.1,
+                                              info="Before Fallback to Regr3D!")
             # adjust the confidence threshold
             min_conf_thr = gradio.Slider(label="min_conf_thr", value=1.5, minimum=0.0, maximum=10, step=0.1)
             # adjust the camera size in the output pointcloud
             cam_size = gradio.Slider(label="cam_size", value=0.2, minimum=0.001, maximum=1.0, step=0.001)
-            TSDF_thresh = gradio.Slider(label="TSDF Threshold", value=0., minimum=0., maximum=1., step=0.01)
         with gradio.Row():
             as_pointcloud = gradio.Checkbox(value=True, label="As pointcloud")
-            # two post process implemented
-            mask_sky = gradio.Checkbox(value=False, label="Mask sky")
-            clean_depth = gradio.Checkbox(value=True, label="Clean-up depthmaps")
-            transparent_cams = gradio.Checkbox(value=False, label="Transparent cameras")
-
+            shared_intrinsics = gradio.Checkbox(value=False, label="Shared intrinsics",
+                                                info="Only optimize one set of intrinsics for all views")
+        run_btn = gradio.Button("Run")
         outmodel = gradio.Model3D()
 
         # events
-        scenegraph_type.change(set_scenegraph_options,
-                               inputs=[inputfiles, win_cyclic, refid, scenegraph_type],
-                               outputs=[win_col, winsize, win_cyclic, refid])
-        inputfiles.change(set_scenegraph_options,
-                          inputs=[inputfiles, win_cyclic, refid, scenegraph_type],
-                          outputs=[win_col, winsize, win_cyclic, refid])
-        win_cyclic.change(set_scenegraph_options,
-                          inputs=[inputfiles, win_cyclic, refid, scenegraph_type],
-                          outputs=[win_col, winsize, win_cyclic, refid])
-        run_btn.click(fn=recon_fun,
-                      inputs=[scene, inputfiles, optim_level, lr1, niter1, lr2, niter2, min_conf_thr, matching_conf_thr,
-                              as_pointcloud, mask_sky, clean_depth, transparent_cams, cam_size,
-                              scenegraph_type, winsize, win_cyclic, refid, TSDF_thresh, shared_intrinsics],
-                      outputs=[scene, outmodel])
-        min_conf_thr.release(fn=model_from_scene_fun,
-                             inputs=[scene, min_conf_thr, as_pointcloud, mask_sky,
-                                     clean_depth, transparent_cams, cam_size, TSDF_thresh],
-                             outputs=outmodel)
-        cam_size.change(fn=model_from_scene_fun,
-                        inputs=[scene, min_conf_thr, as_pointcloud, mask_sky,
-                                clean_depth, transparent_cams, cam_size, TSDF_thresh],
-                        outputs=outmodel)
-        TSDF_thresh.change(fn=model_from_scene_fun,
-                           inputs=[scene, min_conf_thr, as_pointcloud, mask_sky,
-                                   clean_depth, transparent_cams, cam_size, TSDF_thresh],
-                           outputs=outmodel)
-        as_pointcloud.change(fn=model_from_scene_fun,
-                             inputs=[scene, min_conf_thr, as_pointcloud, mask_sky,
-                                     clean_depth, transparent_cams, cam_size, TSDF_thresh],
-                             outputs=outmodel)
-        mask_sky.change(fn=model_from_scene_fun,
-                        inputs=[scene, min_conf_thr, as_pointcloud, mask_sky,
-                                clean_depth, transparent_cams, cam_size, TSDF_thresh],
-                        outputs=outmodel)
-        clean_depth.change(fn=model_from_scene_fun,
-                           inputs=[scene, min_conf_thr, as_pointcloud, mask_sky,
-                                   clean_depth, transparent_cams, cam_size, TSDF_thresh],
-                           outputs=outmodel)
-        transparent_cams.change(model_from_scene_fun,
-                                inputs=[scene, min_conf_thr, as_pointcloud, mask_sky,
-                                        clean_depth, transparent_cams, cam_size, TSDF_thresh],
-                                outputs=outmodel)
+        run_btn.click(fn=local_get_reconstructed_scene,
+                      inputs=[inputfiles, min_conf_thr, matching_conf_thr,
+                              as_pointcloud,
+                              cam_size, shared_intrinsics],
+                      outputs=[filestate, outmodel])
+
 demo.launch(share=None, server_name=None, server_port=None)
 shutil.rmtree(tmpdirname)