update traj and vis save

modules/droid.py

@@ -35,6 +35,7 @@ class Options:
     backend_nms: int = 3
     upsample: bool = False
     reconstruction_path: Path = None
+    vis_save: Path = None
     # new options
     intrinsic: np.ndarray = None
     focal: float = None

modules/droid_core/droid.py

@@ -35,7 +35,7 @@ def __init__(self, args):
         # visualizer
         if not self.disable_vis:
             from .visualization import droid_visualization
-            self.visualizer = Process(target=droid_visualization, args=(self.video,))
+            self.visualizer = Process(target=droid_visualization, args=(self.video, 'cuda:0', self.args.vis_save))
             self.visualizer.start()
 
         # post processor - fill in poses for non-keyframes
@@ -72,7 +72,7 @@ def track(self, tstamp, image, depth=None, intrinsics=None):
             # self.backend()
 
     def terminate(self, stream=None):
-        """ terminate the visualization process, return poses [t, q] """
+        """ terminate the visualization process, return timestamp and poses [t, q] """
 
         del self.frontend
 
@@ -85,5 +85,11 @@ def terminate(self, stream=None):
         self.backend(12)
 
         camera_trajectory = self.traj_filler(stream)
-        return camera_trajectory.inv().data.cpu().numpy()
+        camera_trajectory = camera_trajectory.inv().data.cpu().numpy()
+
+        # fill timestamp
+        timestamps = np.arange(len(camera_trajectory)).reshape(-1, 1)
+        traj_tum = np.concatenate([timestamps, camera_trajectory], axis=1)
+
+        return traj_tum
 

modules/droid_core/visualization.py

@@ -51,7 +51,7 @@ def create_point_actor(points, colors):
     point_cloud.colors = o3d.utility.Vector3dVector(colors)
     return point_cloud
 
-def droid_visualization(video, device="cuda:0"):
+def droid_visualization(video, device="cuda:0", vis_save=None):
     """ DROID visualization frontend """
 
     torch.cuda.set_device(device)
@@ -66,6 +66,7 @@ def droid_visualization(video, device="cuda:0"):
 
     def increase_filter(vis):
         droid_visualization.filter_thresh *= 2
+        print(f"filter threshold: {droid_visualization.filter_thresh}")
         with droid_visualization.video.get_lock():
             droid_visualization.video.dirty[:droid_visualization.video.counter.value] = True
 
@@ -153,4 +154,19 @@ def animation_callback(vis):
     vis.get_render_option().load_from_json(str(Path(__file__).parent / "misc/renderoption.json"))
 
     vis.run()
+
+    # save traj ply for visualization
+    if vis_save is None:
+        vis_save = Path(vis_save)
+        pcd = o3d.geometry.PointCloud()
+        for actor in droid_visualization.points.values():
+            pcd += actor
+        o3d.io.write_point_cloud(str(vis_save / 'scene.plt'), pcd)
+
+        # save traj
+        lineset = o3d.geometry.LineSet()
+        for cam in droid_visualization.cameras.values():
+            lineset += cam
+        o3d.io.write_line_set(str(vis_save / 'traj.ply'), lineset)
+
     vis.destroy_window()

modules/metric.py

@@ -26,8 +26,8 @@ class Metric3D:
 
     def __init__(
         self,
-        checkpoint: Union[str, Path] = './weights/metric_depth_vit_large_800k.pth',
-        model_name: str = 'v2-L',
+        checkpoint: Union[str, Path] = './weights/metric_depth_vit_giant2_800k.pth',
+        model_name: str = 'v2-g',
     ) -> None:
         checkpoint = Path(checkpoint).resolve()
         cfg:Config = self._load_config_(model_name, checkpoint)

modules/utils.py

@@ -135,7 +135,7 @@ def trajectory_to_poses(
 
     for i in range(len(traj)):
         pose = traj[i]
-        t, q = pose[:3], pose[3:]
+        t, q = pose[1:4], pose[4:]
         R = quaternion_to_matrix(q)
         T = np.eye(4)
         # Twc = [R | t]

run.py

@@ -14,6 +14,7 @@
     parser.add_argument("--traj", type=str, help="trajectory file", default="./trajectory.txt")
     parser.add_argument("--poses", type=str, help="result directory", default=None)
     parser.add_argument("--viz", action='store_true', help="visualize", default=False)
+    parser.add_argument("--viz-save", type=str, help="save visualization", default=None)
     parser.add_argument("--focal", type=float, default=None, help="focal length")
     parser.add_argument("--calib", type=str, default=None, help="calib file, overwrite focal")
     parser.add_argument("--weight", type=str, default='./weights/droid.pth', help="checkpoint file")
@@ -31,6 +32,7 @@
     opt.weights = Path(args.weight)         # checkpoint file
     opt.disable_vis = not args.viz          # visualization
     opt.focal = args.focal                  # focal length
+    opt.vis_save = args.viz_save            # save visualization
 
     # global ba on frontend, 0 (set to off) by default
     opt.global_ba_frontend = args.global_ba_frontend   # frequency to run global ba on frontend

scripts/viz_scene.py

@@ -0,0 +1,17 @@
+import argparse
+import open3d as o3d
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--scene', type=str, help='scene file')
+    parser.add_argument('--traj', type=str, help='trajectory file')
+    args = parser.parse_args()
+
+    pcd = o3d.io.read_point_cloud(args.scene)
+    lineset = o3d.io.read_line_set(args.traj)
+    # paint lineset
+    lineset.paint_uniform_color([1, 0, 0])
+
+    # visualize
+    o3d.visualization.draw_geometries([pcd, lineset])