diff --git a/configs/generate_config.yaml b/configs/generate_config.yaml index 70a808a..0ff838b 100644 --- a/configs/generate_config.yaml +++ b/configs/generate_config.yaml @@ -1,23 +1,24 @@ -runners: +runner: general: seed: 0 device: cpu cuda_visible_devices: "0,1,2,3,4,5,6,7" + generate: + voxel_threshold: 0.005 + overlap_threshold: 0.3 + experiment: name: debug root_dir: "experiments" - generate: - - name: OmniObject3d_train - component: OmniObject3d - data_type: train + dataset_list: + - OmniObject3d datasets: - general: - components: OmniObject3d: - root_dir: "C:\\Document\\Local Project\\nbv_rec\\output" + root_dir: "C:\\Document\\Local Project\\nbv_rec\\sample_dataset" + output_dir: "C:\\Document\\Local Project\\nbv_rec\\sample_output" diff --git a/core/dataset.py b/core/dataset.py new file mode 100644 index 0000000..7b377d8 --- /dev/null +++ b/core/dataset.py @@ -0,0 +1,17 @@ +from PytorchBoot.dataset import BaseDataset +import PytorchBoot.stereotype as stereotype + +@stereotype.dataset("nbv_reconstruction_dataset", comment="unfinished") +class NBVReconstructionDataset(BaseDataset): + def __init__(self, config): + super(NBVReconstructionDataset, self).__init__(config) + self.config = config + + def get_datalist(self): + pass + + def load_view(path): + pass + + def load_data_item(self, idx): + pass \ No newline at end of file diff --git a/core/evaluation.py b/core/evaluation.py new file mode 100644 index 0000000..8c6f7e0 --- /dev/null +++ b/core/evaluation.py @@ -0,0 +1,20 @@ +import torch +import PytorchBoot.stereotype as stereotype + +@stereotype.evaluation_method("delta_pose_diff", comment="unfinished") +class DeltaPoseDiff: + def __init__(self, config): + pass + + def evaluate(self, output_list, data_list): + return + + +@stereotype.evaluation_method("coverage_rate_increase",comment="unfinished") +class ConverageRateIncrease: + def __init__(self, config): + pass + + def evaluate(self, output_list, data_list): + return + \ No newline at end of file diff --git a/core/pipeline.py b/core/pipeline.py index 9ad21f6..033598d 100644 --- a/core/pipeline.py +++ b/core/pipeline.py @@ -5,7 +5,7 @@ import PytorchBoot.stereotype as stereotype from PytorchBoot.factory.component_factory import ComponentFactory from PytorchBoot.utils import Log -@stereotype.pipeline("nbv_reconstruction_pipeline") +@stereotype.pipeline("nbv_reconstruction_pipeline", comment="should be tested") class NBVReconstructionPipeline(nn.Module): def __init__(self, config): super(NBVReconstructionPipeline, self).__init__() diff --git a/modules/func_lib/__init__.py b/modules/func_lib/__init__.py index 5d3879a..d139a11 100644 --- a/modules/func_lib/__init__.py +++ b/modules/func_lib/__init__.py @@ -1,5 +1,4 @@ from modules.func_lib.samplers import ( - cond_pc_sampler, cond_ode_sampler ) from modules.func_lib.sde import ( diff --git a/runners/strategy_generator.py b/runners/strategy_generator.py index adc1888..f0b05b2 100644 --- a/runners/strategy_generator.py +++ b/runners/strategy_generator.py @@ -1,16 +1,30 @@ import os +import json from PytorchBoot.runners.runner import Runner from PytorchBoot.config import ConfigManager import PytorchBoot.stereotype as stereotype -@stereotype.runner("strategy_generator") +from utils.data_load import DataLoadUtil +from utils.reconstruction import ReconstructionUtil + +@stereotype.runner("strategy_generator", comment="unfinished") class StrategyGenerator(Runner): def __init__(self, config): super().__init__(config) self.load_experiment("generate") def run(self): - self.demo(seq=16,num=100) + dataset_name_list = ConfigManager.get("runner", "dataset_list") + voxel_threshold, overlap_threshold = ConfigManager.get("runner","generate","voxel_threshold"), ConfigManager.get("runner","generate","overlap_threshold") + for dataset_name in dataset_name_list: + root_dir = ConfigManager.get("datasets", dataset_name, "root_dir") + output_dir = ConfigManager.get("datasets", dataset_name, "output_dir") + if not os.path.exists(output_dir): + os.makedirs(output_dir) + + scene_idx_list = DataLoadUtil.get_scene_idx_list(root_dir) + for scene_idx in scene_idx_list: + self.generate_sequence(root_dir, output_dir, scene_idx,voxel_threshold, overlap_threshold) def create_experiment(self, backup_name=None): super().create_experiment(backup_name) @@ -20,54 +34,34 @@ class StrategyGenerator(Runner): def load_experiment(self, backup_name=None): super().load_experiment(backup_name) - def demo(self, seq, num=100): - import os - from utils.data_load import DataLoadUtil - from utils.reconstruction import ReconstructionUtil - import numpy as np - - component = self.config["generate"][0]["component"] #r"C:\Document\Local Project\nbv_rec\output" - data_dir = ConfigManager.get("datasets", "components", component, "root_dir") - model_path = os.path.join(data_dir, f"sequence.{seq}\\world_points.txt") - model_pts = np.loadtxt(model_path) - - output_dir = os.path.join(str(self.experiment_path), "output") + def generate_sequence(self,root, output_dir, seq, voxel_threshold, overlap_threshold): + frame_idx_list = DataLoadUtil.get_frame_idx_list(root, seq) + model_pts = DataLoadUtil.load_model_points(root, seq) pts_list = [] - for idx in range(0,num): - path = DataLoadUtil.get_path(data_dir, seq, idx) + for frame_idx in frame_idx_list: + path = DataLoadUtil.get_path(root, seq, frame_idx) point_cloud = DataLoadUtil.get_point_cloud_world_from_path(path) - - sampled_point_cloud = ReconstructionUtil.downsample_point_cloud(point_cloud, 0.005) + sampled_point_cloud = ReconstructionUtil.downsample_point_cloud(point_cloud, voxel_threshold) pts_list.append(sampled_point_cloud) - - sampled_model_pts = ReconstructionUtil.downsample_point_cloud(model_pts, 0.005) - np.savetxt(os.path.join(output_dir,"sampled_model_points.txt"), sampled_model_pts) - thre = 0.005 - - useful_view, useless_view = ReconstructionUtil.compute_next_best_view_sequence(model_pts, pts_list, threshold=thre) - print("useful:", useful_view) - print("useless:", useless_view) - - selected_full_views = ReconstructionUtil.combine_point_with_view_sequence(pts_list, useful_view) - downsampled_selected_full_views = ReconstructionUtil.downsample_point_cloud(selected_full_views, thre) - np.savetxt(os.path.join(output_dir,"selected_full_views.txt"), downsampled_selected_full_views) - - limited_useful_view, limited_useless_view = ReconstructionUtil.compute_next_best_view_sequence_with_overlap(model_pts, pts_list, threshold=thre, overlap_threshold=0.3) - print("limited_useful:", limited_useful_view) - print("limited_useless:", limited_useless_view) - - limited_selected_full_views = ReconstructionUtil.combine_point_with_view_sequence(pts_list, limited_useful_view) - downsampled_limited_selected_full_views = ReconstructionUtil.downsample_point_cloud(limited_selected_full_views, thre) - np.savetxt(os.path.join(output_dir,"selected_full_views_limited.txt"), downsampled_limited_selected_full_views) - import json - for idx, score in limited_useful_view: - path = DataLoadUtil.get_path(data_dir, seq, idx) - point_cloud = DataLoadUtil.get_point_cloud_world_from_path(path) - print("saving useful view: ", idx, " | score: ", score) - np.savetxt(os.path.join(output_dir,f"useful_view_{idx}.txt"), point_cloud) - with open(os.path.join(output_dir,f"useful_view.json"), 'w') as f: - json.dump(limited_useful_view, f) - print("seq length: ", len(useful_view), "limited seq length: ", len(limited_useful_view)) + limited_useful_view, _ = ReconstructionUtil.compute_next_best_view_sequence_with_overlap(model_pts, pts_list, threshold=voxel_threshold, overlap_threshold=overlap_threshold) + 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] + } + output_label_path = DataLoadUtil.get_label_path(output_dir, seq) + with open(output_label_path, 'w') as f: + json.dump(seq_save_data, f) + + def generate_data_pairs(self, useful_view): + data_pairs = [] + for next_view_idx in range(len(useful_view)): + scanned_views = useful_view[:next_view_idx] + next_view = useful_view[next_view_idx] + data_pairs.append((scanned_views, next_view)) + return data_pairs + \ No newline at end of file diff --git a/utils/data_load.py b/utils/data_load.py index c48726e..f487e9c 100644 --- a/utils/data_load.py +++ b/utils/data_load.py @@ -4,6 +4,7 @@ import Imath import numpy as np import json import cv2 +import re class DataLoadUtil: @@ -12,6 +13,38 @@ class DataLoadUtil: path = os.path.join(root, f"sequence.{scene_idx}", f"step{frame_idx}") return path + @staticmethod + def get_label_path(root, scene_idx): + path = os.path.join(root, f"sequence.{scene_idx}_label.json") + return path + + @staticmethod + def get_scene_idx_list(root): + scene_dir = os.listdir(root) + scene_idx_list = [] + for scene in scene_dir: + if "sequence" in scene: + scene_idx = int(re.search(r'\d+', scene).group()) + scene_idx_list.append(scene_idx) + return scene_idx_list + + @staticmethod + def get_frame_idx_list(root, scene_idx): + scene_path = os.path.join(root, f"sequence.{scene_idx}") + view_dir = os.listdir(scene_path) + seen_frame_idx = set() + for view in view_dir: + if "step" in view: + frame_idx = int(re.search(r'\d+', view).group()) + seen_frame_idx.add(frame_idx) + return list(seen_frame_idx) + + @staticmethod + def load_model_points(root,scene_idx): + model_path = os.path.join(root, f"sequence.{scene_idx}", "world_points.txt") + model_pts = np.loadtxt(model_path) + return model_pts + @staticmethod def read_exr_depth(depth_path): file = OpenEXR.InputFile(depth_path)