From e0fb9a76170923442cdee3cf31ca06f3ad870317 Mon Sep 17 00:00:00 2001 From: hofee <64160135+GitHofee@users.noreply.github.com> Date: Mon, 2 Sep 2024 23:47:52 +0800 Subject: [PATCH] update gf_view_finder --- configs/strategy_generate_config.yaml | 6 ++-- modules/func_lib/samplers.py | 6 ++-- modules/gf_view_finder.py | 40 +++++++++++++++++++-------- runners/data_splitor.py | 10 ++++--- runners/strategy_generator.py | 23 +++++++++++++-- utils/pose.py | 6 ++-- utils/reconstruction.py | 18 ++++++++++-- 7 files changed, 78 insertions(+), 31 deletions(-) diff --git a/configs/strategy_generate_config.yaml b/configs/strategy_generate_config.yaml index ec0789f..70e4876 100644 --- a/configs/strategy_generate_config.yaml +++ b/configs/strategy_generate_config.yaml @@ -13,13 +13,13 @@ runner: generate: voxel_threshold: 0.005 overlap_threshold: 0.5 - save_points: True + save_points: False dataset_list: - OmniObject3d datasets: OmniObject3d: - model_dir: "/media/hofee/data/data/scaled_object_meshes" - root_dir: "/media/hofee/data/data/nbv_rec/sample" + model_dir: "H:\\AI\\Datasets\\scaled_object_meshes" + root_dir: "C:\\Document\\Local Project\\nbv_rec\\data\\sample" diff --git a/modules/func_lib/samplers.py b/modules/func_lib/samplers.py index b0274fa..ae5c33c 100644 --- a/modules/func_lib/samplers.py +++ b/modules/func_lib/samplers.py @@ -88,8 +88,8 @@ def cond_ode_sampler( x = mean_x num_steps = xs.shape[0] - xs = xs.reshape(batch_size * num_steps, -1) - xs = PoseUtil.normalize_rotation(xs, pose_mode) + xs = xs.reshape(batch_size*num_steps, -1) + xs[:, :-3] = PoseUtil.normalize_rotation(xs[:, :-3], pose_mode) xs = xs.reshape(num_steps, batch_size, -1) - x = PoseUtil.normalize_rotation(x, pose_mode) + x[:, :-3] = PoseUtil.normalize_rotation(x[:, :-3], pose_mode) return xs.permute(1, 0, 2), x diff --git a/modules/gf_view_finder.py b/modules/gf_view_finder.py index 21ada32..1e1cceb 100644 --- a/modules/gf_view_finder.py +++ b/modules/gf_view_finder.py @@ -2,6 +2,9 @@ import torch import torch.nn as nn import PytorchBoot.stereotype as stereotype +import sys +sys.path.append(r"C:\Document\Local Project\nbv_rec\nbv_reconstruction") + from utils.pose import PoseUtil import modules.module_lib as mlib import modules.func_lib as flib @@ -47,7 +50,7 @@ class GradientFieldViewFinder(nn.Module): ) ''' fusion tail ''' - if self.regression_head == 'Rx_Ry': + if self.regression_head == 'Rx_Ry_and_T': if self.pose_mode != 'rot_matrix': raise NotImplementedError if not self.per_point_feature: @@ -62,6 +65,12 @@ class GradientFieldViewFinder(nn.Module): self.act, zero_module(nn.Linear(256, 3)), ) + ''' tranalation regress head ''' + self.fusion_tail_trans = nn.Sequential( + nn.Linear(128 + 256 + 2048, 256), + self.act, + zero_module(nn.Linear(256, 3)), + ) else: raise NotImplementedError else: @@ -89,10 +98,11 @@ class GradientFieldViewFinder(nn.Module): total_feat = torch.cat([seq_feat, t_feat, pose_feat], dim=-1) _, std = self.marginal_prob_fn(total_feat, t) - if self.regression_head == 'Rx_Ry': + if self.regression_head == 'Rx_Ry_and_T': rot_x = self.fusion_tail_rot_x(total_feat) rot_y = self.fusion_tail_rot_y(total_feat) - out_score = torch.cat([rot_x, rot_y], dim=-1) / (std + 1e-7) # normalisation + trans = self.fusion_tail_trans(total_feat) + out_score = torch.cat([rot_x, rot_y, trans], dim=-1) / (std+1e-7) # normalisation else: raise NotImplementedError @@ -134,18 +144,24 @@ class GradientFieldViewFinder(nn.Module): ''' ----------- DEBUG -----------''' if __name__ == "__main__": - test_scene_feat = torch.rand(32, 1024).to("cuda:0") - test_target_feat = torch.rand(32, 1024).to("cuda:0") - test_pose = torch.rand(32, 6).to("cuda:0") + config = { + "regression_head": "Rx_Ry_and_T", + "per_point_feature": False, + "pose_mode": "rot_matrix", + "sde_mode": "ve", + "sampling_steps": 500, + "sample_mode": "ode" + } + test_seq_feat = torch.rand(32, 2048).to("cuda:0") + test_pose = torch.rand(32, 9).to("cuda:0") test_t = torch.rand(32, 1).to("cuda:0") - view_finder = GradientFieldViewFinder().to("cuda:0") + view_finder = GradientFieldViewFinder(config).to("cuda:0") test_data = { - 'target_feat': test_target_feat, - 'scene_feat': test_scene_feat, + 'seq_feat': test_seq_feat, 'sampled_pose': test_pose, 't': test_t } score = view_finder(test_data) - - result = view_finder.next_best_view(test_scene_feat, test_target_feat) - print(result) + print(score.shape) + res, inprocess = view_finder.next_best_view(test_seq_feat) + print(res.shape, inprocess.shape) diff --git a/runners/data_splitor.py b/runners/data_splitor.py index 2a8bb87..eb4a661 100644 --- a/runners/data_splitor.py +++ b/runners/data_splitor.py @@ -4,9 +4,9 @@ from PytorchBoot.runners.runner import Runner from PytorchBoot.config import ConfigManager from PytorchBoot.utils import Log import PytorchBoot.stereotype as stereotype +from PytorchBoot.status import status_manager - -@stereotype.runner("data_splitor", comment="unfinished") +@stereotype.runner("data_splitor") class DataSplitor(Runner): def __init__(self, config): super().__init__(config) @@ -23,15 +23,17 @@ class DataSplitor(Runner): random.shuffle(self.datapath_list) start_idx = 0 - for dataset in self.datasets: + for dataset_idx in range(len(self.datasets)): + dataset = list(self.datasets.keys())[dataset_idx] ratio = self.datasets[dataset]["ratio"] path = self.datasets[dataset]["path"] split_size = int(len(self.datapath_list) * ratio) split_files = self.datapath_list[start_idx:start_idx + split_size] start_idx += split_size self.save_split_files(path, split_files) + status_manager.set_progress("split", "data_splitor", "split dataset", dataset_idx, len(self.datasets)) Log.success(f"save {dataset} split files to {path}") - + status_manager.set_progress("split", "data_splitor", "split dataset", len(self.datasets), len(self.datasets)) def save_split_files(self, path, split_files): os.makedirs(os.path.dirname(path), exist_ok=True) with open(path, "w") as f: diff --git a/runners/strategy_generator.py b/runners/strategy_generator.py index 4318441..b1955f4 100644 --- a/runners/strategy_generator.py +++ b/runners/strategy_generator.py @@ -6,6 +6,7 @@ from PytorchBoot.runners.runner import Runner from PytorchBoot.config import ConfigManager from PytorchBoot.utils import Log import PytorchBoot.stereotype as stereotype +from PytorchBoot.status import status_manager from utils.data_load import DataLoadUtil from utils.reconstruction import ReconstructionUtil @@ -16,12 +17,19 @@ class StrategyGenerator(Runner): def __init__(self, config): super().__init__(config) self.load_experiment("generate") + self.status_info = { + "status_manager": status_manager, + "app_name": "generate", + "runner_name": "strategy_generator" + } def run(self): dataset_name_list = ConfigManager.get("runner", "generate", "dataset_list") voxel_threshold, overlap_threshold = ConfigManager.get("runner","generate","voxel_threshold"), ConfigManager.get("runner","generate","overlap_threshold") self.save_pts = ConfigManager.get("runner","generate","save_points") - for dataset_name in dataset_name_list: + for dataset_idx in range(len(dataset_name_list)): + dataset_name = dataset_name_list[dataset_idx] + status_manager.set_progress("generate", "strategy_generator", "dataset", dataset_idx, len(dataset_name_list)) root_dir = ConfigManager.get("datasets", dataset_name, "root_dir") model_dir = ConfigManager.get("datasets", dataset_name, "model_dir") scene_name_list = os.listdir(root_dir) @@ -29,8 +37,12 @@ class StrategyGenerator(Runner): total = len(scene_name_list) for scene_name in scene_name_list: Log.info(f"({dataset_name})Processing [{cnt}/{total}]: {scene_name}") + status_manager.set_progress("generate", "strategy_generator", "scene", cnt, total) self.generate_sequence(root_dir, model_dir, scene_name,voxel_threshold, overlap_threshold) + cnt += 1 + status_manager.set_progress("generate", "strategy_generator", "scene", total, total) + status_manager.set_progress("generate", "strategy_generator", "dataset", len(dataset_name_list), len(dataset_name_list)) def create_experiment(self, backup_name=None): super().create_experiment(backup_name) @@ -41,6 +53,7 @@ class StrategyGenerator(Runner): super().load_experiment(backup_name) def generate_sequence(self, root, model_dir, scene_name, voxel_threshold, overlap_threshold): + status_manager.set_status("generate", "strategy_generator", "scene", scene_name) frame_num = DataLoadUtil.get_scene_seq_length(root, scene_name) model_pts = DataLoadUtil.load_original_model_points(model_dir, scene_name) down_sampled_model_pts = PtsUtil.voxel_downsample_point_cloud(model_pts, voxel_threshold) @@ -50,7 +63,7 @@ class StrategyGenerator(Runner): for frame_idx in range(frame_num): path = DataLoadUtil.get_path(root, scene_name, frame_idx) - + status_manager.set_progress("generate", "strategy_generator", "loading frame", frame_idx, frame_num) point_cloud = DataLoadUtil.get_point_cloud_world_from_path(path) sampled_point_cloud = PtsUtil.voxel_downsample_point_cloud(point_cloud, voxel_threshold) if self.save_pts: @@ -59,13 +72,17 @@ class StrategyGenerator(Runner): os.makedirs(pts_dir) np.savetxt(os.path.join(pts_dir, f"{frame_idx}.txt"), sampled_point_cloud) pts_list.append(sampled_point_cloud) - limited_useful_view, _ = ReconstructionUtil.compute_next_best_view_sequence_with_overlap(down_sampled_transformed_model_pts, pts_list, threshold=voxel_threshold, overlap_threshold=overlap_threshold) + status_manager.set_progress("generate", "strategy_generator", "loading frame", frame_num, frame_num) + + limited_useful_view, _ = ReconstructionUtil.compute_next_best_view_sequence_with_overlap(down_sampled_transformed_model_pts, pts_list, threshold=voxel_threshold, overlap_threshold=overlap_threshold, status_info=self.status_info) data_pairs = self.generate_data_pairs(limited_useful_view) seq_save_data = { "data_pairs": data_pairs, "best_sequence": limited_useful_view, "max_coverage_rate": limited_useful_view[-1][1] } + + status_manager.set_status("generate", "strategy_generator", "max_coverage_rate", limited_useful_view[-1][1]) Log.success(f"Scene <{scene_name}> Finished, Max Coverage Rate: {limited_useful_view[-1][1]}, Best Sequence length: {len(limited_useful_view)}") output_label_path = DataLoadUtil.get_label_path(root, scene_name) diff --git a/utils/pose.py b/utils/pose.py index dd044f9..01db630 100644 --- a/utils/pose.py +++ b/utils/pose.py @@ -184,11 +184,11 @@ class PoseUtil: ], f"the rotation mode {rot_mode} is not supported!" if rot_mode == "quat_wxyz" or rot_mode == "quat_xyzw": - pose_dim = 4 + pose_dim = 7 elif rot_mode == "euler_xyz": - pose_dim = 3 - elif rot_mode == "euler_xyz_sx_cx" or rot_mode == "rot_matrix": pose_dim = 6 + elif rot_mode == "euler_xyz_sx_cx" or rot_mode == "rot_matrix": + pose_dim = 9 else: raise NotImplementedError return pose_dim diff --git a/utils/reconstruction.py b/utils/reconstruction.py index dc236b3..9a3464e 100644 --- a/utils/reconstruction.py +++ b/utils/reconstruction.py @@ -45,12 +45,12 @@ class ReconstructionUtil: @staticmethod - def compute_next_best_view_sequence_with_overlap(target_point_cloud, point_cloud_list, threshold=0.01, overlap_threshold=0.3): + def compute_next_best_view_sequence_with_overlap(target_point_cloud, point_cloud_list, threshold=0.01, overlap_threshold=0.3, status_info=None): selected_views = [] current_coverage = 0.0 remaining_views = list(range(len(point_cloud_list))) view_sequence = [] - + cnt_processed_view = 0 while remaining_views: best_view = None best_coverage_increase = -1 @@ -74,6 +74,14 @@ class ReconstructionUtil: if coverage_increase > best_coverage_increase: best_coverage_increase = coverage_increase best_view = view_index + cnt_processed_view += 1 + if status_info is not None: + + sm = status_info["status_manager"] + app_name = status_info["app_name"] + runner_name = status_info["runner_name"] + sm.set_status(app_name, runner_name, "current coverage", current_coverage) + sm.set_progress(app_name, runner_name, "processed view", cnt_processed_view, len(point_cloud_list)) if best_view is not None: if best_coverage_increase <=1e-3: @@ -87,7 +95,11 @@ class ReconstructionUtil: else: break - + if status_info is not None: + sm = status_info["status_manager"] + app_name = status_info["app_name"] + runner_name = status_info["runner_name"] + sm.set_progress(app_name, runner_name, "processed view", len(point_cloud_list), len(point_cloud_list)) return view_sequence, remaining_views \ No newline at end of file