update sparse_ga and expose parameter in demo

demo.py

@@ -116,9 +116,10 @@ def get_3D_model_from_scene(outdir, silent, scene, min_conf_thr=2, as_pointcloud
                                         transparent_cams=transparent_cams, cam_size=cam_size, silent=silent)
 
 
-def get_reconstructed_scene(outdir, model, device, silent, image_size, filelist, optim_level, lr1, niter1, lr2, niter2, min_conf_thr,
-                            as_pointcloud, mask_sky, clean_depth, transparent_cams, cam_size,
-                            scenegraph_type, winsize, win_cyclic, refid, TSDF_thresh, shared_intrinsics, **kw):
+def get_reconstructed_scene(outdir, model, device, silent, image_size, 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):
     """
     from a list of images, run mast3r inference, sparse global aligner.
     then run get_3D_model_from_scene
@@ -143,7 +144,8 @@ def get_reconstructed_scene(outdir, model, device, silent, image_size, filelist,
     # Sparse GA (forward mast3r -> matching -> 3D optim -> 2D refinement -> triangulation)
     scene = sparse_global_alignment(filelist, pairs, os.path.join(outdir, 'cache'),
                                     model, lr1=lr1, niter1=niter1, lr2=lr2, niter2=niter2, device=device,
-                                    opt_depth='depth' in optim_level, shared_intrinsics=shared_intrinsics, **kw)
+                                    opt_depth='depth' in optim_level, shared_intrinsics=shared_intrinsics,
+                                    matching_conf_thr=matching_conf_thr, **kw)
     outfile = get_3D_model_from_scene(outdir, silent, scene, min_conf_thr, as_pointcloud, mask_sky,
                                       clean_depth, transparent_cams, cam_size, TSDF_thresh)
     return scene, outfile
@@ -188,27 +190,33 @@ def main_demo(tmpdirname, model, device, image_size, server_name, server_port, s
         with gradio.Column():
             inputfiles = gradio.File(file_count="multiple")
             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")
-                shared_intrinsics = gradio.Checkbox(value=False, label="Shared intrinsics",
-                                                    info="Only optimize one set of intrinsics for all views")
-
-                scenegraph_type = gradio.Dropdown(["complete", "swin", "logwin", "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)
+                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", "swin", "logwin", "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")
 
@@ -238,8 +246,8 @@ def main_demo(tmpdirname, model, device, image_size, server_name, server_port, s
                               inputs=[inputfiles, win_cyclic, refid, scenegraph_type],
                               outputs=[win_col, winsize, win_cyclic, refid])
             run_btn.click(fn=recon_fun,
-                          inputs=[inputfiles, optim_level, lr1, niter1, lr2, niter2, min_conf_thr, as_pointcloud,
-                                  mask_sky, clean_depth, transparent_cams, cam_size,
+                          inputs=[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,

mast3r/cloud_opt/sparse_ga.py

@@ -144,8 +144,8 @@ def sparse_global_alignment(imgs, pairs_in, cache_path, model, subsample=8, desc
     # tmp_pairs = {(a,b):v for (a,b),v in tmp_pairs.items() if {(a,b),(b,a)} & min_spanning_tree}
 
     # smartly combine all usefull data
-    imsizes, pps, base_focals, core_depth, anchors, corres, corres2d = \
-        condense_data(imgs, tmp_pairs, canonical_views, dtype)
+    imsizes, pps, base_focals, core_depth, anchors, corres, corres2d, preds_21 = \
+        condense_data(imgs, tmp_pairs, canonical_views, preds_21, dtype)
 
     imgs, res_coarse, res_fine = sparse_scene_optimizer(
         imgs, subsample, imsizes, pps, base_focals, core_depth, anchors, corres, corres2d, preds_21, canonical_paths, mst,
@@ -345,7 +345,7 @@ def sparse_scene_optimizer(imgs, subsample, imsizes, pps, base_focals, core_dept
             if init[imgs[s.img1]].get('freeze') and init[imgs[s.img2]].get('freeze'):
                 continue
             # fallback to dust3r regression
-            tgt_pts, tgt_confs = subsamp_preds_21[imgs[s.img2]][imgs[s.img1]]
+            tgt_pts, tgt_confs = preds_21[imgs[s.img2]][imgs[s.img1]]
             tgt_pts = geotrf(cam2w[s.img2], tgt_pts)
             cf_sum += tgt_confs.sum()
             loss += tgt_confs @ pix_loss(pts3d[s.img1], tgt_pts)
@@ -733,7 +733,7 @@ PairOfSlices = namedtuple(
     'ImgPair', 'img1, slice1, pix1, anchor_idxs1, img2, slice2, pix2, anchor_idxs2, confs, confs_sum')
 
 
-def condense_data(imgs, tmp_paths, canonical_views, dtype=torch.float32):
+def condense_data(imgs, tmp_paths, canonical_views, preds_21, dtype=torch.float32):
     # aggregate all data properly
     set_imgs = set(imgs)
 
@@ -809,7 +809,16 @@ def condense_data(imgs, tmp_paths, canonical_views, dtype=torch.float32):
     imsizes = torch.tensor([(W, H) for H, W in shapes], device=pp.device)  # (W,H)
     principal_points = torch.stack(principal_points)
     focals = torch.cat(focals)
-    return imsizes, principal_points, focals, core_depth, img_anchors, corres, corres2d
+
+    # Subsample preds_21
+    subsamp_preds_21 = {}
+    for imk, imv in preds_21.items():
+        subsamp_preds_21[imk] = {}
+        for im2k, (pred, conf) in preds_21[imk].items():
+            idxs = img_anchors[imgs.index(im2k)][1]
+            subsamp_preds_21[imk][im2k] = (pred[idxs], conf[idxs])  # anchors subsample
+
+    return imsizes, principal_points, focals, core_depth, img_anchors, corres, corres2d, subsamp_preds_21
 
 
 def canonical_view(ptmaps11, confs11, subsample, mode='avg-angle'):