This commit is contained in:
2025-06-12 14:31:41 +08:00
parent c0c4f9b610
commit 6cb1368ffa
14 changed files with 9353 additions and 26 deletions

View File

@@ -6,6 +6,7 @@ from runners.inference_server import InferencerServer
from runners.evaluate_uncertainty_guide import EvaluateUncertaintyGuide
from runners.evaluate_pbnbv import EvaluatePBNBV
from runners.evaluate_heuristic import Heuristic
from runners.pb_global_points_inferencer import PBGlobalPointsInferencer
@PytorchBootApplication("global_points_inference")
class GlobalPointsInferenceApp:
@@ -109,3 +110,9 @@ class EvaluateHeuristicApp:
@staticmethod
def start():
Heuristic("./configs/local/heuristic_evaluation.yaml").run()
@PytorchBootApplication("pb_global_points_inference")
class PBGlobalPointsInferenceApp:
@staticmethod
def start():
PBGlobalPointsInferencer("./configs/local/global_only_inference_config.yaml").run()

View File

@@ -15,7 +15,7 @@ runner:
- OmniObject3d_test
blender_script_path: "/media/hofee/data/project/python/nbv_reconstruction/blender/data_renderer.py"
output_dir: "/media/hofee/data/project/exp/old_no_cluster_ab_global_pts_and_local_pose"
output_dir: "/media/hofee/data/project/exp/new_ab_pb_global_pts_and_local_pose"
pipeline: nbv_reconstruction_pipeline_global
voxel_size: 0.003
min_new_area: 1.0
@@ -23,8 +23,8 @@ runner:
enable_cluster: False
dataset:
OmniObject3d_test:
root_dir: "/media/hofee/data/data/new_testset_output"
model_dir: "/media/hofee/data/data/scaled_object_meshes"
root_dir: "/media/hofee/repository/final_test_set/preprocessed_dataset"
model_dir: "/media/hofee/data/data/target/target_formulated_view"
source: seq_reconstruction_dataset_preprocessed
type: test
filter_degree: 75

View File

@@ -0,0 +1,43 @@
runner:
general:
seed: 0
device: cuda
cuda_visible_devices: "0,1,2,3,4,5,6,7"
experiment:
name: train_ab_global_only
root_dir: "experiments"
epoch: -1 # -1 stands for last epoch
pipeline: nbv_reconstruction_pipeline
voxel_size: 0.003
enable_cluster: False
heuristic_method: hemisphere_random
pipeline:
nbv_reconstruction_pipeline:
modules:
pts_encoder: pointnet_encoder
seq_encoder: transformer_seq_encoder
pose_encoder: pose_encoder
view_finder: gf_view_finder
eps: 1e-5
global_scanned_feat: True
heuristic_methods:
hemisphere_random:
center: [0, 0, 0]
radius_fixed: True
fixed_radius: 0.6
min_radius: 0.4
max_radius: 0.8
hemisphere_circle_trajectory:
center: [0, 0, 0]
radius_fixed: False
fixed_radius: 0.6
min_radius: 0.4
max_radius: 0.8
phi_list: [15, 45, 75]
circle_times: 12

View File

@@ -21,7 +21,7 @@ class SeqReconstructionDatasetPreprocessed(BaseDataset):
super(SeqReconstructionDatasetPreprocessed, self).__init__(config)
self.config = config
self.root_dir = config["root_dir"]
self.real_root_dir = r"/media/hofee/data/data/new_testset"
self.real_root_dir = r"/media/hofee/repository/final_test_set/view"
self.item_list = os.listdir(self.root_dir)
def __getitem__(self, index):

Binary file not shown.

View File

@@ -24,18 +24,20 @@ class VoxelStruct:
self.empty_voxels = []
self.unknown_voxels = []
self.frontier_voxels = []
self.curr_camera_pose = None
self.bbx_min = None
self.bbx_max = None
self.voxel_types: Dict[Tuple[float, float, float], VoxelType] = {}
def update_voxel_map(self, points: np.ndarray,
camera_pose: np.ndarray) -> Tuple[List[np.ndarray], List[np.ndarray]]:
points = self.transform_points(points, camera_pose)
#points = self.transform_points(points, camera_pose)
new_occupied = self.voxelize_points(points)
self.occupied_voxels.extend(new_occupied)
self.update_bounding_box()
self.ray_tracing(camera_pose[:3, 3], camera_pose[:3, :3])
self.update_frontier_voxels()
self.curr_camera_pose = camera_pose
return self.frontier_voxels, self.occupied_voxels
def ray_tracing(self, camera_position: np.ndarray, camera_rotation: np.ndarray):
@@ -91,14 +93,10 @@ class VoxelStruct:
hemisphere_radius = self.camera_working_distance + bbx_diagonal / 2
else:
hemisphere_radius = self.camera_working_distance
theta_step = np.pi / (12 * self.num_parallels)
phi_step = np.pi / (6 * self.viewpoints_per_parallel)
# 使用更密集的采样
theta_step = np.pi / (6 * self.num_parallels) # 减小theta的步长
phi_step = np.pi / (6 * self.viewpoints_per_parallel) # 减小phi的步长
# 从顶部到底部采样
for theta in np.arange(0, np.pi/6 + theta_step, theta_step):
# 在每个纬度上采样
for phi in np.arange(0, 2*np.pi, phi_step):
x = hemisphere_radius * np.sin(theta) * np.cos(phi)
y = hemisphere_radius * np.sin(theta) * np.sin(phi)
@@ -273,7 +271,6 @@ class PBNBV:
def capture(self, point_cloud: np.ndarray, camera_pose: np.ndarray):
frontier_voxels, occupied_voxels = self.voxel_struct.update_voxel_map(point_cloud, camera_pose)
# self.voxel_struct.visualize_voxel_struct(camera_pose)
self.fit_ellipsoids(frontier_voxels, occupied_voxels)
def reset(self):
@@ -382,8 +379,8 @@ class PBNBV:
return []
center = (self.voxel_struct.bbx_min + self.voxel_struct.bbx_max) / 2
radius = np.linalg.norm(self.voxel_struct.bbx_max - self.voxel_struct.bbx_min) / 2 + self.focal_length
#import ipdb; ipdb.set_trace()
radius = np.linalg.norm(self.voxel_struct.bbx_max - self.voxel_struct.bbx_min) / 2 + 0.3
candidate_views = []
latitudes = np.linspace(np.deg2rad(40), np.deg2rad(90), longitude_num)
@@ -432,15 +429,15 @@ class PBNBV:
scores = [self.evaluate_viewpoint(view) for view in candidate_views]
best_idx = np.argmax(scores)
return candidate_views[best_idx]
return candidate_views[best_idx],candidate_views
def execute(self) -> Tuple[np.ndarray, bool]:
best_view = self.select_best_view()
best_view,candidate_views = self.select_best_view()
has_frontier = any(e["type"] == "frontier" for e in self.ellipsoids)
done = not has_frontier
return best_view, done
return best_view, candidate_views, done
import os
import json
@@ -478,8 +475,8 @@ class EvaluatePBNBV(Runner):
CM = 0.01
self.min_new_pts_num = self.min_new_area * (CM / self.voxel_size) ** 2
self.overlap_limit = ConfigManager.get(namespace.Stereotype.RUNNER, "overlap_limit")
self.pbnbv = PBNBV()
self.pbnbv = PBNBV(self.voxel_size)
''' Experiment '''
self.load_experiment("nbv_evaluator")
self.stat_result_path = os.path.join(self.output_dir, "stat.json")
@@ -541,11 +538,49 @@ class EvaluatePBNBV(Runner):
status_manager.set_progress("inference", "inferencer", f"dataset", len(self.test_set_list), len(self.test_set_list))
def get_output_data(self):
pose_matrix, done = self.pbnbv.execute()
def visualize(self, best_view, candidate_views):
import plotly.graph_objects as go
fig = go.Figure()
colors = ['aggrnyl', 'agsunset', 'algae', 'amp', 'armyrose', 'balance',
'blackbody', 'bluered', 'blues', 'blugrn', 'bluyl', 'brbg']
color = colors[0]
candidate_pose = best_view
origin_candidate = candidate_pose[:3, 3]
z_axis_candidate = candidate_pose[:3, 2]
fig.add_trace(go.Cone(
x=[origin_candidate[0]], y=[origin_candidate[1]], z=[origin_candidate[2]],
u=[z_axis_candidate[0]], v=[z_axis_candidate[1]], w=[z_axis_candidate[2]],
colorscale=color,
sizemode="absolute", sizeref=0.1, anchor="tail", showscale=False
))
for candidate_pose in candidate_views:
origin_candidate = candidate_pose[:3, 3]
z_axis_candidate = candidate_pose[:3, 2]
color = colors[1]
fig.add_trace(go.Cone(
x=[origin_candidate[0]], y=[origin_candidate[1]], z=[origin_candidate[2]],
u=[z_axis_candidate[0]], v=[z_axis_candidate[1]], w=[z_axis_candidate[2]],
colorscale=color,
sizemode="absolute", sizeref=0.05, anchor="tail", showscale=False
))
fig.update_layout(
title="Clustered Poses and Input Points",
scene=dict(
xaxis_title='X',
yaxis_title='Y',
zaxis_title='Z'
),
margin=dict(l=0, r=0, b=0, t=40),
scene_camera=dict(eye=dict(x=1.25, y=1.25, z=1.25))
)
fig.show()
def get_output_data(self):
pose_matrix,candidate_views, done = self.pbnbv.execute()
offset = np.asarray([[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]])
pose_matrix = pose_matrix @ offset
rot = pose_matrix[:3,:3]
pose_6d = PoseUtil.matrix_to_rotation_6d_numpy(rot)
@@ -554,7 +589,9 @@ class EvaluatePBNBV(Runner):
pose_9d = np.concatenate([pose_6d, translation], axis=0).reshape(1,9)
pose_9d = pose_9d.repeat(50, axis=0)
#import ipdb; ipdb.set_trace()
return {"pred_pose_9d": pose_9d}
#self.visualize(pose_matrix,candidate_views)
return {"pred_pose_9d": pose_9d, "candidate_views": candidate_views}
def predict_sequence(self, data, cr_increase_threshold=0, overlap_area_threshold=25, scan_points_threshold=10, max_iter=50, max_retry = 10, max_success=3):
scene_name = data["scene_name"]
@@ -611,8 +648,8 @@ class EvaluatePBNBV(Runner):
predict_result = PredictResult(pred_pose_9d, input_pts=input_data["combined_scanned_pts"], cluster_params=dict(eps=0.25, min_samples=3))
# -----------------------
import ipdb; ipdb.set_trace()
predict_result.visualize()
#import ipdb; ipdb.set_trace()
#predict_result.visualize()
# -----------------------
pred_pose_9d_candidates = predict_result.candidate_9d_poses
#import ipdb; ipdb.set_trace()

View File

@@ -0,0 +1,216 @@
import os
import json
import torch
import numpy as np
from flask import Flask, request, jsonify
from PytorchBoot.config import ConfigManager
import PytorchBoot.namespace as namespace
import PytorchBoot.stereotype as stereotype
from PytorchBoot.factory import ComponentFactory
from PytorchBoot.runners.runner import Runner
from PytorchBoot.utils import Log
from utils.pts import PtsUtil
from beans.predict_result import PredictResult
@stereotype.runner("heuristic_inferencer_server")
class HeuristicInferencerServer(Runner):
def __init__(self, config_path):
super().__init__(config_path)
self.heuristic_method = ConfigManager.get(namespace.Stereotype.RUNNER, "heuristic_method")
self.heuristic_method_config = ConfigManager.get("heuristic_methods", self.heuristic_method)
''' Web Server '''
self.app = Flask(__name__)
''' Pipeline '''
self.pipeline_name = self.config[namespace.Stereotype.PIPELINE]
self.pipeline:torch.nn.Module = ComponentFactory.create(namespace.Stereotype.PIPELINE, self.pipeline_name)
self.pipeline = self.pipeline.to(self.device)
self.pts_num = 8192
self.voxel_size = 0.002
''' Experiment '''
self.load_experiment("inferencer_server")
def generate_hemisphere_random_sequence(self, max_iter, config):
"""Generate a random hemisphere sampling sequence"""
radius_fixed = config["radius_fixed"]
fixed_radius = config["fixed_radius"]
min_radius = config["min_radius"]
max_radius = config["max_radius"]
poses = []
center = np.array(config["center"])
for _ in range(max_iter):
# 随机采样方向
direction = np.random.randn(3)
direction[2] = abs(direction[2]) # 确保在上半球
direction = direction / np.linalg.norm(direction)
# 确定半径
if radius_fixed:
radius = fixed_radius
else:
radius = np.random.uniform(min_radius, max_radius)
# 计算位置和朝向
position = center + direction * radius
z_axis = -direction
y_axis = np.array([0, 0, 1])
x_axis = np.cross(y_axis, z_axis)
x_axis = x_axis / np.linalg.norm(x_axis)
y_axis = np.cross(z_axis, x_axis)
pose = np.eye(4)
pose[:3,:3] = np.stack([x_axis, y_axis, z_axis], axis=1)
pose[:3,3] = position
poses.append(pose)
return poses
def generate_hemisphere_circle_sequence(self, config):
"""Generate a circular trajectory sampling sequence"""
radius_fixed = config["radius_fixed"]
fixed_radius = config["fixed_radius"]
min_radius = config["min_radius"]
max_radius = config["max_radius"]
phi_list = config["phi_list"]
circle_times = config["circle_times"]
poses = []
center = np.array(config["center"])
for phi in phi_list: # 仰角
phi_rad = np.deg2rad(phi)
for i in range(circle_times): # 方位角
theta = i * (2 * np.pi / circle_times)
# 确定半径
if radius_fixed:
radius = fixed_radius
else:
radius = np.random.uniform(min_radius, max_radius)
# 球坐标转笛卡尔坐标
x = radius * np.cos(theta) * np.sin(phi_rad)
y = radius * np.sin(theta) * np.sin(phi_rad)
z = radius * np.cos(phi_rad)
position = center + np.array([x, y, z])
# 计算朝向
direction = (center - position) / np.linalg.norm(center - position)
z_axis = direction
y_axis = np.array([0, 0, 1])
x_axis = np.cross(y_axis, z_axis)
x_axis = x_axis / np.linalg.norm(x_axis)
y_axis = np.cross(z_axis, x_axis)
pose = np.eye(4)
pose[:3,:3] = np.stack([x_axis, y_axis, z_axis], axis=1)
pose[:3,3] = position
poses.append(pose)
return poses
def generate_seq(self, max_iter=50):
if self.heuristic_method == "hemisphere_random":
pose_sequence = self.generate_hemisphere_random_sequence(
max_iter,
self.heuristic_method_config
)
elif self.heuristic_method == "hemisphere_circle_trajectory":
pose_sequence = self.generate_hemisphere_circle_sequence(
self.heuristic_method_config
)
else:
raise ValueError(f"Unknown heuristic method: {self.heuristic_method}")
return pose_sequence
def get_input_data(self, data):
input_data = {}
scanned_pts = data["scanned_pts"]
scanned_n_to_world_pose_9d = data["scanned_n_to_world_pose_9d"]
combined_scanned_views_pts = np.concatenate(scanned_pts, axis=0)
voxel_downsampled_combined_scanned_pts = PtsUtil.voxel_downsample_point_cloud(
combined_scanned_views_pts, self.voxel_size
)
fps_downsampled_combined_scanned_pts, fps_idx = PtsUtil.fps_downsample_point_cloud(
voxel_downsampled_combined_scanned_pts, self.pts_num, require_idx=True
)
input_data["scanned_pts"] = scanned_pts
input_data["scanned_n_to_world_pose_9d"] = np.asarray(scanned_n_to_world_pose_9d, dtype=np.float32)
input_data["combined_scanned_pts"] = np.asarray(fps_downsampled_combined_scanned_pts, dtype=np.float32)
return input_data
def get_result(self, output_data):
pred_pose_9d = output_data["pred_pose_9d"]
pred_pose_9d = np.asarray(PredictResult(pred_pose_9d.cpu().numpy(), None, cluster_params=dict(eps=0.25, min_samples=3)).candidate_9d_poses, dtype=np.float32)
result = {
"pred_pose_9d": pred_pose_9d.tolist()
}
return result
def collate_input(self, input_data):
collated_input_data = {}
collated_input_data["scanned_pts"] = [torch.tensor(input_data["scanned_pts"], dtype=torch.float32, device=self.device)]
collated_input_data["scanned_n_to_world_pose_9d"] = [torch.tensor(input_data["scanned_n_to_world_pose_9d"], dtype=torch.float32, device=self.device)]
collated_input_data["combined_scanned_pts"] = torch.tensor(input_data["combined_scanned_pts"], dtype=torch.float32, device=self.device).unsqueeze(0)
return collated_input_data
def do_inference(self, input_data):
scanned_pts = input_data["scanned_pts"]
def run(self):
Log.info("Loading from epoch {}.".format(self.current_epoch))
@self.app.route("/inference", methods=["POST"])
def inference():
data = request.json
input_data = self.get_input_data(data)
collated_input_data = self.collate_input(input_data)
output_data = self.do_inference(collated_input_data)
result = self.get_result(output_data)
return jsonify(result)
self.app.run(host="0.0.0.0", port=5000)
def get_checkpoint_path(self, is_last=False):
return os.path.join(self.experiment_path, namespace.Direcotry.CHECKPOINT_DIR_NAME,
"Epoch_{}.pth".format(
self.current_epoch if self.current_epoch != -1 and not is_last else "last"))
def load_checkpoint(self, is_last=False):
self.load(self.get_checkpoint_path(is_last))
Log.success(f"Loaded checkpoint from {self.get_checkpoint_path(is_last)}")
if is_last:
checkpoint_root = os.path.join(self.experiment_path, namespace.Direcotry.CHECKPOINT_DIR_NAME)
meta_path = os.path.join(checkpoint_root, "meta.json")
if not os.path.exists(meta_path):
raise FileNotFoundError(
"No checkpoint meta.json file in the experiment {}".format(self.experiments_config["name"]))
file_path = os.path.join(checkpoint_root, "meta.json")
with open(file_path, "r") as f:
meta = json.load(f)
self.current_epoch = meta["last_epoch"]
self.current_iter = meta["last_iter"]
def load_experiment(self, backup_name=None):
super().load_experiment(backup_name)
self.current_epoch = self.experiments_config["epoch"]
self.load_checkpoint(is_last=(self.current_epoch == -1))
def create_experiment(self, backup_name=None):
super().create_experiment(backup_name)
def load(self, path):
state_dict = torch.load(path)
self.pipeline.load_state_dict(state_dict)

View File

@@ -64,6 +64,9 @@ class InferencerServer(Runner):
collated_input_data["combined_scanned_pts"] = torch.tensor(input_data["combined_scanned_pts"], dtype=torch.float32, device=self.device).unsqueeze(0)
return collated_input_data
def do_inference(self, input_data):
return self.pipeline.forward_test(input_data)
def run(self):
Log.info("Loading from epoch {}.".format(self.current_epoch))
@@ -72,7 +75,7 @@ class InferencerServer(Runner):
data = request.json
input_data = self.get_input_data(data)
collated_input_data = self.collate_input(input_data)
output_data = self.pipeline.forward_test(collated_input_data)
output_data = self.do_inference(collated_input_data)
result = self.get_result(output_data)
return jsonify(result)

View File

@@ -0,0 +1,829 @@
import os
import json
from utils.render import RenderUtil
from utils.pose import PoseUtil
from utils.pts import PtsUtil
from utils.reconstruction import ReconstructionUtil
from beans.predict_result import PredictResult
import torch
from tqdm import tqdm
import numpy as np
import pickle
from PytorchBoot.config import ConfigManager
import PytorchBoot.namespace as namespace
import PytorchBoot.stereotype as stereotype
from PytorchBoot.factory import ComponentFactory
from PytorchBoot.dataset import BaseDataset
from PytorchBoot.runners.runner import Runner
from PytorchBoot.utils import Log
from PytorchBoot.status import status_manager
from utils.data_load import DataLoadUtil
from sklearn.mixture import GaussianMixture
from typing import List, Tuple, Dict
from enum import Enum
class VoxelType(Enum):
NONE = 0
OCCUPIED = 1
EMPTY = 2
UNKNOWN = 3
FRONTIER = 4
class VoxelStruct:
def __init__(self, voxel_resolution=0.01, ray_trace_step=0.01, surrounding_radius=1,
num_parallels=10, viewpoints_per_parallel=10, camera_working_distance=0.5):
self.voxel_resolution = voxel_resolution
self.ray_trace_step = ray_trace_step
self.surrounding_radius = surrounding_radius
self.num_parallels = num_parallels
self.viewpoints_per_parallel = viewpoints_per_parallel
self.camera_working_distance = camera_working_distance
self.occupied_voxels = []
self.empty_voxels = []
self.unknown_voxels = []
self.frontier_voxels = []
self.curr_camera_pose = None
self.bbx_min = None
self.bbx_max = None
self.voxel_types: Dict[Tuple[float, float, float], VoxelType] = {}
def update_voxel_map(self, points: np.ndarray,
camera_pose: np.ndarray) -> Tuple[List[np.ndarray], List[np.ndarray]]:
#points = self.transform_points(points, camera_pose)
new_occupied = self.voxelize_points(points)
self.occupied_voxels.extend(new_occupied)
self.update_bounding_box()
self.ray_tracing(camera_pose[:3, 3], camera_pose[:3, :3])
self.update_frontier_voxels()
self.curr_camera_pose = camera_pose
return self.frontier_voxels, self.occupied_voxels
def ray_tracing(self, camera_position: np.ndarray, camera_rotation: np.ndarray):
if self.bbx_min is None or self.bbx_max is None:
return
directions = self.generate_ray_directions()
for direction in directions:
direction_cam = camera_rotation @ direction
current_pos = camera_position.copy()
cnt = 0
while not self.is_in_bounding_box(current_pos):
current_pos -= direction_cam * self.ray_trace_step*2
cnt += 1
if cnt > 200:
break
occupied_flag = False
maybe_unknown_voxels = []
while self.is_in_bounding_box(current_pos):
voxel = self.get_voxel_coordinate(current_pos)
voxel_key = tuple(voxel)
if self.is_occupied(voxel):
current_pos -= direction_cam * self.ray_trace_step
occupied_flag = True
continue
if not occupied_flag:
if voxel_key not in self.voxel_types or self.voxel_types[voxel_key] == VoxelType.NONE or self.voxel_types[voxel_key] == VoxelType.UNKNOWN:
maybe_unknown_voxels.append(voxel)
else:
if voxel_key not in self.voxel_types or self.voxel_types[voxel_key] == VoxelType.NONE:
self.voxel_types[voxel_key] = VoxelType.UNKNOWN
self.unknown_voxels.append(voxel)
current_pos -= direction_cam * self.ray_trace_step
if not occupied_flag:
for voxel in maybe_unknown_voxels:
self.voxel_types[tuple(voxel)] = VoxelType.UNKNOWN
self.unknown_voxels.append(voxel)
else:
for voxel in maybe_unknown_voxels:
voxel_key = tuple(voxel)
if voxel_key in self.voxel_types and self.voxel_types[voxel_key] == VoxelType.UNKNOWN:
self.unknown_voxels = [v for v in self.unknown_voxels if not np.array_equal(v, voxel)]
self.voxel_types[voxel_key] = VoxelType.EMPTY
self.empty_voxels.append(voxel)
def generate_ray_directions(self):
directions = []
if self.bbx_min is not None and self.bbx_max is not None:
bbx_diagonal = np.linalg.norm(self.bbx_max - self.bbx_min)
hemisphere_radius = self.camera_working_distance + bbx_diagonal / 2
else:
hemisphere_radius = self.camera_working_distance
theta_step = np.pi / (6 * self.num_parallels)
phi_step = np.pi / (6 * self.viewpoints_per_parallel)
for theta in np.arange(0, np.pi/6 + theta_step, theta_step):
for phi in np.arange(0, 2*np.pi, phi_step):
x = hemisphere_radius * np.sin(theta) * np.cos(phi)
y = hemisphere_radius * np.sin(theta) * np.sin(phi)
z = hemisphere_radius * np.cos(theta)
direction = np.array([-x, -y, -z])
direction = direction / np.linalg.norm(direction)
directions.append(direction)
return directions
def update_frontier_voxels(self):
self.frontier_voxels = []
remaining_unknown = []
for voxel in self.unknown_voxels:
neighbors = self.find_neighbors(voxel)
has_empty = any(self.voxel_types.get(tuple(n), VoxelType.NONE) == VoxelType.EMPTY for n in neighbors)
has_occupied = any(self.voxel_types.get(tuple(n), VoxelType.NONE) == VoxelType.OCCUPIED for n in neighbors)
if has_empty and has_occupied:
self.voxel_types[tuple(voxel)] = VoxelType.FRONTIER
self.frontier_voxels.append(voxel)
else:
remaining_unknown.append(voxel)
self.unknown_voxels = remaining_unknown
def is_in_bounding_box(self, point: np.ndarray) -> bool:
if self.bbx_min is None or self.bbx_max is None:
return False
return np.all(point >= self.bbx_min) and np.all(point <= self.bbx_max)
def get_voxel_coordinate(self, point: np.ndarray) -> np.ndarray:
return (point / self.voxel_resolution).astype(int) * self.voxel_resolution
def voxelize_points(self, points: np.ndarray) -> List[np.ndarray]:
voxel_coords = (points / self.voxel_resolution).astype(int)
unique_voxels = np.unique(voxel_coords, axis=0)
voxels = [voxel * self.voxel_resolution for voxel in unique_voxels]
for voxel in voxels:
self.voxel_types[tuple(voxel)] = VoxelType.OCCUPIED
return voxels
def is_occupied(self, voxel: np.ndarray) -> bool:
return self.voxel_types.get(tuple(voxel), VoxelType.NONE) == VoxelType.OCCUPIED
def find_neighbors(self, voxel: np.ndarray) -> List[np.ndarray]:
neighbors = []
for dx in [-1, 0, 1]:
for dy in [-1, 0, 1]:
for dz in [-1, 0, 1]:
if dx == 0 and dy == 0 and dz == 0:
continue
neighbor = voxel + np.array([dx, dy, dz]) * self.voxel_resolution
neighbors.append(neighbor)
return neighbors
def update_bounding_box(self):
if not self.occupied_voxels:
return
occupied_array = np.array(self.occupied_voxels)
self.bbx_min = occupied_array.min(axis=0) - 2 * self.voxel_resolution
self.bbx_max = occupied_array.max(axis=0) + 2 * self.voxel_resolution
def transform_points(self, points: np.ndarray, transform: np.ndarray) -> np.ndarray:
ones = np.ones((points.shape[0], 1))
points_homo = np.hstack((points, ones))
transformed = (transform @ points_homo.T).T
return transformed[:, :3]
def create_voxel_geometry(self,voxels, color, voxel_size):
import open3d as o3d
points = np.array(voxels)
if len(points) == 0:
return None
pcd = o3d.geometry.PointCloud()
pcd.points = o3d.utility.Vector3dVector(points)
pcd.colors = o3d.utility.Vector3dVector(np.tile(color, (len(points), 1)))
return pcd
def create_ray_geometry(self,camera_pos, directions, camera_rot, length=1.0):
import open3d as o3d
lines = []
colors = []
for direction in directions:
# 将方向向量转换到相机坐标系
direction_cam = camera_rot @ direction
end_point = camera_pos - direction_cam * length
lines.append([camera_pos, end_point])
colors.append([0.5, 0.5, 0.5]) # 灰色光线
line_set = o3d.geometry.LineSet()
line_set.points = o3d.utility.Vector3dVector(np.array(lines).reshape(-1, 3))
line_set.lines = o3d.utility.Vector2iVector(np.array([[i*2, i*2+1] for i in range(len(lines))]))
line_set.colors = o3d.utility.Vector3dVector(colors)
return line_set
def visualize_voxel_struct(self, camera_pose: np.ndarray = None):
if camera_pose is None:
camera_pose = self.curr_camera_pose
import open3d as o3d
vis = o3d.visualization.Visualizer()
vis.create_window()
coordinate_frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=0.1, origin=[0, 0, 0])
vis.add_geometry(coordinate_frame)
# 显示已占据的体素(蓝色)
occupied_voxels = self.create_voxel_geometry(
self.occupied_voxels,
[0, 0, 1],
self.voxel_resolution
)
if occupied_voxels:
vis.add_geometry(occupied_voxels)
# 显示空体素(绿色)
empty_voxels = self.create_voxel_geometry(
self.empty_voxels,
[0, 1, 0],
self.voxel_resolution
)
if empty_voxels:
vis.add_geometry(empty_voxels)
# 显示未知体素(灰色)
unknown_voxels = self.create_voxel_geometry(
self.unknown_voxels,
[0.5, 0.5, 0.5],
self.voxel_resolution
)
if unknown_voxels:
vis.add_geometry(unknown_voxels)
# 显示frontier体素红色
frontier_voxels = self.create_voxel_geometry(
self.frontier_voxels,
[1, 0, 0],
self.voxel_resolution
)
if frontier_voxels:
vis.add_geometry(frontier_voxels)
# 显示光线
if camera_pose is not None:
directions = self.generate_ray_directions()
rays = self.create_ray_geometry(
camera_pose[:3, 3],
directions,
camera_pose[:3, :3],
length=0.5 # 光线长度
)
vis.add_geometry(rays)
opt = vis.get_render_option()
opt.background_color = np.asarray([0.8, 0.8, 0.8])
opt.point_size = 5.0
vis.run()
vis.destroy_window()
class PBNBV:
def __init__(self, voxel_resolution=0.01, camera_intrinsic=None):
self.voxel_resolution = voxel_resolution
self.voxel_struct = VoxelStruct(voxel_resolution)
self.camera_intrinsic = camera_intrinsic or np.array([
[902.14, 0, 320],
[0, 902.14, 200],
[0, 0, 1]
])
self.focal_length = (self.camera_intrinsic[0,0] + self.camera_intrinsic[1,1]) / 2 / 1000
self.ellipsoids = []
def capture(self, point_cloud: np.ndarray, camera_pose: np.ndarray):
# ------ Debug Start ------
#import ipdb;ipdb.set_trace()
# ------ Debug End ------
frontier_voxels, occupied_voxels = self.voxel_struct.update_voxel_map(point_cloud, camera_pose)
self.fit_ellipsoids(frontier_voxels, occupied_voxels)
def reset(self):
self.ellipsoids = []
self.voxel_struct = VoxelStruct(self.voxel_resolution)
def fit_ellipsoids(self, frontier_voxels: List[np.ndarray], occupied_voxels: List[np.ndarray],
max_ellipsoids=10):
self.ellipsoids = []
if not frontier_voxels and not occupied_voxels:
return
if frontier_voxels:
frontier_gmm = self.fit_gmm(np.array(frontier_voxels), max_ellipsoids)
self.ellipsoids.extend(self.gmm_to_ellipsoids(frontier_gmm, "frontier"))
if occupied_voxels:
occupied_gmm = self.fit_gmm(np.array(occupied_voxels), max_ellipsoids)
self.ellipsoids.extend(self.gmm_to_ellipsoids(occupied_gmm, "occupied"))
def fit_gmm(self, data: np.ndarray, max_components: int) -> GaussianMixture:
best_gmm = None
best_bic = np.inf
for n in range(1, min(max_components, len(data)) + 1):
gmm = GaussianMixture(n_components=n, covariance_type='full')
gmm.fit(data)
bic = gmm.bic(data)
if bic < best_bic:
best_bic = bic
best_gmm = gmm
return best_gmm
def gmm_to_ellipsoids(self, gmm: GaussianMixture, ellipsoid_type: str) -> List[Dict]:
ellipsoids = []
for i in range(gmm.n_components):
mean = gmm.means_[i]
cov = gmm.covariances_[i]
eigvals, eigvecs = np.linalg.eigh(cov)
radii = np.sqrt(eigvals) * 3
rotation = eigvecs
pose = np.eye(4)
pose[:3, :3] = rotation
pose[:3, 3] = mean
ellipsoids.append({
"type": ellipsoid_type,
"pose": pose,
"radii": radii
})
return ellipsoids
def evaluate_viewpoint(self, viewpoint_pose: np.ndarray) -> float:
if not self.ellipsoids:
return 0.0
ellipsoid_weights = self.compute_ellipsoid_weights(viewpoint_pose)
projection_scores = []
for ellipsoid, weight in zip(self.ellipsoids, ellipsoid_weights):
score = self.project_ellipsoid(ellipsoid, viewpoint_pose) * weight
projection_scores.append((ellipsoid["type"], score))
frontier_score = sum(s for t, s in projection_scores if t == "frontier")
occupied_score = sum(s for t, s in projection_scores if t == "occupied")
return frontier_score - occupied_score
def compute_ellipsoid_weights(self, viewpoint_pose: np.ndarray) -> List[float]:
centers_world = np.array([e["pose"][:3, 3] for e in self.ellipsoids])
centers_homo = np.hstack((centers_world, np.ones((len(centers_world), 1))))
centers_cam = (np.linalg.inv(viewpoint_pose) @ centers_homo.T).T[:, :3]
z_coords = centers_cam[:, 2]
sorted_indices = np.argsort(z_coords)
weights = np.zeros(len(self.ellipsoids))
for rank, idx in enumerate(sorted_indices):
weights[idx] = 0.5 ** rank
return weights.tolist()
def project_ellipsoid(self, ellipsoid: Dict, viewpoint_pose: np.ndarray) -> float:
ellipsoid_pose_cam = np.linalg.inv(viewpoint_pose) @ ellipsoid["pose"]
radii = ellipsoid["radii"]
rotation = ellipsoid_pose_cam[:3, :3]
scales = np.diag(radii)
transform = rotation @ scales
major_axis = np.linalg.norm(transform[:, 0])
minor_axis = np.linalg.norm(transform[:, 1])
area = np.pi * major_axis * minor_axis
return area
def generate_candidate_views(self, num_views=100, longitude_num=5) -> List[np.ndarray]:
if self.voxel_struct.bbx_min is None:
return []
center = (self.voxel_struct.bbx_min + self.voxel_struct.bbx_max) / 2
#import ipdb; ipdb.set_trace()
radius = np.linalg.norm(self.voxel_struct.bbx_max - self.voxel_struct.bbx_min) / 2 + 0.3
candidate_views = []
latitudes = np.linspace(np.deg2rad(40), np.deg2rad(90), longitude_num)
lengths = [2 * np.pi * np.sin(lat) * radius for lat in latitudes]
total_length = sum(lengths)
points_per_lat = [int(round(num_views * l / total_length)) for l in lengths]
for lat, n in zip(latitudes, points_per_lat):
if n == 0:
continue
longitudes = np.linspace(0, 2*np.pi, n, endpoint=False)
for lon in longitudes:
x = radius * np.sin(lat) * np.cos(lon)
y = radius * np.sin(lat) * np.sin(lon)
z = radius * np.cos(lat)
position = np.array([x, y, z]) + center
z_axis = center - position
z_axis /= np.linalg.norm(z_axis)
x_axis = np.cross(z_axis, np.array([0, 0, 1]))
if np.linalg.norm(x_axis) < 1e-6:
x_axis = np.array([1, 0, 0])
x_axis /= np.linalg.norm(x_axis)
y_axis = np.cross(z_axis, x_axis)
y_axis /= np.linalg.norm(y_axis)
rotation = np.column_stack((x_axis, y_axis, z_axis))
view_pose = np.eye(4)
view_pose[:3, :3] = rotation
view_pose[:3, 3] = position
candidate_views.append(view_pose)
return candidate_views
def select_best_view(self, pred_pose_mat: np.ndarray) -> np.ndarray:
# ------ Debug Start ------
# import ipdb;ipdb.set_trace()
# ------ Debug End ------
candidate_views = pred_pose_mat
scores = [self.evaluate_viewpoint(view) for view in candidate_views]
best_idx = np.argmax(scores)
return candidate_views[best_idx],candidate_views
def execute(self, pred_pose_mat: np.ndarray) -> Tuple[np.ndarray, bool]:
best_view,candidate_views = self.select_best_view(pred_pose_mat)
has_frontier = any(e["type"] == "frontier" for e in self.ellipsoids)
done = not has_frontier
return best_view, candidate_views, done
def visualize(self):
self.voxel_struct.visualize_voxel_struct(self.voxel_struct.curr_camera_pose)
@stereotype.runner("pb_global_points_inferencer")
class PBGlobalPointsInferencer(Runner):
def __init__(self, config_path):
super().__init__(config_path)
self.script_path = ConfigManager.get(namespace.Stereotype.RUNNER, "blender_script_path")
self.output_dir = ConfigManager.get(namespace.Stereotype.RUNNER, "output_dir")
self.voxel_size = ConfigManager.get(namespace.Stereotype.RUNNER, "voxel_size")
self.min_new_area = ConfigManager.get(namespace.Stereotype.RUNNER, "min_new_area")
CM = 0.01
self.min_new_pts_num = self.min_new_area * (CM / self.voxel_size) **2
self.overlap_limit = ConfigManager.get(namespace.Stereotype.RUNNER, "overlap_limit")
self.enable_cluster = ConfigManager.get(namespace.Stereotype.RUNNER, "enable_cluster")
self.pbnbv = PBNBV()
''' Pipeline '''
self.pipeline_name = self.config[namespace.Stereotype.PIPELINE]
self.pipeline:torch.nn.Module = ComponentFactory.create(namespace.Stereotype.PIPELINE, self.pipeline_name)
self.pipeline = self.pipeline.to(self.device)
''' Experiment '''
self.load_experiment("nbv_evaluator")
self.stat_result_path = os.path.join(self.output_dir, "stat.json")
if os.path.exists(self.stat_result_path):
with open(self.stat_result_path, "r") as f:
self.stat_result = json.load(f)
else:
self.stat_result = {}
''' Test '''
self.test_config = ConfigManager.get(namespace.Stereotype.RUNNER, namespace.Mode.TEST)
self.test_dataset_name_list = self.test_config["dataset_list"]
self.test_set_list = []
self.test_writer_list = []
seen_name = set()
for test_dataset_name in self.test_dataset_name_list:
if test_dataset_name not in seen_name:
seen_name.add(test_dataset_name)
else:
raise ValueError("Duplicate test dataset name: {}".format(test_dataset_name))
test_set: BaseDataset = ComponentFactory.create(namespace.Stereotype.DATASET, test_dataset_name)
self.test_set_list.append(test_set)
self.print_info()
def run(self):
Log.info("Loading from epoch {}.".format(self.current_epoch))
self.inference()
Log.success("Inference finished.")
def get_output_data(self, inference_output):
pred_pose_9d = inference_output["pred_pose_9d"]
pred_pose_9d = pred_pose_9d.cpu().numpy()
pred_pose_mat = np.zeros((pred_pose_9d.shape[0], 4, 4))
for i in range(pred_pose_9d.shape[0]):
pred_pose_mat[i, :3, :3] = PoseUtil.rotation_6d_to_matrix_numpy(pred_pose_9d[i, :6])
pred_pose_mat[i, :3, 3] = pred_pose_9d[i, 6:]
pred_pose_mat[i, 3, 3] = 1
#import ipdb; ipdb.set_trace()
pose_matrix,candidate_views, done = self.pbnbv.execute(pred_pose_mat)
rot = pose_matrix[:3,:3]
pose_6d = PoseUtil.matrix_to_rotation_6d_numpy(rot)
translation = pose_matrix[:3, 3]
pose_9d = np.concatenate([pose_6d, translation], axis=0).reshape(1,9)
pose_9d = pose_9d.repeat(50, axis=0)
#import ipdb; ipdb.set_trace()
#self.visualize(pose_matrix,candidate_views)
return {"pred_pose_9d": pose_9d, "candidate_views": candidate_views}
def inference(self):
self.pipeline.eval()
with torch.no_grad():
test_set: BaseDataset
for dataset_idx, test_set in enumerate(self.test_set_list):
status_manager.set_progress("inference", "inferencer", f"dataset", dataset_idx, len(self.test_set_list))
test_set_name = test_set.get_name()
total=int(len(test_set))
for i in tqdm(range(total), desc=f"Processing {test_set_name}", ncols=100):
try:
self.pbnbv.reset()
data = test_set.__getitem__(i)
scene_name = data["scene_name"]
inference_result_path = os.path.join(self.output_dir, test_set_name, f"{scene_name}.pkl")
if os.path.exists(inference_result_path):
Log.info(f"Inference result already exists for scene: {scene_name}")
continue
status_manager.set_progress("inference", "inferencer", f"Batch[{test_set_name}]", i+1, total)
output = self.predict_sequence(data)
self.save_inference_result(test_set_name, data["scene_name"], output)
except Exception as e:
print(e)
Log.error(f"Error, {e}")
import traceback
traceback.print_exc()
continue
status_manager.set_progress("inference", "inferencer", f"dataset", len(self.test_set_list), len(self.test_set_list))
def predict_sequence(self, data, cr_increase_threshold=0, overlap_area_threshold=25, scan_points_threshold=10, max_iter=50, max_retry = 10, max_success=3):
scene_name = data["scene_name"]
Log.info(f"Processing scene: {scene_name}")
status_manager.set_status("inference", "inferencer", "scene", scene_name)
''' data for rendering '''
scene_path = data["scene_path"]
O_to_L_pose = data["O_to_L_pose"]
voxel_threshold = self.voxel_size
filter_degree = 75
down_sampled_model_pts = data["gt_pts"]
first_frame_to_world_9d = data["first_scanned_n_to_world_pose_9d"][0]
first_frame_to_world = np.eye(4)
first_frame_to_world[:3,:3] = PoseUtil.rotation_6d_to_matrix_numpy(first_frame_to_world_9d[:6])
first_frame_to_world[:3,3] = first_frame_to_world_9d[6:]
self.pbnbv.capture(data["first_scanned_pts"][0], first_frame_to_world)
''' data for inference '''
input_data = {}
input_data["combined_scanned_pts"] = torch.tensor(data["first_scanned_pts"][0], dtype=torch.float32).to(self.device).unsqueeze(0)
input_data["scanned_pts_mask"] = [torch.zeros(input_data["combined_scanned_pts"].shape[1], dtype=torch.bool).to(self.device).unsqueeze(0)]
input_data["scanned_n_to_world_pose_9d"] = [torch.tensor(data["first_scanned_n_to_world_pose_9d"], dtype=torch.float32).to(self.device)]
input_data["mode"] = namespace.Mode.TEST
input_pts_N = input_data["combined_scanned_pts"].shape[1]
root = os.path.dirname(scene_path)
display_table_info = DataLoadUtil.get_display_table_info(root, scene_name)
radius = display_table_info["radius"]
scan_points = np.asarray(ReconstructionUtil.generate_scan_points(display_table_top=0,display_table_radius=radius))
first_frame_target_pts, first_frame_target_normals, first_frame_scan_points_indices = RenderUtil.render_pts(first_frame_to_world, scene_path, self.script_path, scan_points, voxel_threshold=voxel_threshold, filter_degree=filter_degree, nO_to_nL_pose=O_to_L_pose)
scanned_view_pts = [first_frame_target_pts]
history_indices = [first_frame_scan_points_indices]
last_pred_cr, added_pts_num = self.compute_coverage_rate(scanned_view_pts, None, down_sampled_model_pts, threshold=voxel_threshold)
retry_duplication_pose = []
retry_no_pts_pose = []
retry_overlap_pose = []
retry = 0
pred_cr_seq = [last_pred_cr]
success = 0
last_pts_num = PtsUtil.voxel_downsample_point_cloud(data["first_scanned_pts"][0], voxel_threshold).shape[0]
import time
while len(pred_cr_seq) < max_iter and retry < max_retry and success < max_success:
#import ipdb; ipdb.set_trace()
Log.green(f"iter: {len(pred_cr_seq)}, retry: {retry}/{max_retry}, success: {success}/{max_success}")
combined_scanned_pts = np.vstack(scanned_view_pts)
voxel_downsampled_combined_scanned_pts_np, inverse = self.voxel_downsample_with_mapping(combined_scanned_pts, voxel_threshold)
inference_output = self.pipeline(input_data)
inference_pred_pose_9d = inference_output["pred_pose_9d"]
output = self.get_output_data(inference_output)
pred_pose_9d = output["pred_pose_9d"]
if not self.enable_cluster:
pred_pose_9d_candidates = [pred_pose_9d[0]]
else:
predict_result = PredictResult(pred_pose_9d.cpu().numpy(), input_pts=input_data["combined_scanned_pts"][0].cpu().numpy(), cluster_params=dict(eps=0.25, min_samples=3))
pred_pose_9d_candidates = predict_result.candidate_9d_poses
pred_pose = np.eye(4)
for pred_pose_9d in pred_pose_9d_candidates:
#import ipdb; ipdb.set_trace()
t1_s = time.time()
pred_pose_9d = np.array(pred_pose_9d, dtype=np.float32)
pred_pose[:3,:3] = PoseUtil.rotation_6d_to_matrix_numpy(pred_pose_9d[:6])
pred_pose[:3,3] = pred_pose_9d[6:]
try:
new_target_pts, new_target_normals, new_scan_points_indices = RenderUtil.render_pts(pred_pose, scene_path, self.script_path, scan_points, voxel_threshold=voxel_threshold, filter_degree=filter_degree, nO_to_nL_pose=O_to_L_pose)
#import ipdb; ipdb.set_trace()
if not ReconstructionUtil.check_scan_points_overlap(history_indices, new_scan_points_indices, scan_points_threshold):
curr_overlap_area_threshold = overlap_area_threshold
else:
curr_overlap_area_threshold = overlap_area_threshold * 0.5
downsampled_new_target_pts = PtsUtil.voxel_downsample_point_cloud(new_target_pts, voxel_threshold)
if self.overlap_limit:
overlap, _ = ReconstructionUtil.check_overlap(downsampled_new_target_pts, voxel_downsampled_combined_scanned_pts_np, overlap_area_threshold = curr_overlap_area_threshold, voxel_size=voxel_threshold, require_new_added_pts_num = True)
if not overlap:
Log.yellow("no overlap!")
retry += 1
retry_overlap_pose.append(pred_pose.tolist())
continue
history_indices.append(new_scan_points_indices)
except Exception as e:
Log.error(f"Error in scene {scene_path}, {e}")
print("current pose: ", pred_pose)
print("curr_pred_cr: ", last_pred_cr)
retry_no_pts_pose.append(pred_pose.tolist())
retry += 1
continue
if new_target_pts.shape[0] == 0:
Log.red("no pts in new target")
retry_no_pts_pose.append(pred_pose.tolist())
retry += 1
continue
pred_cr, _ = self.compute_coverage_rate(scanned_view_pts, new_target_pts, down_sampled_model_pts, threshold=voxel_threshold)
Log.yellow(f"{pred_cr}, {last_pred_cr}, max: , {data['seq_max_coverage_rate']}")
if pred_cr >= data["seq_max_coverage_rate"] - 1e-3:
print("max coverage rate reached!: ", pred_cr)
pred_cr_seq.append(pred_cr)
scanned_view_pts.append(new_target_pts)
input_data["scanned_n_to_world_pose_9d"] = [torch.cat([input_data["scanned_n_to_world_pose_9d"][0], inference_pred_pose_9d], dim=0)]
combined_scanned_pts = np.vstack(scanned_view_pts)
voxel_downsampled_combined_scanned_pts_np = PtsUtil.voxel_downsample_point_cloud(combined_scanned_pts, voxel_threshold)
random_downsampled_combined_scanned_pts_np = PtsUtil.random_downsample_point_cloud(voxel_downsampled_combined_scanned_pts_np, input_pts_N)
self.pbnbv.capture(np.array(random_downsampled_combined_scanned_pts_np, dtype=np.float32), pred_pose)
t1_e = time.time()
t2_s = time.time()
input_data["combined_scanned_pts"] = torch.tensor(random_downsampled_combined_scanned_pts_np, dtype=torch.float32).unsqueeze(0).to(self.device)
last_pred_cr = pred_cr
pts_num = voxel_downsampled_combined_scanned_pts_np.shape[0]
Log.info(f"delta pts num:,{pts_num - last_pts_num },{pts_num}, {last_pts_num}")
if pts_num - last_pts_num < self.min_new_pts_num and pred_cr <= data["seq_max_coverage_rate"] - 1e-2:
retry += 1
retry_duplication_pose.append(pred_pose.tolist())
Log.red(f"delta pts num < {self.min_new_pts_num}:, {pts_num}, {last_pts_num}")
elif pts_num - last_pts_num < self.min_new_pts_num and pred_cr > data["seq_max_coverage_rate"] - 1e-2:
success += 1
Log.success(f"delta pts num < {self.min_new_pts_num}:, {pts_num}, {last_pts_num}")
last_pts_num = pts_num
t2_e = time.time()
print("t1:", t1_e - t1_s)
print("t2:", t2_e - t2_s)
input_data["scanned_n_to_world_pose_9d"] = input_data["scanned_n_to_world_pose_9d"][0].tolist()
result = {
"pred_pose_9d_seq": input_data["scanned_n_to_world_pose_9d"],
"combined_scanned_pts": input_data["combined_scanned_pts"],
"target_pts_seq": scanned_view_pts,
"coverage_rate_seq": pred_cr_seq,
"max_coverage_rate": data["seq_max_coverage_rate"],
"pred_max_coverage_rate": max(pred_cr_seq),
"scene_name": scene_name,
"retry_no_pts_pose": retry_no_pts_pose,
"retry_duplication_pose": retry_duplication_pose,
"retry_overlap_pose": retry_overlap_pose,
"best_seq_len": data["best_seq_len"],
}
self.stat_result[scene_name] = {
"coverage_rate_seq": pred_cr_seq,
"pred_max_coverage_rate": max(pred_cr_seq),
"pred_seq_len": len(pred_cr_seq),
}
print('success rate: ', max(pred_cr_seq))
return result
def voxel_downsample_with_mapping(self, point_cloud, voxel_size=0.003):
voxel_indices = np.floor(point_cloud / voxel_size).astype(np.int32)
unique_voxels, inverse, counts = np.unique(voxel_indices, axis=0, return_inverse=True, return_counts=True)
idx_sort = np.argsort(inverse)
idx_unique = idx_sort[np.cumsum(counts)-counts]
downsampled_points = point_cloud[idx_unique]
return downsampled_points, inverse
def compute_coverage_rate(self, scanned_view_pts, new_pts, model_pts, threshold=0.005):
if new_pts is not None:
new_scanned_view_pts = scanned_view_pts + [new_pts]
else:
new_scanned_view_pts = scanned_view_pts
combined_point_cloud = np.vstack(new_scanned_view_pts)
down_sampled_combined_point_cloud = PtsUtil.voxel_downsample_point_cloud(combined_point_cloud,threshold)
return ReconstructionUtil.compute_coverage_rate(model_pts, down_sampled_combined_point_cloud, threshold)
def voxel_downsample_with_mapping(self, point_cloud, voxel_size=0.003):
voxel_indices = np.floor(point_cloud / voxel_size).astype(np.int32)
unique_voxels, inverse, counts = np.unique(voxel_indices, axis=0, return_inverse=True, return_counts=True)
idx_sort = np.argsort(inverse)
idx_unique = idx_sort[np.cumsum(counts)-counts]
downsampled_points = point_cloud[idx_unique]
return downsampled_points, inverse
def save_inference_result(self, dataset_name, scene_name, output):
dataset_dir = os.path.join(self.output_dir, dataset_name)
if not os.path.exists(dataset_dir):
os.makedirs(dataset_dir)
output_path = os.path.join(dataset_dir, f"{scene_name}.pkl")
pickle.dump(output, open(output_path, "wb"))
with open(self.stat_result_path, "w") as f:
json.dump(self.stat_result, f)
def get_checkpoint_path(self, is_last=False):
return os.path.join(self.experiment_path, namespace.Direcotry.CHECKPOINT_DIR_NAME,
"Epoch_{}.pth".format(
self.current_epoch if self.current_epoch != -1 and not is_last else "last"))
def load_checkpoint(self, is_last=False):
self.load(self.get_checkpoint_path(is_last))
Log.success(f"Loaded checkpoint from {self.get_checkpoint_path(is_last)}")
if is_last:
checkpoint_root = os.path.join(self.experiment_path, namespace.Direcotry.CHECKPOINT_DIR_NAME)
meta_path = os.path.join(checkpoint_root, "meta.json")
if not os.path.exists(meta_path):
raise FileNotFoundError(
"No checkpoint meta.json file in the experiment {}".format(self.experiments_config["name"]))
file_path = os.path.join(checkpoint_root, "meta.json")
with open(file_path, "r") as f:
meta = json.load(f)
self.current_epoch = meta["last_epoch"]
self.current_iter = meta["last_iter"]
def load_experiment(self, backup_name=None):
super().load_experiment(backup_name)
self.current_epoch = self.experiments_config["epoch"]
self.load_checkpoint(is_last=(self.current_epoch == -1))
def create_experiment(self, backup_name=None):
super().create_experiment(backup_name)
def load(self, path):
state_dict = torch.load(path)
self.pipeline.load_state_dict(state_dict)
def print_info(self):
def print_dataset(dataset: BaseDataset):
config = dataset.get_config()
name = dataset.get_name()
Log.blue(f"Dataset: {name}")
for k,v in config.items():
Log.blue(f"\t{k}: {v}")
super().print_info()
table_size = 70
Log.blue(f"{'+' + '-' * (table_size // 2)} Pipeline {'-' * (table_size // 2)}" + '+')
Log.blue(self.pipeline)
Log.blue(f"{'+' + '-' * (table_size // 2)} Datasets {'-' * (table_size // 2)}" + '+')
for i, test_set in enumerate(self.test_set_list):
Log.blue(f"test dataset {i}: ")
print_dataset(test_set)
Log.blue(f"{'+' + '-' * (table_size // 2)}----------{'-' * (table_size // 2)}" + '+')

8192
test.txt Normal file

File diff suppressed because it is too large Load Diff