import numpy as np from PytorchBoot.dataset import BaseDataset import PytorchBoot.stereotype as stereotype from torch.nn.utils.rnn import pad_sequence import sys sys.path.append(r"C:\Document\Local 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"] 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 __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_1_pose = [], [], [] first_frame_idx = scanned_views[0][0] first_frame_to_world = DataLoadUtil.load_cam_info(DataLoadUtil.get_path(self.root_dir, scene_name, first_frame_idx))["cam_to_world"] 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) depth = DataLoadUtil.load_depth(view_path) cam_info = DataLoadUtil.load_cam_info(view_path) mask = DataLoadUtil.load_seg(view_path) frame_curr_to_world = cam_info["cam_to_world"] n_to_1_pose = np.dot(np.linalg.inv(first_frame_to_world), frame_curr_to_world) target_point_cloud = DataLoadUtil.get_target_point_cloud(depth, cam_info["cam_intrinsic"], n_to_1_pose, mask)["points_world"] 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_1_6d = PoseUtil.matrix_to_rotation_6d_numpy(np.asarray(n_to_1_pose[:3,:3])) n_to_1_trans = n_to_1_pose[:3,3] n_to_1_9d = np.concatenate([n_to_1_6d, n_to_1_trans], axis=0) scanned_n_to_1_pose.append(n_to_1_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_1_pose = np.dot(np.linalg.inv(first_frame_to_world), best_frame_to_world) best_to_1_6d = PoseUtil.matrix_to_rotation_6d_numpy(np.asarray(best_to_1_pose[:3,:3])) best_to_1_trans = best_to_1_pose[:3,3] best_to_1_9d = np.concatenate([best_to_1_6d, best_to_1_trans], axis=0) data_item = { "scanned_pts": np.asarray(scanned_views_pts,dtype=np.float32), "scanned_coverage_rate": np.asarray(scanned_coverages_rate,dtype=np.float32), "scanned_n_to_1_pose_9d": np.asarray(scanned_n_to_1_pose,dtype=np.float32), "best_coverage_rate": nbv_coverage_rate, "best_to_1_pose_9d": best_to_1_9d, "max_coverage_rate": max_coverage_rate, "scene_name": scene_name } return data_item def __len__(self): return len(self.datalist) def get_collate_fn(self): def collate_fn(batch): scanned_pts = [item['scanned_pts'] for item in batch] scanned_n_to_1_pose_9d = [item['scanned_n_to_1_pose_9d'] for item in batch] rest = {} for key in batch[0].keys(): if key in ['scanned_pts', 'scanned_n_to_1_pose_9d']: continue if isinstance(batch[0][key], torch.Tensor): rest[key] = torch.stack([item[key] for item in batch]) elif isinstance(batch[0][key], str): rest[key] = [item[key] for item in batch] else: rest[key] = [item[key] for item in batch] return { 'scanned_pts': scanned_pts, 'scanned_n_to_1_pose_9d': scanned_n_to_1_pose_9d, **rest } return collate_fn if __name__ == "__main__": import torch seed = 0 torch.manual_seed(seed) np.random.seed(seed) config = { "root_dir": "C:\\Document\\Local Project\\nbv_rec\\data\\sample", "split_file": "C:\\Document\\Local Project\\nbv_rec\\data\\OmniObject3d_train.txt", "ratio": 0.5, "batch_size": 2, "num_workers": 0, "pts_num": 2048 } ds = NBVReconstructionDataset(config) print(len(ds)) dl = ds.get_loader(shuffle=True) 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_1_pose_9d": for val in value: print(val.shape) if key == "scanned_pts": for val in value: print(val.shape) print()