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 sys.path.append(r"/media/hofee/data/project/python/nbv_reconstruction/nbv_reconstruction") from utils.data_load import DataLoadUtil from utils.pose import PoseUtil from utils.pts import PtsUtil from utils.reconstruction import ReconstructionUtil @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: self.datalist = self.datalist*100 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: label_path = DataLoadUtil.get_label_path(self.root_dir, scene_name) label_data = DataLoadUtil.load_label(label_path) for data_pair in label_data["data_pairs"]: scanned_views = data_pair[0] next_best_view = data_pair[1] max_coverage_rate = label_data["max_coverage_rate"] datalist.append( { "scanned_views": scanned_views, "next_best_view": next_best_view, "max_coverage_rate": max_coverage_rate, "scene_name": scene_name, } ) 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}") # ----- Debug Trace ----- # import ipdb; ipdb.set_trace() # ------------------------ # 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["max_coverage_rate"] scene_name = data_item_info["scene_name"] scanned_views_pts, scanned_coverages_rate, scanned_n_to_world_pose = [], [], [] 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"] nR_to_world_pose = cam_info["cam_to_world_R"] if self.load_from_preprocess: downsampled_target_point_cloud = DataLoadUtil.load_from_preprocessed_pts(view_path) else: cached_data = None if self.cache: cached_data = self.load_from_cache(scene_name, frame_idx) if cached_data is None: print("load depth") depth_L, depth_R = DataLoadUtil.load_depth(view_path, cam_info['near_plane'], cam_info['far_plane'], binocular=True) point_cloud_L = DataLoadUtil.get_point_cloud(depth_L, cam_info['cam_intrinsic'], n_to_world_pose)['points_world'] point_cloud_R = DataLoadUtil.get_point_cloud(depth_R, cam_info['cam_intrinsic'], nR_to_world_pose)['points_world'] point_cloud_L = PtsUtil.random_downsample_point_cloud(point_cloud_L, 65536) point_cloud_R = PtsUtil.random_downsample_point_cloud(point_cloud_R, 65536) overlap_points = DataLoadUtil.get_overlapping_points(point_cloud_L, point_cloud_R) downsampled_target_point_cloud = PtsUtil.random_downsample_point_cloud(overlap_points, self.pts_num) if self.cache: self.save_to_cache(scene_name, frame_idx, downsampled_target_point_cloud) else: downsampled_target_point_cloud = cached_data 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) 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) data_item = { "scanned_pts": np.asarray(scanned_views_pts,dtype=np.float32), "scanned_coverage_rate": scanned_coverages_rate, "scanned_n_to_world_pose_9d": np.asarray(scanned_n_to_world_pose,dtype=np.float32), "best_coverage_rate": nbv_coverage_rate, "best_to_world_pose_9d": np.asarray(best_to_world_9d,dtype=np.float32), "max_coverage_rate": max_coverage_rate, "scene_name": scene_name } # if self.type == namespace.Mode.TEST: # diag = DataLoadUtil.get_bbox_diag(self.model_dir, scene_name) # voxel_threshold = diag*0.02 # model_points_normals = DataLoadUtil.load_points_normals(self.root_dir, scene_name) # pts_list = [] # for view in scanned_views: # frame_idx = view[0] # view_path = DataLoadUtil.get_path(self.root_dir, scene_name, frame_idx) # point_cloud = DataLoadUtil.get_target_point_cloud_world_from_path(view_path, binocular=True) # cam_params = DataLoadUtil.load_cam_info(view_path, binocular=True) # sampled_point_cloud = ReconstructionUtil.filter_points(point_cloud, model_points_normals, cam_pose=cam_params["cam_to_world"], voxel_size=voxel_threshold, theta=self.filter_degree) # pts_list.append(sampled_point_cloud) # nL_to_world_pose = cam_params["cam_to_world"] # nO_to_world_pose = cam_params["cam_to_world_O"] # nO_to_nL_pose = np.dot(np.linalg.inv(nL_to_world_pose), nO_to_world_pose) # data_item["scanned_target_pts_list"] = pts_list # data_item["model_points_normals"] = model_points_normals # data_item["voxel_threshold"] = voxel_threshold # data_item["filter_degree"] = self.filter_degree # data_item["scene_path"] = os.path.join(self.root_dir, scene_name) # data_item["first_frame_to_world"] = np.asarray(first_frame_to_world, dtype=np.float32) # data_item["nO_to_nL_pose"] = np.asarray(nO_to_nL_pose, dtype=np.float32) return data_item def __len__(self): return len(self.datalist) def get_collate_fn(self): def collate_fn(batch): collate_data = {} 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["best_to_world_pose_9d"] = torch.stack([torch.tensor(item['best_to_world_pose_9d']) for item in batch]) if "first_frame_to_world" in batch[0]: collate_data["first_frame_to_world"] = torch.stack([torch.tensor(item["first_frame_to_world"]) 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", "first_frame_to_world"]: 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": "/media/hofee/data/project/python/nbv_reconstruction/sample_for_training/preprocessed_scenes/", "model_dir": "/media/hofee/data/data/scaled_object_meshes", "source": "nbv_reconstruction_dataset", "split_file": "/media/hofee/data/project/python/nbv_reconstruction/sample_for_training/OmniObject3d_train.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 ------ # # for idx, data in enumerate(dl): # cnt=0 # print(data["scene_name"]) # print(data["scanned_coverage_rate"]) # print(data["best_coverage_rate"]) # for pts in data["scanned_pts"][0]: # #np.savetxt(f"pts_{cnt}.txt", pts) # cnt+=1 # #np.savetxt("best_pts.txt", best_pts) # for key, value in data.items(): # if isinstance(value, torch.Tensor): # print(key, ":" ,value.shape) # else: # print(key, ":" ,len(value)) # if key == "scanned_n_to_world_pose_9d": # for val in value: # print(val.shape) # if key == "scanned_pts": # print("scanned_pts") # for val in value: # print(val.shape) # cnt = 0 # for v in val: # import ipdb;ipdb.set_trace() # np.savetxt(f"pts_{cnt}.txt", v) # cnt+=1 # print()