From 4bcfd1ee110ab6b27f03e7143a4ae8c4634fe774 Mon Sep 17 00:00:00 2001 From: atticuszz <1831768457@qq.com> Date: Fri, 7 Jun 2024 21:35:54 +0800 Subject: [PATCH] fix: device selection in pytorch --- README.md | 4 +- src/component/tracker.py | 49 +++++++--- src/{gicp_tensor.py => depth_loss_test.py} | 3 +- src/eval/experiment.py | 16 ++-- src/gicp/depth_loss.py | 14 +-- src/gicp_downsample.py | 103 --------------------- src/icps_test.py | 64 +++++++++++++ src/slam.py | 67 -------------- 8 files changed, 116 insertions(+), 204 deletions(-) rename src/{gicp_tensor.py => depth_loss_test.py} (93%) delete mode 100644 src/gicp_downsample.py create mode 100644 src/icps_test.py delete mode 100644 src/slam.py diff --git a/README.md b/README.md index 31c90c6..083e00a 100644 --- a/README.md +++ b/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 diff --git a/src/component/tracker.py b/src/component/tracker.py index b6ca7d9..a4e34db 100644 --- a/src/component/tracker.py +++ b/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,7 +175,6 @@ 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 @@ -180,9 +182,11 @@ def align( 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: diff --git a/src/gicp_tensor.py b/src/depth_loss_test.py similarity index 93% rename from src/gicp_tensor.py rename to src/depth_loss_test.py index 5534150..df7a7aa 100644 --- a/src/gicp_tensor.py +++ b/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() diff --git a/src/eval/experiment.py b/src/eval/experiment.py index df68302..af45e91 100644 --- a/src/eval/experiment.py +++ b/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, ) diff --git a/src/gicp/depth_loss.py b/src/gicp/depth_loss.py index e8fa547..0d12493 100644 --- a/src/gicp/depth_loss.py +++ b/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,8 +145,8 @@ 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 @@ -154,8 +154,8 @@ def train_model( 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 diff --git a/src/gicp_downsample.py b/src/gicp_downsample.py deleted file mode 100644 index e6909d5..0000000 --- a/src/gicp_downsample.py +++ /dev/null @@ -1,103 +0,0 @@ -import json -from pathlib import Path - -from src.eval.experiment import Experiment, RegistrationConfig, WandbConfig - - -def load_finished_experiments(file_path: Path): - """Load the list of finished experiments from the file.""" - if not file_path.exists(): - return [] - with file_path.open("r") as file: - return [tuple(json.loads(line)) for line in file] - - -def save_finished_experiment(file_path: Path, finished: tuple): - """Append a new finished experiment to the file.""" - with file_path.open("a") as file: - file.write(json.dumps(finished) + "\n") - - -def points_left(v_ds, at_least: int = 1000): - - left = v_ds * 1200 * 680 - if left < at_least: - print("v_ds:", v_ds, "points:", left) - return left < at_least - - -# TODO: downsample with different backends -if __name__ == "__main__": - file_path = Path("grip_o3d_finished_experiments.json") - methods = [ - "GICP", - "PLANE_ICP", - "COLORED_ICP", - "ICP", - ] - implements = "open3d", "small_gicp" - # implements = ( - # "small_gicp", - # "open3d", - # ) - voxel_downsampling_resolutions = [ - 1, - 0.9, - 0.8, - 0.7, - 0.6, - 0.5, - 0.45, - 0.4, - 0.35, - 0.3, - 0.25, - 0.2, - 0.15, - 0.1, - 0.05, - 0.01, - ] - voxel_downsampling_resolutions.reverse() - # filter too less - too_less_points = [] - for v_ds in voxel_downsampling_resolutions: - if points_left(v_ds): - too_less_points.append(v_ds) - rooms = ["room" + str(i) for i in range(3)] + ["office" + str(i) for i in range(5)] - - # finished = [ - # (room, method, v_ds, grid_ds) - # for room in rooms[:1] - # for method in methods[:1] - # for v_ds in voxel_downsampling_resolutions[:3] - # for grid_ds in grid_downsample_resolutions - # ] - - finished = load_finished_experiments(file_path) - for room in rooms: - for method in methods: - for imp in implements: - for v_ds in voxel_downsampling_resolutions: - # skip finished - config_tuple = (room, method, imp, v_ds) - if config_tuple in finished or v_ds in too_less_points: - continue - # small_gicp does not have color icp - if imp == "small_gicp" and method == "COLORED_ICP": - continue - registration_config = RegistrationConfig( - registration_type=method, - voxel_downsampling_resolutions=v_ds, - implementation=imp, - ) - experiment = Experiment( - registration_config=registration_config, - wandb_config=WandbConfig( - method, - sub_set=room, - description="small_gicp and o3d gcip for test voxel dowmsample influence for accuracy", - ), - ) - experiment.run() - save_finished_experiment(file_path, config_tuple) diff --git a/src/icps_test.py b/src/icps_test.py new file mode 100644 index 0000000..1f55b6b --- /dev/null +++ b/src/icps_test.py @@ -0,0 +1,64 @@ +import json +from pathlib import Path + +from src.eval.experiment import ICPExperiment, RegistrationConfig, WandbConfig + + +def load_finished_experiments(file_path: Path): + """Load the list of finished experiments from the file.""" + if not file_path.exists(): + return [] + with file_path.open("r") as file: + return [tuple(json.loads(line)) for line in file] + + +def save_finished_experiment(file_path: Path, finished: tuple): + """Append a new finished experiment to the file.""" + with file_path.open("a") as file: + file.write(json.dumps(finished) + "\n") + + +if __name__ == "__main__": + + file_path = Path("grip_o3d_finished_experiments.json") + methods = [ + "GICP", + "PLANE_ICP", + "COLORED_ICP", + "ICP", + ] + implements = ( + "small_gicp", + "open3d", + ) + + rooms = ["room" + str(i) for i in range(3)] + ["office" + str(i) for i in range(5)] + + finished = load_finished_experiments(file_path) + + for room in rooms: + for method in methods: + for imp in implements: + # skip finished + config_tuple = (room, method, imp) + if config_tuple in finished: + continue + # small_gicp does not have color icp + if imp == "small_gicp" and method == "COLORED_ICP": + continue + registration_config = RegistrationConfig( + registration_type=method, + voxel_downsampling_resolutions=0.0, + # knn=20, + ) + experiment = ICPExperiment( + registration_config=registration_config, + wandb_config=WandbConfig( + method, + sub_set=room, + implementation=imp, + description="icps test for accuracy", + ), + ) + experiment.run() + save_finished_experiment(file_path, config_tuple) diff --git a/src/slam.py b/src/slam.py deleted file mode 100644 index bd8c3e1..0000000 --- a/src/slam.py +++ /dev/null @@ -1,67 +0,0 @@ -from collections import deque - -import cv2 -import numpy as np - -from component import Mapper -from component.tracker import Scan2ScanICP -from slam_data import Camera, Replica, RGBDImage - - -class Slam2D: - - def __init__(self, input_folder, cfg_file: str, use_camera: bool = False) -> None: - if use_camera: - self.slam_data = Camera(input_folder, cfg_file) - else: - self.slam_data = Replica(input_folder, cfg_file) - self.use_camera = use_camera - self.tracker = Scan2ScanICP() - self.mapper = Mapper() - # fps - self.stamps = deque(maxlen=100) - # self.vis = PcdVisualizer(intrinsic_matrix=self.slam_data.K) - - # for vis - self.gt_poses = [] - self.estimated_poses = [] - - def run(self) -> None: - """ - tracking and mapping - """ - - for i, rgb_d in enumerate(self.slam_data): - rgb_d: RGBDImage - self.gt_poses.append(rgb_d.pose) - - start = cv2.getTickCount() - pcd_c = rgb_d.depth_to_pointcloud(8) - if i == 0: - pose = self.tracker.align_pcd(pcd_c, rgb_d.pose) - else: - pose = self.tracker.align_pcd(pcd_c) - - # pose = rgb_d.pose - - kf = self.tracker.keyframe() - if kf: - pcd_w = rgb_d.camera_to_world(pose, pcd_c) - self.mapper.build_map_2d(pcd_w) - self.mapper.show() - # self.vis.update_render(pcd_c, pose) - # self.vis.vis_trajectory( - # gt_poses=self.gt_poses, - # estimated_poses=self.estimated_poses, - # downsampling_resolution=5, - # fps=fps, - # ) - end = cv2.getTickCount() - self.stamps.append((end - start) / cv2.getTickFrequency()) - - self.estimated_poses.append(pose) - fps = 1 / np.mean(self.stamps) - print(f"Average FPS: {fps}") - - self.mapper.show() - # self.vis.close()