From b221036e8bed94ca23739a3a3836a1c27ce49392 Mon Sep 17 00:00:00 2001 From: hofee Date: Thu, 31 Oct 2024 16:02:26 +0000 Subject: [PATCH] global: upd --- core/nbv_dataset.py | 2 +- core/old_seq_dataset.py | 154 +++++++++++++++++++ core/seq_dataset.py | 331 +++++++++++++++++++++------------------- 3 files changed, 332 insertions(+), 155 deletions(-) create mode 100644 core/old_seq_dataset.py diff --git a/core/nbv_dataset.py b/core/nbv_dataset.py index 36f4ed1..ca9c0c7 100644 --- a/core/nbv_dataset.py +++ b/core/nbv_dataset.py @@ -48,7 +48,7 @@ class NBVReconstructionDataset(BaseDataset): for line in f: scene_name = line.strip() scene_name_list.append(scene_name) - return scene_name_list[:10] + return scene_name_list def get_datalist(self): datalist = [] diff --git a/core/old_seq_dataset.py b/core/old_seq_dataset.py new file mode 100644 index 0000000..753636e --- /dev/null +++ b/core/old_seq_dataset.py @@ -0,0 +1,154 @@ +import numpy as np +from PytorchBoot.dataset import BaseDataset +import PytorchBoot.namespace as namespace +import PytorchBoot.stereotype as stereotype +from PytorchBoot.utils.log_util import Log +import torch +import os +import sys +sys.path.append(r"/home/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("old_seq_nbv_reconstruction_dataset") +class SeqNBVReconstructionDataset(BaseDataset): + def __init__(self, config): + super(SeqNBVReconstructionDataset, self).__init__(config) + self.type = config["type"] + if self.type != namespace.Mode.TEST: + Log.error("Dataset Only support test mode", terminate=True) + 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.model_dir = config["model_dir"] + self.filter_degree = config["filter_degree"] + self.load_from_preprocess = config.get("load_from_preprocess", False) + + + 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 + scene_max_cr_idx = 0 + + 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 + scene_max_cr_idx = seq_idx + + label_path = DataLoadUtil.get_label_path(self.root_dir, scene_name, scene_max_cr_idx) + label_data = DataLoadUtil.load_label(label_path) + first_frame = label_data["best_sequence"][0] + best_seq_len = len(label_data["best_sequence"]) + datalist.append({ + "scene_name": scene_name, + "first_frame": first_frame, + "max_coverage_rate": scene_max_coverage_rate, + "best_seq_len": best_seq_len, + "label_idx": scene_max_cr_idx, + }) + return datalist + + def __getitem__(self, index): + data_item_info = self.datalist[index] + first_frame_idx = data_item_info["first_frame"][0] + first_frame_coverage = data_item_info["first_frame"][1] + max_coverage_rate = data_item_info["max_coverage_rate"] + scene_name = data_item_info["scene_name"] + first_cam_info = DataLoadUtil.load_cam_info(DataLoadUtil.get_path(self.root_dir, scene_name, first_frame_idx), binocular=True) + first_view_path = DataLoadUtil.get_path(self.root_dir, scene_name, first_frame_idx) + first_left_cam_pose = first_cam_info["cam_to_world"] + first_center_cam_pose = first_cam_info["cam_to_world_O"] + first_target_point_cloud = DataLoadUtil.load_from_preprocessed_pts(first_view_path) + first_pts_num = first_target_point_cloud.shape[0] + first_downsampled_target_point_cloud = PtsUtil.random_downsample_point_cloud(first_target_point_cloud, self.pts_num) + first_to_world_rot_6d = PoseUtil.matrix_to_rotation_6d_numpy(np.asarray(first_left_cam_pose[:3,:3])) + first_to_world_trans = first_left_cam_pose[:3,3] + first_to_world_9d = np.concatenate([first_to_world_rot_6d, first_to_world_trans], axis=0) + diag = DataLoadUtil.get_bbox_diag(self.model_dir, scene_name) + voxel_threshold = diag*0.02 + first_O_to_first_L_pose = np.dot(np.linalg.inv(first_left_cam_pose), first_center_cam_pose) + scene_path = os.path.join(self.root_dir, scene_name) + model_points_normals = DataLoadUtil.load_points_normals(self.root_dir, scene_name) + + data_item = { + "first_pts_num": np.asarray( + first_pts_num, dtype=np.int32 + ), + "first_pts": np.asarray([first_downsampled_target_point_cloud],dtype=np.float32), + "combined_scanned_pts": np.asarray(first_downsampled_target_point_cloud,dtype=np.float32), + "first_to_world_9d": np.asarray([first_to_world_9d],dtype=np.float32), + "scene_name": scene_name, + "max_coverage_rate": max_coverage_rate, + "voxel_threshold": voxel_threshold, + "filter_degree": self.filter_degree, + "O_to_L_pose": first_O_to_first_L_pose, + "first_frame_coverage": first_frame_coverage, + "scene_path": scene_path, + "model_points_normals": model_points_normals, + "best_seq_len": data_item_info["best_seq_len"], + "first_frame_id": first_frame_idx, + } + return data_item + + def __len__(self): + return len(self.datalist) + + def get_collate_fn(self): + def collate_fn(batch): + collate_data = {} + collate_data["first_pts"] = [torch.tensor(item['first_pts']) for item in batch] + collate_data["first_to_world_9d"] = [torch.tensor(item['first_to_world_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 ["first_pts", "first_to_world_9d", "combined_scanned_pts"]: + 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": "/home/data/hofee/project/nbv_rec/data/nbv_rec_data_512_preproc_npy", + "split_file": "/home/data/hofee/project/nbv_rec/data/OmniObject3d_train.txt", + "model_dir": "/home/data/hofee/project/nbv_rec/data/scaled_object_meshes", + "ratio": 0.005, + "batch_size": 2, + "filter_degree": 75, + "num_workers": 0, + "pts_num": 32684, + "type": namespace.Mode.TEST, + "load_from_preprocess": True + } + ds = SeqNBVReconstructionDataset(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 ------+ \ No newline at end of file diff --git a/core/seq_dataset.py b/core/seq_dataset.py index 3196949..1e816df 100644 --- a/core/seq_dataset.py +++ b/core/seq_dataset.py @@ -1,154 +1,177 @@ -import numpy as np -from PytorchBoot.dataset import BaseDataset -import PytorchBoot.namespace as namespace -import PytorchBoot.stereotype as stereotype -from PytorchBoot.utils.log_util import Log -import torch -import os -import sys -sys.path.append(r"/home/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("seq_nbv_reconstruction_dataset") -class SeqNBVReconstructionDataset(BaseDataset): - def __init__(self, config): - super(SeqNBVReconstructionDataset, self).__init__(config) - self.type = config["type"] - if self.type != namespace.Mode.TEST: - Log.error("Dataset Only support test mode", terminate=True) - 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.model_dir = config["model_dir"] - self.filter_degree = config["filter_degree"] - self.load_from_preprocess = config.get("load_from_preprocess", False) - - - 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 - scene_max_cr_idx = 0 - - 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 - scene_max_cr_idx = seq_idx - - label_path = DataLoadUtil.get_label_path(self.root_dir, scene_name, scene_max_cr_idx) - label_data = DataLoadUtil.load_label(label_path) - first_frame = label_data["best_sequence"][0] - best_seq_len = len(label_data["best_sequence"]) - datalist.append({ - "scene_name": scene_name, - "first_frame": first_frame, - "max_coverage_rate": scene_max_coverage_rate, - "best_seq_len": best_seq_len, - "label_idx": scene_max_cr_idx, - }) - return datalist - - def __getitem__(self, index): - data_item_info = self.datalist[index] - first_frame_idx = data_item_info["first_frame"][0] - first_frame_coverage = data_item_info["first_frame"][1] - max_coverage_rate = data_item_info["max_coverage_rate"] - scene_name = data_item_info["scene_name"] - first_cam_info = DataLoadUtil.load_cam_info(DataLoadUtil.get_path(self.root_dir, scene_name, first_frame_idx), binocular=True) - first_view_path = DataLoadUtil.get_path(self.root_dir, scene_name, first_frame_idx) - first_left_cam_pose = first_cam_info["cam_to_world"] - first_center_cam_pose = first_cam_info["cam_to_world_O"] - first_target_point_cloud = DataLoadUtil.load_from_preprocessed_pts(first_view_path) - first_pts_num = first_target_point_cloud.shape[0] - first_downsampled_target_point_cloud = PtsUtil.random_downsample_point_cloud(first_target_point_cloud, self.pts_num) - first_to_world_rot_6d = PoseUtil.matrix_to_rotation_6d_numpy(np.asarray(first_left_cam_pose[:3,:3])) - first_to_world_trans = first_left_cam_pose[:3,3] - first_to_world_9d = np.concatenate([first_to_world_rot_6d, first_to_world_trans], axis=0) - diag = DataLoadUtil.get_bbox_diag(self.model_dir, scene_name) - voxel_threshold = diag*0.02 - first_O_to_first_L_pose = np.dot(np.linalg.inv(first_left_cam_pose), first_center_cam_pose) - scene_path = os.path.join(self.root_dir, scene_name) - model_points_normals = DataLoadUtil.load_points_normals(self.root_dir, scene_name) - - data_item = { - "first_pts_num": np.asarray( - first_pts_num, dtype=np.int32 - ), - "first_pts": np.asarray([first_downsampled_target_point_cloud],dtype=np.float32), - "combined_scanned_pts": np.asarray(first_downsampled_target_point_cloud,dtype=np.float32), - "first_to_world_9d": np.asarray([first_to_world_9d],dtype=np.float32), - "scene_name": scene_name, - "max_coverage_rate": max_coverage_rate, - "voxel_threshold": voxel_threshold, - "filter_degree": self.filter_degree, - "O_to_L_pose": first_O_to_first_L_pose, - "first_frame_coverage": first_frame_coverage, - "scene_path": scene_path, - "model_points_normals": model_points_normals, - "best_seq_len": data_item_info["best_seq_len"], - "first_frame_id": first_frame_idx, - } - return data_item - - def __len__(self): - return len(self.datalist) - - def get_collate_fn(self): - def collate_fn(batch): - collate_data = {} - collate_data["first_pts"] = [torch.tensor(item['first_pts']) for item in batch] - collate_data["first_to_world_9d"] = [torch.tensor(item['first_to_world_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 ["first_pts", "first_to_world_9d", "combined_scanned_pts"]: - 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": "/home/data/hofee/project/nbv_rec/data/nbv_rec_data_512_preproc_npy", - "split_file": "/home/data/hofee/project/nbv_rec/data/OmniObject3d_train.txt", - "model_dir": "/home/data/hofee/project/nbv_rec/data/scaled_object_meshes", - "ratio": 0.005, - "batch_size": 2, - "filter_degree": 75, - "num_workers": 0, - "pts_num": 32684, - "type": namespace.Mode.TEST, - "load_from_preprocess": True - } - ds = SeqNBVReconstructionDataset(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 ------+ \ No newline at end of file +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"/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("seq_reconstruction_dataset") +class SeqReconstructionDataset(BaseDataset): + def __init__(self, config): + super(SeqReconstructionDataset, 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 = [] + scene_max_cr_idx = 0 + 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 + scene_max_cr_idx = seq_idx + max_coverage_rate_list.append(max_coverage_rate) + best_label_path = DataLoadUtil.get_label_path(self.root_dir, scene_name, scene_max_cr_idx) + best_label_data = DataLoadUtil.load_label(best_label_path) + first_frame = best_label_data["best_sequence"][0] + best_seq_len = len(best_label_data["best_sequence"]) + datalist.append({ + "scene_name": scene_name, + "first_frame": first_frame, + "best_seq_len": best_seq_len, + "max_coverage_rate": scene_max_coverage_rate, + "label_idx": scene_max_cr_idx, + }) + 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 __getitem__(self, index): + data_item_info = self.datalist[index] + 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, + ) = ([], [], []) + view = data_item_info["first_frame"] + 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) + + # combined_scanned_views_pts = np.concatenate(scanned_views_pts, axis=0) + # voxel_downsampled_combined_scanned_pts_np = PtsUtil.voxel_downsample_point_cloud(combined_scanned_views_pts, 0.002) + # random_downsampled_combined_scanned_pts_np = PtsUtil.random_downsample_point_cloud(voxel_downsampled_combined_scanned_pts_np, self.pts_num) + + data_item = { + "first_scanned_pts": np.asarray(scanned_views_pts, dtype=np.float32), # Ndarray(S x Nv x 3) + "first_scanned_coverage_rate": scanned_coverages_rate, # List(S): Float, range(0, 1) + "first_scanned_n_to_world_pose_9d": np.asarray(scanned_n_to_world_pose, dtype=np.float32), # Ndarray(S x 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) + + +# -------------- Debug ---------------- # +if __name__ == "__main__": + import torch + + seed = 0 + torch.manual_seed(seed) + np.random.seed(seed) + config = { + "root_dir": "/data/hofee/data/new_full_data", + "source": "seq_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 = SeqReconstructionDataset(config) + print(len(ds)) + print(ds.__getitem__(10)) +