import numpy as np from PytorchBoot.dataset import BaseDataset import PytorchBoot.namespace as namespace import PytorchBoot.stereotype as stereotype from PytorchBoot.config import ConfigManager from PytorchBoot.utils.log_util import Log import torch import os import sys import time sys.path.append(r"/data/hofee/project/nbv_rec/nbv_reconstruction") from utils.data_load import DataLoadUtil from utils.pose import PoseUtil from utils.pts import PtsUtil @stereotype.dataset("nbv_reconstruction_dataset") class NBVReconstructionDataset(BaseDataset): def __init__(self, config): super(NBVReconstructionDataset, self).__init__(config) self.config = config self.root_dir = config["root_dir"] self.split_file_path = config["split_file"] self.scene_name_list = self.load_scene_name_list() self.datalist = self.get_datalist() self.pts_num = config["pts_num"] self.type = config["type"] self.cache = config.get("cache") self.load_from_preprocess = config.get("load_from_preprocess", False) if self.type == namespace.Mode.TEST: #self.model_dir = config["model_dir"] self.filter_degree = config["filter_degree"] if self.type == namespace.Mode.TRAIN: scale_ratio = 1 self.datalist = self.datalist*scale_ratio if self.cache: expr_root = ConfigManager.get("runner", "experiment", "root_dir") expr_name = ConfigManager.get("runner", "experiment", "name") self.cache_dir = os.path.join(expr_root, expr_name, "cache") # self.preprocess_cache() def load_scene_name_list(self): scene_name_list = [] with open(self.split_file_path, "r") as f: for line in f: scene_name = line.strip() scene_name_list.append(scene_name) return scene_name_list def get_datalist(self): datalist = [] for scene_name in self.scene_name_list: seq_num = DataLoadUtil.get_label_num(self.root_dir, scene_name) scene_max_coverage_rate = 0 max_coverage_rate_list = [] for seq_idx in range(seq_num): label_path = DataLoadUtil.get_label_path( self.root_dir, scene_name, seq_idx ) label_data = DataLoadUtil.load_label(label_path) max_coverage_rate = label_data["max_coverage_rate"] if max_coverage_rate > scene_max_coverage_rate: scene_max_coverage_rate = max_coverage_rate max_coverage_rate_list.append(max_coverage_rate) if max_coverage_rate_list: mean_coverage_rate = np.mean(max_coverage_rate_list) for seq_idx in range(seq_num): label_path = DataLoadUtil.get_label_path( self.root_dir, scene_name, seq_idx ) label_data = DataLoadUtil.load_label(label_path) if max_coverage_rate_list[seq_idx] > mean_coverage_rate - 0.1: for data_pair in label_data["data_pairs"]: scanned_views = data_pair[0] next_best_view = data_pair[1] datalist.append( { "scanned_views": scanned_views, "next_best_view": next_best_view, "seq_max_coverage_rate": max_coverage_rate, "scene_name": scene_name, "label_idx": seq_idx, "scene_max_coverage_rate": scene_max_coverage_rate, } ) return datalist def preprocess_cache(self): Log.info("preprocessing cache...") for item_idx in range(len(self.datalist)): self.__getitem__(item_idx) Log.success("finish preprocessing cache.") def load_from_cache(self, scene_name, curr_frame_idx): cache_name = f"{scene_name}_{curr_frame_idx}.txt" cache_path = os.path.join(self.cache_dir, cache_name) if os.path.exists(cache_path): data = np.loadtxt(cache_path) return data else: return None def save_to_cache(self, scene_name, curr_frame_idx, data): cache_name = f"{scene_name}_{curr_frame_idx}.txt" cache_path = os.path.join(self.cache_dir, cache_name) try: np.savetxt(cache_path, data) except Exception as e: Log.error(f"Save cache failed: {e}") def voxel_downsample_with_mapping(self, point_cloud, voxel_size=0.003): voxel_indices = np.floor(point_cloud / voxel_size).astype(np.int32) unique_voxels, inverse, counts = np.unique(voxel_indices, axis=0, return_inverse=True, return_counts=True) idx_sort = np.argsort(inverse) idx_unique = idx_sort[np.cumsum(counts)-counts] downsampled_points = point_cloud[idx_unique] return downsampled_points, inverse def __getitem__(self, index): data_item_info = self.datalist[index] scanned_views = data_item_info["scanned_views"] nbv = data_item_info["next_best_view"] max_coverage_rate = data_item_info["seq_max_coverage_rate"] scene_name = data_item_info["scene_name"] ( scanned_views_pts, scanned_coverages_rate, scanned_n_to_world_pose, ) = ([], [], []) #start_time = time.time() start_indices = [0] total_points = 0 for view in scanned_views: frame_idx = view[0] coverage_rate = view[1] view_path = DataLoadUtil.get_path(self.root_dir, scene_name, frame_idx) cam_info = DataLoadUtil.load_cam_info(view_path, binocular=True) n_to_world_pose = cam_info["cam_to_world"] target_point_cloud = ( DataLoadUtil.load_from_preprocessed_pts(view_path) ) downsampled_target_point_cloud = PtsUtil.random_downsample_point_cloud( target_point_cloud, self.pts_num ) scanned_views_pts.append(downsampled_target_point_cloud) scanned_coverages_rate.append(coverage_rate) n_to_world_6d = PoseUtil.matrix_to_rotation_6d_numpy( np.asarray(n_to_world_pose[:3, :3]) ) n_to_world_trans = n_to_world_pose[:3, 3] n_to_world_9d = np.concatenate([n_to_world_6d, n_to_world_trans], axis=0) scanned_n_to_world_pose.append(n_to_world_9d) total_points += len(downsampled_target_point_cloud) start_indices.append(total_points) #end_time = time.time() #Log.info(f"load data time: {end_time - start_time}") nbv_idx, nbv_coverage_rate = nbv[0], nbv[1] nbv_path = DataLoadUtil.get_path(self.root_dir, scene_name, nbv_idx) cam_info = DataLoadUtil.load_cam_info(nbv_path) best_frame_to_world = cam_info["cam_to_world"] best_to_world_6d = PoseUtil.matrix_to_rotation_6d_numpy( np.asarray(best_frame_to_world[:3, :3]) ) best_to_world_trans = best_frame_to_world[:3, 3] best_to_world_9d = np.concatenate( [best_to_world_6d, best_to_world_trans], axis=0 ) combined_scanned_views_pts = np.concatenate(scanned_views_pts, axis=0) voxel_downsampled_combined_scanned_pts_np, inverse = self.voxel_downsample_with_mapping(combined_scanned_views_pts, 0.003) random_downsampled_combined_scanned_pts_np, random_downsample_idx = PtsUtil.random_downsample_point_cloud(voxel_downsampled_combined_scanned_pts_np, self.pts_num, require_idx=True) # all_idx_unique = np.arange(len(voxel_downsampled_combined_scanned_pts_np)) # all_random_downsample_idx = all_idx_unique[random_downsample_idx] # scanned_pts_mask = [] # for idx, start_idx in enumerate(start_indices): # if idx == len(start_indices) - 1: # break # end_idx = start_indices[idx+1] # view_inverse = inverse[start_idx:end_idx] # view_unique_downsampled_idx = np.unique(view_inverse) # view_unique_downsampled_idx_set = set(view_unique_downsampled_idx) # mask = np.array([idx in view_unique_downsampled_idx_set for idx in all_random_downsample_idx]) # #scanned_pts_mask.append(mask) data_item = { "scanned_pts": np.asarray(scanned_views_pts, dtype=np.float32), # Ndarray(S x Nv x 3) "combined_scanned_pts": np.asarray(random_downsampled_combined_scanned_pts_np, dtype=np.float32), # Ndarray(N x 3) #"scanned_pts_mask": np.asarray(scanned_pts_mask, dtype=np.bool), # Ndarray(N) "scanned_coverage_rate": scanned_coverages_rate, # List(S): Float, range(0, 1) "scanned_n_to_world_pose_9d": np.asarray(scanned_n_to_world_pose, dtype=np.float32), # Ndarray(S x 9) "best_coverage_rate": nbv_coverage_rate, # Float, range(0, 1) "best_to_world_pose_9d": np.asarray(best_to_world_9d, dtype=np.float32), # Ndarray(9) "seq_max_coverage_rate": max_coverage_rate, # Float, range(0, 1) "scene_name": scene_name, # String } return data_item def __len__(self): return len(self.datalist) def get_collate_fn(self): def collate_fn(batch): collate_data = {} ''' ------ Varialbe Length ------ ''' collate_data["scanned_pts"] = [ torch.tensor(item["scanned_pts"]) for item in batch ] collate_data["scanned_n_to_world_pose_9d"] = [ torch.tensor(item["scanned_n_to_world_pose_9d"]) for item in batch ] # collate_data["scanned_pts_mask"] = [ # torch.tensor(item["scanned_pts_mask"]) for item in batch # ] ''' ------ Fixed Length ------ ''' collate_data["best_to_world_pose_9d"] = torch.stack( [torch.tensor(item["best_to_world_pose_9d"]) for item in batch] ) collate_data["combined_scanned_pts"] = torch.stack( [torch.tensor(item["combined_scanned_pts"]) for item in batch] ) for key in batch[0].keys(): if key not in [ "scanned_pts", "scanned_n_to_world_pose_9d", "best_to_world_pose_9d", "combined_scanned_pts", "scanned_pts_mask", ]: collate_data[key] = [item[key] for item in batch] return collate_data return collate_fn # -------------- Debug ---------------- # if __name__ == "__main__": import torch seed = 0 torch.manual_seed(seed) np.random.seed(seed) config = { "root_dir": "/data/hofee/nbv_rec_part2_preprocessed", "source": "nbv_reconstruction_dataset", "split_file": "/data/hofee/data/sample.txt", "load_from_preprocess": True, "ratio": 0.5, "batch_size": 2, "filter_degree": 75, "num_workers": 0, "pts_num": 4096, "type": namespace.Mode.TRAIN, } ds = NBVReconstructionDataset(config) print(len(ds)) # ds.__getitem__(10) dl = ds.get_loader(shuffle=True) for idx, data in enumerate(dl): data = ds.process_batch(data, "cuda:0") print(data) # ------ Debug Start ------ import ipdb ipdb.set_trace() # ------ Debug End ------