add DPVO

openxrlab · Apr 23, 2024 · 2c94bb3 · 2c94bb3
1 parent 2d066a1
commit 2c94bb3
Show file tree

Hide file tree

Showing 61 changed files with 8,203 additions and 277 deletions.
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -1,4 +1,4 @@
-exclude: '^slam/model_components/slam_helpers_splatam\.py$|^scripts/download_data\.py$|^test/data/'
+exclude: '^slam/model_components/slam_helpers_splatam\.py$|^scripts/download_data\.py$|^test/data/|^third_party'
 
 repos:
   - repo: https://github.com/pycqa/flake8.git

diff --git a/README.md b/README.md
@@ -19,10 +19,10 @@ OpenXRLab Deep-learning based SLAM Toolbox and Benchmark. It is a part of the Op
 
 We provide a set of pre-implemented deep-learning based SLAM algorithms.
 
-| Replica/office0                                              |                                                              |                                                              |                                                              |                                                              |
-| :----------------------------------------------------------- | ------------------------------------------------------------ | ------------------------------------------------------------ | ------------------------------------------------------------ | ------------------------------------------------------------ |
-| [nice-slam](https://github.com/cvg/nice-slam)                | [co-slam](https://github.com/HengyiWang/Co-SLAM)             | [Vox-Fusion](https://github.com/zju3dv/Vox-Fusion)           | [Point_SLAM](https://github.com/eriksandstroem/Point-SLAM)   | [splaTAM](https://github.com/spla-tam/SplaTAM)               |
-| <img src="./docs/imgs/nice_slam.gif" alt="nice-slam" style="zoom: 50%;" /> | <img src="./docs/imgs/co_slam.gif" alt="nice-slam" style="zoom: 50%;" /> | <img src="./docs/imgs/vox_fusion.gif" alt="nice-slam" style="zoom: 50%;" /> | <img src="./docs/imgs/point_slam.gif" alt="nice-slam" style="zoom: 50%;" /> | <img src="./docs/imgs/splatam.gif" alt="nice-slam" style="zoom: 50%;" /> |
+| Replica/office0                                              |                                                              |                                                              |                                                              |                                                              |                                                              |
+| :----------------------------------------------------------- | ------------------------------------------------------------ | ------------------------------------------------------------ | ------------------------------------------------------------ | ------------------------------------------------------------ | ------------------------------------------------------------ |
+| [nice-slam](https://github.com/cvg/nice-slam)                | [co-slam](https://github.com/HengyiWang/Co-SLAM)             | [Vox-Fusion](https://github.com/zju3dv/Vox-Fusion)           | [Point_SLAM](https://github.com/eriksandstroem/Point-SLAM)   | [splaTAM](https://github.com/spla-tam/SplaTAM)               | [DPVO](https://github.com/princeton-vl/DPVO)                 |
+| <img src="./docs/imgs/nice_slam.gif" alt="nice-slam" style="zoom: 50%;" /> | <img src="./docs/imgs/co_slam.gif" alt="nice-slam" style="zoom: 50%;" /> | <img src="./docs/imgs/vox_fusion.gif" alt="nice-slam" style="zoom: 50%;" /> | <img src="./docs/imgs/point_slam.gif" alt="nice-slam" style="zoom: 50%;" /> | <img src="./docs/imgs/splatam.gif" alt="nice-slam" style="zoom: 50%;" /> | <img src="./docs/imgs/dpvo.gif" alt="dpvo" style="zoom: 50%;" /> |
 
 ## Quickstart
 
@@ -49,7 +49,6 @@ conda activate xrdslam
 **Build extension**
 
 ```
-# for vox-fusion 3rd-party
 cd third_party
 bash install.sh
 ```
@@ -188,6 +187,7 @@ Note: The default configuration in the algorithm is suitable for Replica. If you
 | Point-SLAM_X | 0.47                | 34.10 | 0.97  | 0.10   | 99.30           | 83.78        | 90.86   | 0.38           | 1.25           | 3.12           | 88.15                  |
 | SplaTAM      | 0.36                | 34.11 | 0.97  | 0.10   | -               | -            | -       | -              | -              | -              | -                      |
 | SplaTAM_X    | 0.40                | 34.44 | 0.96  | 0.09   | -               | -            | -       | -              | -              | -              | -                      |
+| DPVO_X | 0.36 | - | - | - | - | - | - | - | - | - | - |
 
 ## License
 
@@ -196,7 +196,7 @@ Please note that this license only applies to the code in our library, the depen
 
 ## Acknowledgement
 
-In addition to the implemented algorithm ([nice-slam](https://github.com/cvg/nice-slam),[co-slam](https://github.com/HengyiWang/Co-SLAM),[Vox-Fusion](https://github.com/zju3dv/Vox-Fusion),[Point_SLAM](https://github.com/eriksandstroem/Point-SLAM),[splaTAM](https://github.com/spla-tam/SplaTAM)), our code also adapt some codes from [nerfStudio](https://github.com/nerfstudio-project/nerfstudio/), [sdfstudio](https://autonomousvision.github.io/sdfstudio/). Thanks for making the code available.
+In addition to the implemented algorithm ([nice-slam](https://github.com/cvg/nice-slam),[co-slam](https://github.com/HengyiWang/Co-SLAM),[Vox-Fusion](https://github.com/zju3dv/Vox-Fusion),[Point_SLAM](https://github.com/eriksandstroem/Point-SLAM),[splaTAM](https://github.com/spla-tam/SplaTAM), [DPVO](https://github.com/princeton-vl/DPVO)), our code also adapt some codes from [nerfStudio](https://github.com/nerfstudio-project/nerfstudio/), [sdfstudio](https://autonomousvision.github.io/sdfstudio/). Thanks for making the code available.
 
 ## Built On
 

diff --git a/docs/imgs/dpvo.gif b/docs/imgs/dpvo.gif
diff --git a/environment.yml b/environment.yml
@@ -189,6 +189,7 @@ dependencies:
       - termcolor==2.4.0
       - threadpoolctl==3.2.0
       - tifffile==2023.12.9
+      - torch-scatter==2.1.2
       - torchmetrics==1.3.0.post0
       - tqdm==4.66.1
       - traitlets==5.14.1
@@ -201,7 +202,6 @@ dependencies:
       - yacs==0.1.8
       - zipp==3.17.0
       - faiss-gpu==1.7.2
-      - git+https://github.com/tfy14esa/evaluate_3d_reconstruction_lib.git
       - git+https://github.com/facebookresearch/pytorch3d.git
       - git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch
       - git+https://github.com/JonathonLuiten/diff-gaussian-rasterization-w-depth.git@cb65e4b86bc3bd8ed42174b72a62e8d3a3a71110
diff --git a/pretrained/dpvo/dpvo.pth b/pretrained/dpvo/dpvo.pth
diff --git a/requirements.txt b/requirements.txt
@@ -18,6 +18,7 @@ scikit-image
 scikit-learn
 scipy
 seaborn
+torch_scatter
 torchmetrics
 tqdm
 trimesh==3.21.5

diff --git a/scripts/eval.py b/scripts/eval.py
@@ -25,6 +25,7 @@ class EvalMatrics:
 
     eval_traj: bool = True
     eval_recon: bool = True
+    correct_scale: bool = False
 
     def main(self) -> None:
         """Main function."""
@@ -45,7 +46,8 @@ def main(self) -> None:
             poses_est = poses_est[mask]
             traj_result = evaluate(poses_gt,
                                    poses_est,
-                                   plot=f'{output}/eval_ate_plot.png')
+                                   plot=f'{output}/eval_ate_plot.png',
+                                   correct_scale=self.correct_scale)
             print(traj_result)
         if self.eval_recon:
             if self.gt_mesh is None or not os.path.exists(self.gt_mesh):

diff --git a/scripts/utils/eval_ate.py b/scripts/utils/eval_ate.py
@@ -61,7 +61,7 @@ def associate(first_list, second_list, offset=0.0, max_difference=0.02):
     return matches
 
 
-def align(model, data):
+def align(model, data, correct_scale=False):
     """Align two trajectories using the method of Horn (closed-form).
 
     Input:
@@ -74,27 +74,47 @@ def align(model, data):
     trans_error -- translational error per point (1xn)
     """
     numpy.set_printoptions(precision=3, suppress=True)
+
+    # Compute rotation matrix
     model_zerocentered = model - model.mean(1)
     data_zerocentered = data - data.mean(1)
 
     W = numpy.zeros((3, 3))
     for column in range(model.shape[1]):
         W += numpy.outer(model_zerocentered[:, column],
                          data_zerocentered[:, column])
-    U, d, Vh = numpy.linalg.linalg.svd(W.transpose())
+    U, d, Vh = numpy.linalg.svd(W.transpose())
     S = numpy.matrix(numpy.identity(3))
-    if (numpy.linalg.det(U) * numpy.linalg.det(Vh) < 0):
+    if numpy.linalg.det(U) * numpy.linalg.det(Vh) < 0:
         S[2, 2] = -1
     rot = U * S * Vh
-    trans = data.mean(1) - rot * model.mean(1)
 
-    model_aligned = rot * model + trans
+    if correct_scale:
+        # calculate scale
+        rotmodel = rot * model_zerocentered
+        dots = 0.0
+        norms = 0.0
+        for column in range(data_zerocentered.shape[1]):
+            dots += numpy.dot(data_zerocentered[:, column].transpose(),
+                              rotmodel[:, column])
+            norms_i = numpy.linalg.norm(model_zerocentered[:, column])
+            norms += norms_i * norms_i
+        scale = float(dots / norms)
+    else:
+        scale = 1.0
+
+    # Compute translation vector
+    trans = data.mean(1) - scale * rot * model.mean(1)
+
+    # Align model with data
+    model_aligned = scale * rot * model + trans
     alignment_error = model_aligned - data
 
+    # Compute translational error
     trans_error = numpy.sqrt(
         numpy.sum(numpy.multiply(alignment_error, alignment_error), 0)).A[0]
 
-    return rot, trans, trans_error
+    return rot, trans, trans_error, scale
 
 
 def plot_traj(ax, stamps, traj, style, color, label):
@@ -127,7 +147,11 @@ def plot_traj(ax, stamps, traj, style, color, label):
         ax.plot(x, y, style, color=color, label=label)
 
 
-def evaluate_ate(first_list, second_list, plot='', _args=''):
+def evaluate_ate(first_list,
+                 second_list,
+                 plot='',
+                 correct_scale=False,
+                 _args=''):
     # parse command line
     parser = argparse.ArgumentParser(
         description='This script computes the absolute trajectory error \
@@ -184,9 +208,11 @@ def evaluate_ate(first_list, second_list, plot='', _args=''):
         [[float(value) * float(args.scale) for value in second_list[b][0:3]]
          for a, b in matches]).transpose()
 
-    rot, trans, trans_error = align(second_xyz, first_xyz)
+    rot, trans, trans_error, scale = align(second_xyz,
+                                           first_xyz,
+                                           correct_scale=correct_scale)
 
-    second_xyz_aligned = rot * second_xyz + trans
+    second_xyz_aligned = rot * second_xyz * scale + trans
 
     first_stamps = list(first_list.keys())
     first_stamps.sort()
@@ -199,7 +225,7 @@ def evaluate_ate(first_list, second_list, plot='', _args=''):
     second_xyz_full = numpy.matrix(
         [[float(value) * float(args.scale) for value in second_list[b][0:3]]
          for b in second_stamps]).transpose()
-    second_xyz_full_aligned = rot * second_xyz_full + trans
+    second_xyz_full_aligned = rot * second_xyz_full * scale + trans
 
     if args.verbose:
         print('compared_pose_pairs %d pairs' % (len(trans_error)))
@@ -274,10 +300,12 @@ def evaluate_ate(first_list, second_list, plot='', _args=''):
         rot,
         'trans':
         trans,
+        'scale':
+        scale,
     }
 
 
-def evaluate(poses_gt, poses_est, plot):
+def evaluate(poses_gt, poses_est, plot, correct_scale=False):
 
     poses_gt = poses_gt.cpu().numpy()
     poses_est = poses_est.cpu().numpy()
@@ -286,7 +314,7 @@ def evaluate(poses_gt, poses_est, plot):
     poses_gt = dict([(i, poses_gt[i]) for i in range(N)])
     poses_est = dict([(i, poses_est[i]) for i in range(N)])
 
-    results = evaluate_ate(poses_gt, poses_est, plot)
+    results = evaluate_ate(poses_gt, poses_est, plot, correct_scale)
     return results
 
 

diff --git a/scripts/utils/viz_utils.py b/scripts/utils/viz_utils.py
@@ -164,7 +164,7 @@ def animation_callback(vis):
 
     # set the viewer's pose in the back of the first frame's pose
     param = ctr.convert_to_pinhole_camera_parameters()
-    if algorithm_name != 'splaTAM':
+    if algorithm_name != 'splaTAM' and algorithm_name != 'dpvo':
         init_pose[:3, 3] += 6 * normalize(init_pose[:3, 2])
         init_pose[:3, 2] *= -1
         init_pose[:3, 1] *= -1
@@ -202,7 +202,7 @@ def update_pose(self, index, pose, gt=False):
             pose = pose.cpu().numpy()
 
         # Note: splaTAM should not use pose[:3, 2] *= -1
-        if self.algorithm_name != 'splaTAM':
+        if self.algorithm_name != 'splaTAM' and self.algorithm_name != 'dpvo':
             pose[:3, 2] *= -1
         self.queue.put_nowait(('pose', index, pose, gt))
 

diff --git a/scripts/viewer.py b/scripts/viewer.py
@@ -75,14 +75,14 @@ def main(self) -> None:
                     frontend.update_cloud(cloudfile)
             frontend.update_pose(1, estimate_c2w_list[i], gt=False)
             # Note: not show gt_traj for splaTAM
-            if algorithm_name != 'splaTAM':
+            if algorithm_name != 'splaTAM' and algorithm_name != 'dpvo':
                 frontend.update_pose(1, gt_c2w_list[i], gt=True)
             # the visualizer might get stuck if update every frame
             # with a long sequence (10000+ frames)
             if i % 10 == 0:
                 frontend.update_cam_trajectory(i, gt=False)
                 # Note: not show gt_traj for splaTAM
-                if algorithm_name != 'splaTAM':
+                if algorithm_name != 'splaTAM' and algorithm_name != 'dpvo':
                     frontend.update_cam_trajectory(i, gt=True)
 
         if self.config.save_rendering:

diff --git a/slam/algorithms/base_algorithm.py b/slam/algorithms/base_algorithm.py
@@ -209,10 +209,11 @@ def setup_optimizers(self,
                 })
 
     def do_tracking(self, cur_frame):
-        optimize_frames = [cur_frame]
-        return self.optimize_update(self.config.tracking_n_iters,
-                                    optimize_frames,
-                                    is_mapping=False)
+        if self.is_initialized():
+            optimize_frames = [cur_frame]
+            return self.optimize_update(self.config.tracking_n_iters,
+                                        optimize_frames,
+                                        is_mapping=False)
 
     def do_mapping(self, cur_frame):
         if not self.is_initialized():