fix: device selection in pytorch

README.md

@@ -5,6 +5,8 @@
 3. [x] test gicp-small_gicp with different voxel_ds and knn for estimation covs and normals
    - knn=10 is the best as for voxel greater than 0.01
    - voxel less than 0.001 leads to extremely slow for processing
-4. [ ] test voxel downsample with different ratio  with open3d small_gicp ,different icps
+4. [x] test voxel downsample with different ratio  with open3d small_gicp ,different icps
+   - no ds for cmp with my method,and set knn=20,max_co=0.1
 5. [x] build depth_loss for model
 6. [ ] build silhoutte_loss and color loss(or lab loss) for model
+7. [ ] test for icps as 4 init params

src/component/tracker.py

@@ -47,8 +47,9 @@ class Scan2ScanICP:
 
     def __init__(
         self,
-        voxel_downsampling_resolutions=0.1,  # Adjusted for potentially denser point clouds from depth images
         max_corresponding_distance=0.1,
+        voxel_downsampling_resolutions: float = 0.0,
+        knn: int = 20,
         num_threads=32,
         registration_type: Literal["ICP", "PLANE_ICP", "GICP", "COLORED_ICP"] = "GICP",
         implementation: Literal["small_gicp", "open3d"] = "small_gicp",
@@ -78,6 +79,8 @@ def __init__(
         self.kf_tree: small_gicp.KdTree | None = None
         self.kf_T_last_current = np.identity(4)
 
+        self.knn = knn
+
     def align_pcd(
         self,
         raw_points: NDArray[np.float64],
@@ -172,17 +175,18 @@ def align(
         raw_points: NDArray[np.float64],
         init_gt_pose: NDArray[np.float64] | None = None,
         T_last_current: NDArray[np.float64] = np.identity(4),
-        knn: int = 10,
     ):
         """
         Parameters
         ----------
         raw_points: shape = (h*w*(3/4/6/7)
         """
         if self.backend == "small_gicp" and self.registration_type != "COLORED_ICP":
-            return self.align_small_gicp(raw_points, init_gt_pose, T_last_current, knn)
+            return self.align_small_gicp(
+                raw_points, init_gt_pose, T_last_current, self.knn
+            )
         elif self.backend == "open3d":
-            return self.align_o3d(raw_points, init_gt_pose, T_last_current)
+            return self.align_o3d(raw_points, init_gt_pose, T_last_current, self.knn)
         else:
             raise ValueError("wrong backend type")
 
@@ -191,16 +195,25 @@ def align_small_gicp(
         raw_points: NDArray[np.float64],
         init_gt_pose: NDArray[np.float64] | None = None,
         T_last_current: NDArray[np.float64] = np.identity(4),
-        knn: int = 10,
+        knn: int = 20,
     ):
         # down sample the point cloud
-        downsampled, tree = small_gicp.preprocess_points(
-            raw_points,
-            self.voxel_downsampling_resolutions,
-            num_threads=self.num_threads,
-            num_neighbors=knn,
-        )
-
+        if self.voxel_downsampling_resolutions > 0.0:
+            downsampled, tree = small_gicp.preprocess_points(
+                raw_points,
+                self.voxel_downsampling_resolutions,
+                num_threads=self.num_threads,
+                num_neighbors=knn,
+            )
+        elif self.voxel_downsampling_resolutions == 0.0:
+            raw_points = small_gicp.PointCloud(raw_points)
+            tree = small_gicp.KdTree(raw_points, num_threads=self.num_threads)
+            small_gicp.estimate_normals_covariances(
+                raw_points, tree, num_neighbors=knn, num_threads=self.num_threads
+            )
+            downsampled = raw_points
+        else:
+            raise ValueError("voxel_downsampling_resolutions must greater than 0.0")
         # first frame
         if self.previous_pcd is None:
             self.previous_pcd = downsampled
@@ -236,6 +249,7 @@ def align_o3d(
         raw_points: NDArray[np.float64],
         init_gt_pose: NDArray[np.float64] | None = None,
         T_last_current: NDArray[np.float64] = np.identity(4),
+        knn: int = 20,
     ):
 
         # 创建 Open3D 点云对象
@@ -244,10 +258,15 @@ def align_o3d(
         if self.registration_type == "COLORED_ICP":
             pcd.colors = o3d.utility.Vector3dVector(raw_points[:, 4:] / 255.0)
             # Compute normals for the point cloud, which are needed for GICP and Colored ICP
-        pcd.estimate_normals(search_param=o3d.geometry.KDTreeSearchParamKNN(knn=10))
+        pcd.estimate_normals(search_param=o3d.geometry.KDTreeSearchParamKNN(knn=knn))
+
         # 体素下采样
-        voxel_size = self.voxel_downsampling_resolutions
-        downsampled = pcd.voxel_down_sample(voxel_size)
+        if self.voxel_downsampling_resolutions > 0.0:
+            downsampled = pcd.voxel_down_sample(self.voxel_downsampling_resolutions)
+        elif self.voxel_downsampling_resolutions == 0.0:
+            downsampled = pcd
+        else:
+            raise ValueError("voxel_downsampling_resolutions must greater than 0.0")
 
         # 如果是第一帧，初始化
         if self.previous_pcd is None:

src/gicp_tensor.py → src/depth_loss_test.py

@@ -20,7 +20,7 @@ def save_finished_experiment(file_path: Path, finished: tuple):
 
 # TODO: downsample with different backends
 if __name__ == "__main__":
-    file_path = Path("grip_o3d_finished_experiments.json")
+    file_path = Path("depth_loss_finished_experiments.json")
     methods = ["depth_loss"]
     implements = "pytorch"
     num_iters = 20
@@ -43,6 +43,7 @@ def save_finished_experiment(file_path: Path, finished: tuple):
                         num_iters=num_iters,
                         learning_rate=learning_rate,
                         description="depth_loss for pose estimation",
+                        implementation="pytorch",
                     ),
                 )
                 experiment.run()

src/eval/experiment.py

@@ -11,7 +11,7 @@
     calculate_translation_error,
     diff_pcd_COM,
 )
-from src.gicp.depth_loss import train_model
+from src.gicp.depth_loss import train_model, DEVICE
 from src.slam_data import Replica, RGBDImage
 from src.slam_data.dataset import DataLoaderBase
 from src.utils import to_tensor
@@ -21,7 +21,6 @@ class RegistrationConfig(NamedTuple):
     max_corresponding_distance: float = 0.1
     num_threads: int = 32
     registration_type: Literal["ICP", "PLANE_ICP", "GICP", "COLORED_ICP"] = ("GICP",)
-    implementation: Literal["small_gicp", "open3d"] = "small_gicp"
     voxel_downsampling_resolutions: float | None = None
     # grid_downsample_resolution: int | None = None
     # for gicp estimate normals and covs 10 is the best after tests
@@ -36,6 +35,7 @@ class WandbConfig(NamedTuple):
     dataset: str = "Replica"
     sub_set: str = "office0"
     description: str = "GICP on Replica dataset"
+    implementation: str | None = None
     num_iters: int | None = None
     learning_rate: float | None = None
 
@@ -95,7 +95,7 @@ def run(self, max_images: int = 2000):
 
             # NOTE: align interface
             if i == 0:
-                res = self.backends.align(new_pcd, rgbd_image.pose, knn=self.knn)
+                # res = self.backends.align(new_pcd, rgbd_image.pose)
                 res = self.backends.align(new_pcd, rgbd_image.pose)
                 continue
             else:
@@ -126,13 +126,9 @@ def run(self, max_images: int = 2000):
 
 class DepthLossExperiment(ExperimentBase):
 
-    def __init__(
-        self,
-        wandb_config: WandbConfig,
-    ):
+    def __init__(self, wandb_config: WandbConfig):
         super().__init__(
-            backends=train_model,
-            wandb_config=wandb_config,
+            backends=train_model, wandb_config=wandb_config, extra_config=kwargs
         )
         self.num_iters = wandb_config.num_iters
         self.learning_rate = wandb_config.learning_rate
@@ -158,7 +154,7 @@ def run(self, max_images: int = 2000):
                 min_loss, pose = self.backends(
                     self.data[i - 1],
                     rgbd_image,
-                    to_tensor(self.data.K, device="cuda"),
+                    to_tensor(self.data.K, device=DEVICE),
                     num_iterations=self.num_iters,
                     learning_rate=self.learning_rate,
                 )

src/gicp/depth_loss.py

@@ -7,12 +7,12 @@
 from src.slam_data import Replica, RGBDImage
 from src.utils import to_tensor
 
-device = (
+DEVICE = (
     "cuda"
     if torch.cuda.is_available()
     else "mps" if torch.backends.mps.is_available() else "cpu"
 )
-print(f"Using {device} device")
+print(f"Using {DEVICE} DEVICE")
 
 
 class PoseEstimationModel(nn.Module):
@@ -145,17 +145,17 @@ def train_model(
     num_iterations=20,
     learning_rate=1e-6,
 ) -> tuple[float, Tensor]:
-    init_pose = to_tensor(tar_rgb_d.pose, device)
-    model = PoseEstimationModel(K, init_pose, device)
+    init_pose = to_tensor(tar_rgb_d.pose, DEVICE)
+    model = PoseEstimationModel(K, init_pose, DEVICE)
     optimizer = optim.Adam(model.parameters(), lr=learning_rate)
     # scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.5)
     # result
     min_loss = float("inf")
     best_pose = init_pose.clone()
     for i in range(num_iterations):
         optimizer.zero_grad()
-        depth_last = to_tensor(tar_rgb_d.depth, device)
-        depth_current = to_tensor(src_rgb_d.depth, device)
+        depth_last = to_tensor(tar_rgb_d.depth, DEVICE)
+        depth_current = to_tensor(src_rgb_d.depth, DEVICE)
         loss = model(depth_last, depth_current)
         loss.backward()
         optimizer.step()
@@ -172,7 +172,7 @@ def train_model(
 
 def eval():
     tar_rgb_d, src_rgb_d = Replica()[0], Replica()[1]
-    _, estimate_pose = train_model(tar_rgb_d, src_rgb_d, to_tensor(tar_rgb_d.K, device))
+    _, estimate_pose = train_model(tar_rgb_d, src_rgb_d, to_tensor(tar_rgb_d.K, DEVICE))
 
     eT = calculate_translation_error(
         estimate_pose.detach().cpu().numpy(), tar_rgb_d.pose