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

update gf_view_finder

Browse Source
This commit is contained in:
hofee
2024-09-02 23:47:52 +08:00
parent 2fcfcd1966
commit e0fb9a7617
7 changed files with 78 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
runners/data_splitor.py
View File

@@ -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:

23
runners/strategy_generator.py
View File

@@ -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)