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

finish generate_sequence.py

Browse Source
This commit is contained in:
hofee
2024-08-22 20:27:21 +08:00
parent ff3e89b17c
commit 7cd1954a25
7 changed files with 121 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files

88
runners/strategy_generator.py
View File

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