nbv_reconstruction/runners/strategy_generator.py
2024-08-21 17:11:56 +08:00

73 lines
3.6 KiB
Python

import os
from PytorchBoot.runners.runner import Runner
from PytorchBoot.config import ConfigManager
import PytorchBoot.stereotype as stereotype
@stereotype.runner("strategy_generator")
class StrategyGenerator(Runner):
def __init__(self, config):
super().__init__(config)
self.load_experiment("generate")
def run(self):
self.demo(seq=16,num=100)
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)
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")
pts_list = []
for idx in range(0,num):
path = DataLoadUtil.get_path(data_dir, seq, idx)
point_cloud = DataLoadUtil.get_point_cloud_world_from_path(path)
sampled_point_cloud = ReconstructionUtil.downsample_point_cloud(point_cloud, 0.005)
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))