allow dataparser poses to be in float64 for higher precision

Author: atonderski

nerfstudio/data/dataparsers/ad_dataparser.py

@@ -310,10 +310,11 @@ def _adjust_times(
     def _adjust_poses(self, cameras: Cameras, lidars: Lidars, trajectories: List[Dict]):
         """Determines a new, centered, world coordinate system, and adjusts all poses."""
         w2m = _get_world_to_mean_transform(cameras, lidars)
-        cameras.camera_to_worlds = pose_multiply(w2m, cameras.camera_to_worlds)
-        lidars.lidar_to_worlds = pose_multiply(w2m, lidars.lidar_to_worlds)
+        # Cast poses to float32 only after transforming to local frame to avoid precision loss
+        cameras.camera_to_worlds = pose_multiply(w2m, cameras.camera_to_worlds).to(torch.float32)
+        lidars.lidar_to_worlds = pose_multiply(w2m, lidars.lidar_to_worlds).to(torch.float32)
         for traj in trajectories:
-            traj["poses"][:, :3] = pose_multiply(w2m, traj["poses"][:, :3])
+            traj["poses"][:, :3] = pose_multiply(w2m, traj["poses"][:, :3]).to(torch.float32)
         return w2m
 
     def _get_train_eval_indices(self, sensors: Union[Cameras, Lidars]) -> Tuple[Tensor, Tensor]:
@@ -617,8 +618,8 @@ def _get_world_to_mean_transform(cameras: Cameras, lidars: Lidars):
         m2w = to4x4(select_poses[0:1])[0]
     else:
         # Otherwise
-        m2w = torch.from_numpy(_get_mean_pose_from_trajectory(select_trajectory).astype(np.float32))
-    return torch.linalg.inv(m2w)[:3]
+        m2w = torch.from_numpy(_get_mean_pose_from_trajectory(select_trajectory))
+    return torch.linalg.inv(m2w)[:3].to(poses.dtype)
 
 
 def _empty_cameras():

nerfstudio/data/dataparsers/argoverse2_dataparser.py

@@ -186,7 +186,7 @@ class Argoverse2(ADDataParser):
     @property
     def actor_transform(self) -> torch.Tensor:
         """Argo uses x-forward, so we need to rotate to x-right."""
-        wlh_to_lwh = np.eye(4, dtype=np.float32)
+        wlh_to_lwh = np.eye(4)
         wlh_to_lwh[:3, :3] = WLH_TO_LWH
         return torch.from_numpy(wlh_to_lwh)[:3, :]
 
@@ -360,11 +360,11 @@ def _read_lidars(self, lidars: Lidars, filepaths: List[Path]) -> List[torch.Tens
             assert sweep is not None
             uplidar2ego = sweep.ego_SE3_up_lidar
             all_lup2w = torch.tensor(
-                np.array([e2w.compose(uplidar2ego).transform_matrix for e2w in all_ego2w]), dtype=torch.float32
+                np.array([e2w.compose(uplidar2ego).transform_matrix for e2w in all_ego2w]), dtype=torch.float64
             )
             downlidar2ego = sweep.ego_SE3_down_lidar
             all_ldown2w = torch.tensor(
-                np.array([e2w.compose(downlidar2ego).transform_matrix for e2w in all_ego2w]), dtype=torch.float32
+                np.array([e2w.compose(downlidar2ego).transform_matrix for e2w in all_ego2w]), dtype=torch.float64
             )
             all_times = torch.from_numpy(log_pose_df["timestamp_ns"].to_numpy() / 1e9)
 

nerfstudio/data/dataparsers/base_dataparser.py

@@ -68,7 +68,7 @@ class DataparserOutputs:
     """
     dataparser_transform: Float[Tensor, "3 4"] = torch.eye(4)[:3, :]
     """Transform applied by the dataparser to the entire scene."""
-    actor_transform: Float[Tensor, "3 4"] = torch.eye(4, dtype=torch.float32)[:3, :]
+    actor_transform: Float[Tensor, "3 4"] = torch.eye(4)[:3, :]
     """Transform applied by the dataparser to each actor's local frame."""
     dataparser_scale: float = 1.0
     """Scale applied by the dataparser."""

nerfstudio/data/dataparsers/kittimot_dataparser.py

@@ -107,7 +107,7 @@ class KittiMot(ADDataParser):
     @property
     def actor_transform(self) -> Tensor:
         """The transform needed to convert the actor poses to our desired format (x-right, y-forward, z-up)."""
-        return torch.from_numpy(RIGHT_FRONT_UP2RIGHT_DOWN_FRONT)[:3, :]
+        return torch.from_numpy(RIGHT_FRONT_UP2RIGHT_DOWN_FRONT)
 
     def _get_cameras(self) -> Tuple[Cameras, List[Path]]:
         """Returns camera info and image filenames."""

nerfstudio/data/dataparsers/nuscenes_dataparser.py

@@ -185,7 +185,7 @@ class NuScenes(ADDataParser):
     @property
     def actor_transform(self) -> torch.Tensor:
         """Nuscenes uses x-forward, so we need to rotate to x-right."""
-        return torch.from_numpy(WLH_TO_LWH)[:3, :]
+        return torch.from_numpy(WLH_TO_LWH)
 
     def _get_cameras(self) -> Tuple[Cameras, List[Path]]:
         if "all" in self.config.cameras:

nerfstudio/data/dataparsers/zod_dataparser.py

@@ -191,7 +191,7 @@ class Zod(ADDataParser):
     @property
     def actor_transform(self) -> torch.Tensor:
         """ZOD uses x-forward, so we need to rotate to x-right."""
-        return torch.from_numpy(WLH_TO_LWH)[:3, :]
+        return torch.from_numpy(WLH_TO_LWH)
 
     def _get_lane_shift_sign(self, sequence: str) -> Literal[-1, 1]:
         return LANE_SHIFT_SIGN.get(sequence, 1)

nerfstudio/utils/poses.py

@@ -174,7 +174,7 @@ def interpolate_trajectories(poses, pose_times, query_times, pose_valid_mask=Non
     right_time = pose_times[right_idx]
     left_time = pose_times[left_idx]
     time_diff = right_time - left_time + 1e-6
-    fraction = (qt - left_time) / time_diff  # 0 = all left, 1 = all right
+    fraction = ((qt - left_time) / time_diff).to(poses.dtype)  # 0 = all left, 1 = all right
     if clamp_frac:
         fraction = fraction.clamp(0.0, 1.0)  # clamp to handle out of bounds