upd
This commit is contained in:
@@ -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()
|
||||
|
@@ -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
|
||||
|
43
configs/server/server_heuristic_inference_server_config.yaml
Normal file
43
configs/server/server_heuristic_inference_server_config.yaml
Normal 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
|
Binary file not shown.
@@ -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.
Binary file not shown.
BIN
runners/__pycache__/pb_global_points_inferencer.cpython-39.pyc
Normal file
BIN
runners/__pycache__/pb_global_points_inferencer.cpython-39.pyc
Normal file
Binary file not shown.
BIN
runners/__pycache__/ug_inference_server.cpython-39.pyc
Normal file
BIN
runners/__pycache__/ug_inference_server.cpython-39.pyc
Normal file
Binary file not shown.
@@ -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()
|
||||
|
216
runners/inference_heuristic_server.py
Normal file
216
runners/inference_heuristic_server.py
Normal 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)
|
||||
|
||||
|
||||
|
@@ -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)
|
||||
|
||||
|
829
runners/pb_global_points_inferencer.py
Normal file
829
runners/pb_global_points_inferencer.py
Normal 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)}" + '+')
|
||||
|
Reference in New Issue
Block a user