fix: Update demo

demo.py

@@ -15,6 +15,7 @@ from spann3r.datasets import *
 from torch.utils.data import DataLoader
 from spann3r.tools.eval_recon import accuracy, completion
 from spann3r.tools.vis import render_frames, find_render_cam, vis_pred_and_imgs
+from backend_utils import improved_multiway_registration, pts2normal, point2mesh, combine_and_clean_point_clouds
 
 def get_args_parser():
     parser = argparse.ArgumentParser('Spann3R demo', add_help=False)
@@ -27,9 +28,28 @@ def get_args_parser():
     parser.add_argument('--conf_thresh', type=float, default=1e-3, help='confidence threshold')
     parser.add_argument('--kf_every', type=int, default=10, help='map every kf_every frames')
     parser.add_argument('--vis', action='store_true', help='visualize')
-
+    parser.add_argument('--voxel_size', type=float, default=0.004, help='voxel size for multiway registration')
     return parser
 
+import tempfile
+import subprocess
+def extract_frames(video_path: str, duration: float = 20.0, fps: float = 3.0) -> str:
+    temp_dir = tempfile.mkdtemp()
+    output_path = os.path.join(temp_dir, "%03d.jpg")
+    
+    filter_complex = f"select='if(lt(t,{duration}),1,0)',fps={fps}"
+
+    command = [
+        "ffmpeg",
+        "-i", video_path,
+        "-vf", filter_complex,
+        "-vsync", "0",
+        output_path
+    ]
+    
+    subprocess.run(command, check=True)
+    return temp_dir
+
 @torch.no_grad()
 def main(args):
 
@@ -42,6 +62,10 @@ def main(args):
     
     model.load_state_dict(torch.load(args.ckpt_path)['model'])
     model.eval()
+    
+    if args.demo_path.endswith('.mp4') or args.demo_path.endswith('.avi') or args.demo_path.endswith('.MOV'):
+        args.demo_path = extract_frames(args.demo_path)
+        args.kf_every = 1
 
     ##### Load dataset
     dataset = Demo(ROOT=args.demo_path, resolution=224, full_video=True, kf_every=args.kf_every)
@@ -96,59 +120,51 @@ def main(args):
     os.makedirs(save_demo_path, exist_ok=True)
 
     pts_all = []
+    pts_normal_all = []
     pts_gt_all = []
     images_all = []
     masks_all = []
-    conf_all = []
+    conf_sig_all = []
+    cameras_all = []
 
+    last_focal = None
     for j, view in enumerate(ordered_batch):
         
         image = view['img'].permute(0, 2, 3, 1).cpu().numpy()[0]
         mask = view['valid_mask'].cpu().numpy()[0]
 
         pts = preds[j]['pts3d' if j==0 else 'pts3d_in_other_view'].detach().cpu().numpy()[0]
+        pts_normal = pts2normal(preds[j]['pts3d' if j==0 else 'pts3d_in_other_view'][0]).cpu().numpy()
         conf = preds[j]['conf'][0].cpu().data.numpy()
-
+        conf_sig = (conf - 1) / conf
         pts_gt = view['pts3d'].cpu().numpy()[0]
 
         images_all.append((image[None, ...] + 1.0)/2.0)
         pts_all.append(pts[None, ...])
+        pts_normal_all.append(pts_normal[None, ...])
         pts_gt_all.append(pts_gt[None, ...])
         masks_all.append(mask[None, ...])
-        conf_all.append(conf[None, ...])
-    
+        conf_sig_all.append(conf_sig[None, ...])
+
     images_all = np.concatenate(images_all, axis=0)
     pts_all = np.concatenate(pts_all, axis=0)
+    pts_normal_all = np.concatenate(pts_normal_all, axis=0)
     pts_gt_all = np.concatenate(pts_gt_all, axis=0)
     masks_all = np.concatenate(masks_all, axis=0)
-    conf_all = np.concatenate(conf_all, axis=0)
-
-    save_params = dict(
-        images_all=images_all,
-        pts_all=pts_all,
-        pts_gt_all=pts_gt_all,
-        masks_all=masks_all,
-        conf_all=conf_all
-        )
-    
-    np.save(os.path.join(save_demo_path, f"{demo_name}.npy"), save_params)
-
-
-    # Save point cloud
-    conf_sig_all = (conf_all-1) / conf_all
-
-    pcd = o3d.geometry.PointCloud()
-    pcd.points = o3d.utility.Vector3dVector(pts_all[conf_sig_all>args.conf_thresh].reshape(-1, 3))
-    pcd.colors = o3d.utility.Vector3dVector(images_all[conf_sig_all>args.conf_thresh].reshape(-1, 3))
-    o3d.io.write_point_cloud(os.path.join(save_demo_path, f"{demo_name}_conf{args.conf_thresh}.ply"), pcd)
-
-
-    if args.vis:
-        camera_parameters = find_render_cam(pcd)
-
-        render_frames(pts_all, images_all, camera_parameters, save_demo_path, mask=conf_sig_all>args.conf_thresh)
-        vis_pred_and_imgs(pts_all, save_demo_path, images_all=images_all, conf_all=conf_sig_all)
-
+    conf_sig_all = np.concatenate(conf_sig_all, axis=0)
+
+    # Create point clouds for multiway registration
+    pcds = []
+    for j in range(len(pts_all)):
+        pcd = o3d.geometry.PointCloud()
+        mask = conf_sig_all[j] > args.conf_thresh
+        pcd.points = o3d.utility.Vector3dVector(pts_all[j][mask])
+        pcd.colors = o3d.utility.Vector3dVector(images_all[j][mask])
+        pcd.normals = o3d.utility.Vector3dVector(pts_normal_all[j][mask])
+        pcds.append(pcd)
+
+    pcd_combined = combine_and_clean_point_clouds(pcds, voxel_size=args.voxel_size * 0.1)
+    mesh_recon = point2mesh(pcd_combined)
 
 
 if __name__ == '__main__':