hofee/nbv_reconstruction
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

update strategy and overlap rate compute method

Browse Source
This commit is contained in:
hofee
2024-08-30 16:49:21 +08:00
parent a14bdc2c55
commit 71676e2f4e
9 changed files with 123 additions and 149 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
runners/data_generator.py
View File

@@ -1,34 +0,0 @@
import os
import json
from PytorchBoot.runners.runner import Runner
from PytorchBoot.config import ConfigManager
from PytorchBoot.utils import Log
import PytorchBoot.stereotype as stereotype
@stereotype.runner("data_generator", comment="unfinished")
class DataGenerator(Runner):
def __init__(self, config):
super().__init__(config)
self.load_experiment("generate")
def run(self):
dataset_name_list = ConfigManager.get("runner", "generate" ,"dataset_list")
for dataset_name in dataset_name_list:
self.generate(dataset_name)
def generate(self, dataset_name):
dataset_config = ConfigManager.get("datasets", dataset_name)
model_dir = dataset_config["model_dir"]
output_dir = dataset_config["output_dir"]
Log.debug(model_dir)
Log.debug(output_dir)
def create_experiment(self, backup_name=None):
super().create_experiment(backup_name)
output_dir = os.path.join(str(self.experiment_path), "output")
os.makedirs(output_dir)
def load_experiment(self, backup_name=None):
super().load_experiment(backup_name)

54
runners/strategy_generator.py
View File

@@ -1,11 +1,15 @@
import os
import json
import numpy as np
from PytorchBoot.runners.runner import Runner
from PytorchBoot.config import ConfigManager
from PytorchBoot.utils import Log
import PytorchBoot.stereotype as stereotype
from utils.data_load import DataLoadUtil
from utils.reconstruction import ReconstructionUtil
from utils.pts import PtsUtil
@stereotype.runner("strategy_generator")
class StrategyGenerator(Runner):
@@ -14,17 +18,19 @@ class StrategyGenerator(Runner):
self.load_experiment("generate")
def run(self):
dataset_name_list = ConfigManager.get("runner", "generate" "dataset_list")
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:
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)
model_dir = ConfigManager.get("datasets", dataset_name, "model_dir")
scene_name_list = os.listdir(root_dir)
cnt = 0
total = len(scene_name_list)
for scene_name in scene_name_list:
Log.info(f"({dataset_name})Processing [{cnt}/{total}]: {scene_name}")
self.generate_sequence(root_dir, model_dir, scene_name,voxel_threshold, overlap_threshold)
cnt += 1
def create_experiment(self, backup_name=None):
super().create_experiment(backup_name)
@@ -34,26 +40,40 @@ class StrategyGenerator(Runner):
def load_experiment(self, backup_name=None):
super().load_experiment(backup_name)
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)
def generate_sequence(self, root, model_dir, scene_name, voxel_threshold, overlap_threshold):
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)
obj_pose = DataLoadUtil.load_target_object_pose(root, scene_name)
down_sampled_transformed_model_pts = PtsUtil.transform_point_cloud(down_sampled_model_pts, obj_pose)
pts_list = []
for frame_idx in frame_idx_list:
path = DataLoadUtil.get_path(root, seq, frame_idx)
for frame_idx in range(frame_num):
path = DataLoadUtil.get_path(root, scene_name, frame_idx)
point_cloud = DataLoadUtil.get_point_cloud_world_from_path(path)
sampled_point_cloud = ReconstructionUtil.downsample_point_cloud(point_cloud, voxel_threshold)
sampled_point_cloud = PtsUtil.voxel_downsample_point_cloud(point_cloud, voxel_threshold)
if self.save_pts:
pts_dir = os.path.join(root,scene_name, "pts")
if not os.path.exists(pts_dir):
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(model_pts, pts_list, threshold=voxel_threshold, overlap_threshold=overlap_threshold)
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)
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)
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)
with open(output_label_path, 'w') as f:
json.dump(seq_save_data, f)
DataLoadUtil.save_downsampled_world_model_points(root, scene_name, down_sampled_transformed_model_pts)
def generate_data_pairs(self, useful_view):
data_pairs = []
for next_view_idx in range(len(useful_view)):