diff --git a/dataset/slam_dataset.py b/dataset/slam_dataset.py index 0971c93..4e335b1 100644 --- a/dataset/slam_dataset.py +++ b/dataset/slam_dataset.py @@ -238,6 +238,7 @@ def read_frame_with_loader(self, frame_id, init_pose: bool = True): # TO ADD if "imus" in dict_keys: self.cur_frame_imus = frame_data["imus"] + # TO ADD self.cur_point_cloud_torch = torch.tensor(points, device=self.device, dtype=self.dtype) @@ -662,22 +663,24 @@ def get_poses_np_for_vis(self): return odom_poses, gt_poses, pgo_poses def write_results(self): - odom_poses = self.odom_poses[:self.processed_frame+1] - odom_poses_out = apply_kitti_format_calib(odom_poses, self.calib["Tr"]) - write_kitti_format_poses(os.path.join(self.run_path, "odom_poses"), odom_poses_out) - write_tum_format_poses(os.path.join(self.run_path, "odom_poses"), odom_poses_out, self.poses_ts, 0.1*self.config.step_frame) - write_traj_as_o3d(odom_poses, os.path.join(self.run_path, "odom_poses.ply")) + + if self.config.track_on: + odom_poses = self.odom_poses[:self.processed_frame+1] + odom_poses_out = apply_kitti_format_calib(odom_poses, self.calib["Tr"]) + write_kitti_format_poses(os.path.join(self.run_path, "odom_poses"), odom_poses_out) + write_tum_format_poses(os.path.join(self.run_path, "odom_poses"), odom_poses_out, self.poses_ts, 0.1*self.config.step_frame) + write_traj_as_o3d(odom_poses, os.path.join(self.run_path, "odom_poses.ply")) - if self.config.pgo_on: - pgo_poses = self.pgo_poses[:self.processed_frame+1] - slam_poses_out = apply_kitti_format_calib(pgo_poses, self.calib["Tr"]) - write_kitti_format_poses( - os.path.join(self.run_path, "slam_poses"), slam_poses_out - ) - write_tum_format_poses( - os.path.join(self.run_path, "slam_poses"), slam_poses_out, self.poses_ts, 0.1*self.config.step_frame - ) - write_traj_as_o3d(pgo_poses, os.path.join(self.run_path, "slam_poses.ply")) + if self.config.pgo_on: + pgo_poses = self.pgo_poses[:self.processed_frame+1] + slam_poses_out = apply_kitti_format_calib(pgo_poses, self.calib["Tr"]) + write_kitti_format_poses( + os.path.join(self.run_path, "slam_poses"), slam_poses_out + ) + write_tum_format_poses( + os.path.join(self.run_path, "slam_poses"), slam_poses_out, self.poses_ts, 0.1*self.config.step_frame + ) + write_traj_as_o3d(pgo_poses, os.path.join(self.run_path, "slam_poses.ply")) # timing report time_table = np.array(self.time_table) @@ -703,136 +706,137 @@ def write_results(self): gt_poses = self.gt_poses[:self.processed_frame+1] write_traj_as_o3d(gt_poses, os.path.join(self.run_path, "gt_poses.ply")) - print("Odometry evaluation:") - avg_tra, avg_rot = relative_error(gt_poses, odom_poses) - ate_rot, ate_trans, align_mat = absolute_error( - gt_poses, odom_poses, self.config.eval_traj_align - ) - if avg_tra == 0: # for rgbd dataset (shorter sequence) - print("Absoulte Trajectory Error (cm):", f"{ate_trans*100.0:.3f}") - else: - print("Average Translation Error (%):", f"{avg_tra:.3f}") - print("Average Rotational Error (deg/100m):", f"{avg_rot*100.0:.3f}") - print("Absoulte Trajectory Error (m):", f"{ate_trans:.3f}") - - if self.config.wandb_vis_on: - wandb_log_content = { - "Average Translation Error [%]": avg_tra, - "Average Rotational Error [deg/m]": avg_rot, - "Absoulte Trajectory Error [m]": ate_trans, - "Absoulte Rotational Error [deg]": ate_rot, - "Consuming time per frame [s]": mean_time_without_init_s, - } - wandb.log(wandb_log_content) - - if self.config.pgo_on: - print("SLAM evaluation:") - avg_tra_slam, avg_rot_slam = relative_error( - gt_poses, pgo_poses - ) - ate_rot_slam, ate_trans_slam, align_mat_slam = absolute_error( - gt_poses, pgo_poses, self.config.eval_traj_align + if self.config.track_on: + print("Odometry evaluation:") + avg_tra, avg_rot = relative_error(gt_poses, odom_poses) + ate_rot, ate_trans, align_mat = absolute_error( + gt_poses, odom_poses, self.config.eval_traj_align ) - if avg_tra_slam == 0: # for rgbd dataset (shorter sequence) - print( - "Absoulte Trajectory Error (cm):", - f"{ate_trans_slam*100.0:.3f}", - ) + if avg_tra == 0: # for rgbd dataset (shorter sequence) + print("Absoulte Trajectory Error (cm):", f"{ate_trans*100.0:.3f}") else: - print("Average Translation Error (%):", f"{avg_tra_slam:.3f}") - print( - "Average Rotational Error (deg/100m):", - f"{avg_rot_slam*100.0:.3f}", - ) - print( - "Absoulte Trajectory Error (m):", f"{ate_trans_slam:.3f}" - ) + print("Average Translation Error (%):", f"{avg_tra:.3f}") + print("Average Rotational Error (deg/100m):", f"{avg_rot*100.0:.3f}") + print("Absoulte Trajectory Error (m):", f"{ate_trans:.3f}") if self.config.wandb_vis_on: wandb_log_content = { - "SLAM Average Translation Error [%]": avg_tra_slam, - "SLAM Average Rotational Error [deg/m]": avg_rot_slam, - "SLAM Absoulte Trajectory Error [m]": ate_trans_slam, - "SLAM Absoulte Rotational Error [deg]": ate_rot_slam, + "Average Translation Error [%]": avg_tra, + "Average Rotational Error [deg/m]": avg_rot, + "Absoulte Trajectory Error [m]": ate_trans, + "Absoulte Rotational Error [deg]": ate_rot, + "Consuming time per frame [s]": mean_time_without_init_s, } wandb.log(wandb_log_content) - csv_columns = [ - "Average Translation Error [%]", - "Average Rotational Error [deg/m]", - "Absoulte Trajectory Error [m]", - "Absoulte Rotational Error [deg]", - "Consuming time per frame [s]", - "Frame count", - ] - pose_eval = [ - { - csv_columns[0]: avg_tra, - csv_columns[1]: avg_rot, - csv_columns[2]: ate_trans, - csv_columns[3]: ate_rot, - csv_columns[4]: mean_time_without_init_s, - csv_columns[5]: int(self.processed_frame), - } - ] - if self.config.pgo_on: - slam_eval_dict = { - csv_columns[0]: avg_tra_slam, - csv_columns[1]: avg_rot_slam, - csv_columns[2]: ate_trans_slam, - csv_columns[3]: ate_rot_slam, - csv_columns[4]: mean_time_without_init_s, - csv_columns[5]: int(self.processed_frame), - } - pose_eval.append(slam_eval_dict) - output_csv_path = os.path.join(self.run_path, "pose_eval.csv") - try: - with open(output_csv_path, "w") as csvfile: - writer = csv.DictWriter(csvfile, fieldnames=csv_columns) - writer.writeheader() - for data in pose_eval: - writer.writerow(data) - except IOError: - print("I/O error") - - # if self.config.o3d_vis_on: # x service issue for remote server - output_traj_plot_path_2d = os.path.join(self.run_path, "traj_plot_2d.png") - output_traj_plot_path_3d = os.path.join(self.run_path, "traj_plot_3d.png") - # trajectory not aligned yet in the plot - # require list of numpy arraies as the input - - gt_position_list = [self.gt_poses[i] for i in range(self.processed_frame)] - odom_position_list = [self.odom_poses[i] for i in range(self.processed_frame)] + if self.config.pgo_on: + print("SLAM evaluation:") + avg_tra_slam, avg_rot_slam = relative_error( + gt_poses, pgo_poses + ) + ate_rot_slam, ate_trans_slam, align_mat_slam = absolute_error( + gt_poses, pgo_poses, self.config.eval_traj_align + ) + if avg_tra_slam == 0: # for rgbd dataset (shorter sequence) + print( + "Absoulte Trajectory Error (cm):", + f"{ate_trans_slam*100.0:.3f}", + ) + else: + print("Average Translation Error (%):", f"{avg_tra_slam:.3f}") + print( + "Average Rotational Error (deg/100m):", + f"{avg_rot_slam*100.0:.3f}", + ) + print( + "Absoulte Trajectory Error (m):", f"{ate_trans_slam:.3f}" + ) - if self.config.pgo_on: - pgo_position_list = [self.pgo_poses[i] for i in range(self.processed_frame)] - plot_trajectories( - output_traj_plot_path_2d, - pgo_position_list, - gt_position_list, - odom_position_list, - plot_3d=False, - ) - plot_trajectories( - output_traj_plot_path_3d, - pgo_position_list, - gt_position_list, - odom_position_list, - plot_3d=True, - ) - else: - plot_trajectories( - output_traj_plot_path_2d, - odom_position_list, - gt_position_list, - plot_3d=False, - ) - plot_trajectories( - output_traj_plot_path_3d, - odom_position_list, - gt_position_list, - plot_3d=True, - ) + if self.config.wandb_vis_on: + wandb_log_content = { + "SLAM Average Translation Error [%]": avg_tra_slam, + "SLAM Average Rotational Error [deg/m]": avg_rot_slam, + "SLAM Absoulte Trajectory Error [m]": ate_trans_slam, + "SLAM Absoulte Rotational Error [deg]": ate_rot_slam, + } + wandb.log(wandb_log_content) + + csv_columns = [ + "Average Translation Error [%]", + "Average Rotational Error [deg/m]", + "Absoulte Trajectory Error [m]", + "Absoulte Rotational Error [deg]", + "Consuming time per frame [s]", + "Frame count", + ] + pose_eval = [ + { + csv_columns[0]: avg_tra, + csv_columns[1]: avg_rot, + csv_columns[2]: ate_trans, + csv_columns[3]: ate_rot, + csv_columns[4]: mean_time_without_init_s, + csv_columns[5]: int(self.processed_frame), + } + ] + if self.config.pgo_on: + slam_eval_dict = { + csv_columns[0]: avg_tra_slam, + csv_columns[1]: avg_rot_slam, + csv_columns[2]: ate_trans_slam, + csv_columns[3]: ate_rot_slam, + csv_columns[4]: mean_time_without_init_s, + csv_columns[5]: int(self.processed_frame), + } + pose_eval.append(slam_eval_dict) + output_csv_path = os.path.join(self.run_path, "pose_eval.csv") + try: + with open(output_csv_path, "w") as csvfile: + writer = csv.DictWriter(csvfile, fieldnames=csv_columns) + writer.writeheader() + for data in pose_eval: + writer.writerow(data) + except IOError: + print("I/O error") + + # if self.config.o3d_vis_on: # x service issue for remote server + output_traj_plot_path_2d = os.path.join(self.run_path, "traj_plot_2d.png") + output_traj_plot_path_3d = os.path.join(self.run_path, "traj_plot_3d.png") + # trajectory not aligned yet in the plot + # require list of numpy arraies as the input + + gt_position_list = [self.gt_poses[i] for i in range(self.processed_frame)] + odom_position_list = [self.odom_poses[i] for i in range(self.processed_frame)] + + if self.config.pgo_on: + pgo_position_list = [self.pgo_poses[i] for i in range(self.processed_frame)] + plot_trajectories( + output_traj_plot_path_2d, + pgo_position_list, + gt_position_list, + odom_position_list, + plot_3d=False, + ) + plot_trajectories( + output_traj_plot_path_3d, + pgo_position_list, + gt_position_list, + odom_position_list, + plot_3d=True, + ) + else: + plot_trajectories( + output_traj_plot_path_2d, + odom_position_list, + gt_position_list, + plot_3d=False, + ) + plot_trajectories( + output_traj_plot_path_3d, + odom_position_list, + gt_position_list, + plot_3d=True, + ) return pose_eval diff --git a/model/neural_points.py b/model/neural_points.py index 627f081..e4708bb 100644 --- a/model/neural_points.py +++ b/model/neural_points.py @@ -22,6 +22,7 @@ transform_batch_torch, voxel_down_sample_min_value_torch, voxel_down_sample_torch, + feature_pca_torch, ) @@ -99,6 +100,10 @@ def __init__(self, config: Config) -> None: ) else: self.color_features = None + + # feature pca + self.geo_feature_pca = self.color_feature_pca = None + # here, the ts represent the actually processed frame id (not neccessarily the frame id of the dataset) self.point_ts_create = torch.empty( (0), device=self.device, dtype=torch.int @@ -159,6 +164,14 @@ def print_memory(self): print("Memory consumption: %f (MB)" % cur_memory) self.memory_footprint.append(cur_memory) + + def compute_feature_principle_components(self, down_rate: int = 1): + _, self.geo_feature_pca = feature_pca_torch((self.geo_features)[:-1], down_rate=down_rate, project_data=False) + + if self.color_features is not None: + _, self.color_feature_pca = feature_pca_torch((self.color_features)[:-1], down_rate=down_rate, project_data=False) + + def get_neural_points_o3d( self, query_global: bool = True, @@ -191,35 +204,43 @@ def get_neural_points_o3d( neural_pc_o3d = o3d.geometry.PointCloud() neural_pc_o3d.points = o3d.utility.Vector3dVector(neural_points_np) - if color_mode == 0: # "geo_feature" + if color_mode == 0 and (self.geo_feature_pca is not None): # "geo_feature" if query_global: - neural_features_vis = self.geo_features[:-1:random_down_ratio].detach() + neural_features_vis = self.geo_features[:-1:random_down_ratio] else: neural_features_vis = self.local_geo_features[ :-1:random_down_ratio ].detach() - neural_features_vis = F.normalize(neural_features_vis, p=2, dim=1) - neural_features_np = neural_features_vis.cpu().numpy().astype(np.float64) - neural_pc_o3d.colors = o3d.utility.Vector3dVector( - neural_features_np[:, 0:3] * ratio_vis - ) - elif color_mode == 1: # "color_feature" - if self.color_features is None: - return neural_pc_o3d + geo_feature_3d, _ = feature_pca_torch(neural_features_vis, principal_components=self.geo_feature_pca) # [0,1] + geo_feature_rgb = geo_feature_3d.cpu().numpy().astype(np.float64) + neural_pc_o3d.colors = o3d.utility.Vector3dVector(geo_feature_rgb) + + # neural_features_vis = F.normalize(neural_features_vis, p=2, dim=1) + # neural_features_np = neural_features_vis.cpu().numpy().astype(np.float64) + # neural_pc_o3d.colors = o3d.utility.Vector3dVector( + # neural_features_np[:, 0:3] * ratio_vis + # ) + + elif color_mode == 1 and (self.color_feature_pca is not None) and (self.color_features is not None): # "color_feature" if query_global: neural_features_vis = self.color_features[ :-1:random_down_ratio - ].detach() + ] else: neural_features_vis = self.local_color_features[ :-1:random_down_ratio ].detach() - neural_features_vis = F.normalize(neural_features_vis, p=2, dim=1) - neural_features_np = neural_features_vis.cpu().numpy().astype(np.float64) - neural_pc_o3d.colors = o3d.utility.Vector3dVector( - neural_features_np[:, 0:3] * ratio_vis - ) + + color_feature_3d, _ = feature_pca_torch(neural_features_vis, principal_components=self.color_feature_pca) # [0,1] + color_feature_rgb = color_feature_3d.cpu().numpy().astype(np.float64) + neural_pc_o3d.colors = o3d.utility.Vector3dVector(color_feature_rgb) + + # neural_features_vis = F.normalize(neural_features_vis, p=2, dim=1) + # neural_features_np = neural_features_vis.cpu().numpy().astype(np.float64) + # neural_pc_o3d.colors = o3d.utility.Vector3dVector( + # neural_features_np[:, 0:3] * ratio_vis + # ) elif color_mode == 2: # "ts": # frame number (ts) as the color if query_global: diff --git a/pin_slam.py b/pin_slam.py index 955df5d..91516f9 100644 --- a/pin_slam.py +++ b/pin_slam.py @@ -317,6 +317,7 @@ def run_pin_slam(config_path=None, dataset_name=None, sequence_name=None, seed=N cur_iter_num = max(1, cur_iter_num-10) if frame_id == config.freeze_after_frame: # freeze the decoder after certain frame freeze_decoders(geo_mlp, sem_mlp, color_mlp, config) + neural_points.compute_feature_principle_components(down_rate = 17) # prime number # conduct local bundle adjustment (with lower frequency) if config.track_on and config.ba_freq_frame > 0 and (frame_id+1) % config.ba_freq_frame == 0: @@ -429,9 +430,7 @@ def run_pin_slam(config_path=None, dataset_name=None, sequence_name=None, seed=N dataset.processed_frame += 1 # VI. Save results - pose_eval_results = None - if config.track_on: - pose_eval_results = dataset.write_results() + pose_eval_results = dataset.write_results() if config.pgo_on and pgm.pgo_count>0: print("# Loop corrected: ", pgm.pgo_count) pgm.write_g2o(os.path.join(run_path, "final_pose_graph.g2o")) diff --git a/utils/config.py b/utils/config.py index 7c436c6..a680114 100644 --- a/utils/config.py +++ b/utils/config.py @@ -42,8 +42,6 @@ def __init__(self): self.num_workers: int = 12 # number of worker for the dataloader self.device: str = "cuda" # use "cuda" or "cpu" self.gpu_id: str = "0" # used GPU id - self.dtype = torch.float32 # default torch tensor data type - self.tran_dtype = torch.float64 # dtype used for all the transformation and poses # dataset specific self.kitti_correction_on: bool = False # intrinsic vertical angle correction # issue 11 @@ -206,7 +204,7 @@ def __init__(self): self.ba_bs: int = 16384 # batch size for ba optimization # tracking (odometry estimation) - self.track_on: bool = True + self.track_on: bool = False self.photometric_loss_on: bool = False # add the color (or intensity) [photometric loss] to the tracking loss self.photometric_loss_weight: float = 0.01 # weight for the photometric loss in tracking self.consist_wieght_on: bool = True # weight for color (intensity) consistency for the measured and queried value @@ -300,6 +298,9 @@ def __init__(self): self.republish_raw_input: bool = False # publish the raw input point cloud or not self.timeout_duration_s: int = 30 # in seconds, exit after receiving no topic for x seconds + def setup_dtype(self): + self.dtype = torch.float32 # default torch tensor data type + self.tran_dtype = torch.float64 # dtype used for all the transformation and poses def load(self, config_file): config_args = yaml.safe_load(open(os.path.abspath(config_file))) @@ -346,8 +347,6 @@ def load(self, config_file): self.stop_frame_thre = config_args["setting"].get("stop_frame_thre", self.stop_frame_thre) self.deskew = config_args["setting"].get("deskew", self.deskew) # apply motion undistortion or not - if self.step_frame > 1: - self.deskew = False # process if "process" in config_args: @@ -438,8 +437,8 @@ def load(self, config_file): self.pool_filter_freq = config_args["continual"].get("pool_filter_freq", 1) # tracker - self.track_on = config_args.get("tracker", False) # only on if indicated - if self.track_on: + if "tracker" in config_args: + self.track_on = True if self.color_on: self.photometric_loss_on = config_args["tracker"].get("photo_loss", self.photometric_loss_on) if self.photometric_loss_on: @@ -461,29 +460,29 @@ def load(self, config_file): # pgo if self.track_on: - self.pgo_on = config_args.get("pgo", False) # only on if indicated - if self.pgo_on: - self.local_map_context = config_args["pgo"].get("map_context", self.local_map_context) - self.loop_with_feature = config_args["pgo"].get("loop_with_feature", self.loop_with_feature) - self.local_map_context_latency = config_args["pgo"].get('local_map_latency', self.local_map_context_latency) - self.context_virtual_side_count = config_args["pgo"].get("virtual_side_count", self.context_virtual_side_count) - self.context_virtual_step_m = config_args["pgo"].get("virtual_step_m", self.voxel_size_m * 4.0) - self.npmc_max_dist = config_args["pgo"].get("npmc_max_dist", self.max_range * 0.7) - self.pgo_freq = config_args["pgo"].get("pgo_freq_frame", self.pgo_freq) - self.pgo_with_pose_prior = config_args["pgo"].get("with_pose_prior", self.pgo_with_pose_prior) - # default cov (constant for all the edges) - self.pgo_tran_std = float(config_args["pgo"].get("tran_std", self.pgo_tran_std)) - self.pgo_rot_std = float(config_args["pgo"].get("rot_std", self.pgo_rot_std)) - # use default or estimated cov - self.use_reg_cov_mat = config_args["pgo"].get("use_reg_cov", False) - # merge the neural point map or not after the loop, merge the map may lead to some holes - self.pgo_error_thre = float(config_args["pgo"].get("pgo_error_thre_frame", self.pgo_error_thre_frame)) - self.pgo_max_iter = config_args["pgo"].get("pgo_max_iter", self.pgo_max_iter) - self.pgo_merge_map = config_args["pgo"].get("merge_map", False) - self.context_cosdist_threshold = config_args["pgo"].get("context_cosdist", self.context_cosdist_threshold) - self.min_loop_travel_dist_ratio = config_args["pgo"].get("min_loop_travel_ratio", self.min_loop_travel_dist_ratio) - self.loop_dist_drift_ratio_thre = config_args["pgo"].get("max_loop_dist_ratio", self.loop_dist_drift_ratio_thre) - self.local_loop_dist_thre = config_args["pgo"].get("local_loop_dist_thre", self.voxel_size_m * 5.0) + if "pgo" in config_args: + self.pgo_on = True + self.local_map_context = config_args["pgo"].get("map_context", self.local_map_context) + self.loop_with_feature = config_args["pgo"].get("loop_with_feature", self.loop_with_feature) + self.local_map_context_latency = config_args["pgo"].get('local_map_latency', self.local_map_context_latency) + self.context_virtual_side_count = config_args["pgo"].get("virtual_side_count", self.context_virtual_side_count) + self.context_virtual_step_m = config_args["pgo"].get("virtual_step_m", self.voxel_size_m * 4.0) + self.npmc_max_dist = config_args["pgo"].get("npmc_max_dist", self.max_range * 0.7) + self.pgo_freq = config_args["pgo"].get("pgo_freq_frame", self.pgo_freq) + self.pgo_with_pose_prior = config_args["pgo"].get("with_pose_prior", self.pgo_with_pose_prior) + # default cov (constant for all the edges) + self.pgo_tran_std = float(config_args["pgo"].get("tran_std", self.pgo_tran_std)) + self.pgo_rot_std = float(config_args["pgo"].get("rot_std", self.pgo_rot_std)) + # use default or estimated cov + self.use_reg_cov_mat = config_args["pgo"].get("use_reg_cov", False) + # merge the neural point map or not after the loop, merge the map may lead to some holes + self.pgo_error_thre = float(config_args["pgo"].get("pgo_error_thre_frame", self.pgo_error_thre_frame)) + self.pgo_max_iter = config_args["pgo"].get("pgo_max_iter", self.pgo_max_iter) + self.pgo_merge_map = config_args["pgo"].get("merge_map", False) + self.context_cosdist_threshold = config_args["pgo"].get("context_cosdist", self.context_cosdist_threshold) + self.min_loop_travel_dist_ratio = config_args["pgo"].get("min_loop_travel_ratio", self.min_loop_travel_dist_ratio) + self.loop_dist_drift_ratio_thre = config_args["pgo"].get("max_loop_dist_ratio", self.loop_dist_drift_ratio_thre) + self.local_loop_dist_thre = config_args["pgo"].get("local_loop_dist_thre", self.voxel_size_m * 5.0) # mapping optimizer if "optimizer" in config_args: diff --git a/utils/tools.py b/utils/tools.py index 1b1f722..3468b60 100644 --- a/utils/tools.py +++ b/utils/tools.py @@ -13,6 +13,7 @@ import sys import time import warnings +import yaml from datetime import datetime from pathlib import Path from typing import List @@ -54,6 +55,9 @@ def setup_experiment(config: Config, argv=None, debug_mode: bool = False): else: torch.cuda.empty_cache() + # set the random seed for all + seed_anything(config.seed) + if not debug_mode: access = 0o755 os.makedirs(run_path, access, exist_ok=True) @@ -67,10 +71,12 @@ def setup_experiment(config: Config, argv=None, debug_mode: bool = False): map_path = os.path.join(run_path, "map") model_path = os.path.join(run_path, "model") log_path = os.path.join(run_path, "log") + meta_data_path = os.path.join(run_path, "meta") os.makedirs(mesh_path, access, exist_ok=True) os.makedirs(map_path, access, exist_ok=True) os.makedirs(model_path, access, exist_ok=True) os.makedirs(log_path, access, exist_ok=True) + os.makedirs(meta_data_path, access, exist_ok=True) if config.wandb_vis_on: # set up wandb @@ -97,16 +103,25 @@ def setup_experiment(config: Config, argv=None, debug_mode: bool = False): run_str = "python3 " + " ".join(argv) reproduce_shell.write(run_str) - # set the random seed for all - torch.set_default_dtype(config.dtype) - # set the random seed for all - setup_seed(config.seed) + # disable lidar deskewing when not input per frame + if config.step_frame > 1: + config.deskew = False + + # write the full configs to yaml file + config_dict = vars(config) + config_out_path = os.path.join(meta_data_path, "config_all.yaml") + with open(config_out_path, 'w') as file: + yaml.dump(config_dict, file, default_flow_style=False) + + # set up dtypes, note that torch stuff cannot be write to yaml, so we set it up after write out the yaml for the whole config + config.setup_dtype() + torch.set_default_dtype(config.dtype) return run_path -def setup_seed(seed): +def seed_anything(seed): os.environ["PYTHONHASHSEED"] = str(seed) torch.manual_seed(seed) np.random.seed(seed) @@ -701,7 +716,7 @@ def tranmat_close_to_identity(mats: np.ndarray, rot_thre: float, tran_thre: floa rot_close_to_identity = np.all(rot_diff < rot_thre) - tran_diff = mats[:3, 3] + tran_diff = np.abs(mats[:3, 3]) tran_close_to_identity = np.all(tran_diff < tran_thre) @@ -710,6 +725,60 @@ def tranmat_close_to_identity(mats: np.ndarray, rot_thre: float, tran_thre: floa else: return False +def feature_pca_torch(data, principal_components = None, + principal_dim: int = 3, + down_rate: int = 1, + project_data: bool = True, + normalize: bool = True): + """ + do PCA to a NxD torch tensor to get the data along the K principle dimensions + N is the data count, D is the dimension of the data + + We can also use a pre-computed principal_components for only the projection of input data + """ + + N, D = data.shape + + # Step 1: Center the data (subtract the mean of each dimension) + data_centered = data - data.mean(dim=0) + + if principal_components is None: + data_centered_for_compute = data_centered[::down_rate] + + assert data_centered_for_compute.shape[0] > principal_dim, "not enough data for PCA computation, down_rate might be too large or original data count is too small" + + # Step 2: Compute the covariance matrix (D x D) + cov_matrix = torch.matmul(data_centered_for_compute.T, data_centered_for_compute) / (N - 1) + + # Step 3: Perform eigen decomposition of the covariance matrix + eigenvalues, eigenvectors = torch.linalg.eig(cov_matrix) + eigenvalues_r = eigenvalues.real.to(data) + eigenvectors_r = eigenvectors.real.to(data) + + # Step 4: Sort eigenvalues and eigenvectors in descending order + sorted_indices = torch.argsort(eigenvalues_r, descending=True) + principal_components = eigenvectors_r[:, sorted_indices[:principal_dim]] # First 3 principal components + + data_pca = None + if project_data: + # Step 5: Project data onto the top 3 principal components + data_pca = torch.matmul(data_centered, principal_components) # N, D @ D, P + + # normalize to show as rgb + if normalize: + # min_vals = data_pca.min(dim=0, keepdim=True).values + # max_vals = data_pca.max(dim=0, keepdim=True).values + + # # deal with outliers + min_vals = torch.quantile(data_pca, 0.02, dim=0, keepdim=True) + max_vals = torch.quantile(data_pca, 0.98, dim=0, keepdim=True) + + # Normalize to range [0, 1] + data_pca = (data_pca - min_vals) / (max_vals - min_vals) + + data_pca = data_pca.clamp(0, 1) + + return data_pca, principal_components def plot_timing_detail(time_table: np.ndarray, saving_path: str, with_loop=False): diff --git a/vis_pin_map.py b/vis_pin_map.py index a100d5e..431e009 100644 --- a/vis_pin_map.py +++ b/vis_pin_map.py @@ -40,19 +40,26 @@ def vis_pin_map(): print("[bold green]Load PIN Map[/bold green]","📍" ) + run_path = setup_experiment(config, sys.argv, debug_mode=True) + + o3d_vis_on = True # set to False if you don't want to visualize # initialize the mlp decoder geo_mlp = Decoder(config, config.geo_mlp_hidden_dim, config.geo_mlp_level, 1) - sem_mlp = Decoder(config, config.sem_mlp_hidden_dim, config.sem_mlp_level, config.sem_class_count + 1) - color_mlp = Decoder(config, config.color_mlp_hidden_dim, config.color_mlp_level, config.color_channel) + color_mlp = None + if config.color_channel > 0: + color_mlp = Decoder(config, config.color_mlp_hidden_dim, config.color_mlp_level, config.color_channel) + sem_mlp = Decoder(config, config.sem_mlp_hidden_dim, config.sem_mlp_level, config.sem_class_count + 1) # initialize the neural point features - neural_points = NeuralPoints(config) + neural_points: NeuralPoints = NeuralPoints(config) # Load the map loaded_model = torch.load(config.model_path) neural_points = loaded_model["neural_points"] + neural_points.temporal_local_map_on = False + neural_points.compute_feature_principle_components(down_rate = 17) # print(loaded_model.keys()) geo_mlp.load_state_dict(loaded_model["geo_decoder"]) @@ -62,7 +69,7 @@ def vis_pin_map(): color_mlp.load_state_dict(loaded_model["color_decoder"]) print("PIN Map loaded") - if config.o3d_vis_on: + if o3d_vis_on: vis = MapVisualizer(config) neural_points.recreate_hash(neural_points.neural_points[0], torch.eye(3).cuda(), False, False) @@ -123,7 +130,7 @@ def vis_pin_map(): cur_mesh = mesher.recon_aabb_collections_mesh(chunks_aabb, mesh_vox_size_m, out_mesh_path, False, config.semantic_on, config.color_on, filter_isolated_mesh=True, mesh_min_nn=mesh_min_nn_used) - if config.o3d_vis_on: + if o3d_vis_on: while True: if vis.render_neural_points: neural_pcd = neural_points.get_neural_points_o3d(query_global=True, color_mode=vis.neural_points_vis_mode, random_down_ratio=down_rate)