upd infernce
This commit is contained in:
hofee
@@ -19,7 +19,7 @@ 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
|
||||
@stereotype.runner("inferencer")
|
||||
class Inferencer(Runner):
|
||||
def __init__(self, config_path):
|
||||
@@ -35,7 +35,12 @@ class Inferencer(Runner):
|
||||
|
||||
''' Experiment '''
|
||||
self.load_experiment("nbv_evaluator")
|
||||
self.stat_result = {}
|
||||
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)
|
||||
@@ -68,22 +73,21 @@ class Inferencer(Runner):
|
||||
test_set_name = test_set.get_name()
|
||||
|
||||
total=int(len(test_set))
|
||||
scene_name_list = test_set.get_scene_name_list()
|
||||
for i in range(total):
|
||||
scene_name = scene_name_list[i]
|
||||
for i in tqdm(range(total), desc=f"Processing {test_set_name}", ncols=100):
|
||||
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
|
||||
data = test_set.__getitem__(i)
|
||||
|
||||
status_manager.set_progress("inference", "inferencer", f"Batch[{test_set_name}]", i+1, total)
|
||||
scene_name = data["scene_name"]
|
||||
output = self.predict_sequence(data)
|
||||
self.save_inference_result(test_set_name, data["scene_name"], output)
|
||||
|
||||
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, max_iter=50, max_retry=5):
|
||||
def predict_sequence(self, data, cr_increase_threshold=0.001, overlap_area_threshold=25, scan_points_threshold=10, max_iter=50, max_retry = 7):
|
||||
scene_name = data["scene_name"]
|
||||
Log.info(f"Processing scene: {scene_name}")
|
||||
status_manager.set_status("inference", "inferencer", "scene", scene_name)
|
||||
@@ -102,16 +106,23 @@ class Inferencer(Runner):
|
||||
|
||||
''' data for inference '''
|
||||
input_data = {}
|
||||
scanned_pts = []
|
||||
input_data["combined_scanned_pts"] = torch.tensor(data["first_scanned_pts"][0], dtype=torch.float32).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]
|
||||
first_frame_target_pts, first_frame_target_normals = RenderUtil.render_pts(first_frame_to_world, scene_path, self.script_path, voxel_threshold=voxel_threshold, filter_degree=filter_degree, nO_to_nL_pose=O_to_L_pose)
|
||||
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)
|
||||
|
||||
scanned_pts.append(first_frame_target_pts)
|
||||
retry_duplication_pose = []
|
||||
retry_no_pts_pose = []
|
||||
retry_overlap_pose = []
|
||||
retry = 0
|
||||
pred_cr_seq = [last_pred_cr]
|
||||
success = 0
|
||||
@@ -129,7 +140,22 @@ class Inferencer(Runner):
|
||||
|
||||
try:
|
||||
start_time = time.time()
|
||||
new_target_pts, new_target_normals = RenderUtil.render_pts(pred_pose, scene_path, self.script_path, voxel_threshold=voxel_threshold, filter_degree=filter_degree, nO_to_nL_pose=O_to_L_pose)
|
||||
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)
|
||||
overlap, new_added_pts_num = ReconstructionUtil.check_overlap(downsampled_new_target_pts, down_sampled_model_pts, overlap_area_threshold = curr_overlap_area_threshold, voxel_size=voxel_threshold, require_new_added_pts_num = True)
|
||||
if not overlap:
|
||||
retry += 1
|
||||
retry_overlap_pose.append(pred_pose.cpu().numpy().tolist())
|
||||
continue
|
||||
|
||||
scanned_pts.append(new_target_pts)
|
||||
history_indices.append(new_scan_points_indices)
|
||||
end_time = time.time()
|
||||
print(f"Time taken for rendering: {end_time - start_time} seconds")
|
||||
except Exception as e:
|
||||
@@ -147,14 +173,16 @@ class Inferencer(Runner):
|
||||
continue
|
||||
|
||||
start_time = time.time()
|
||||
pred_cr, new_added_pts_num = self.compute_coverage_rate(scanned_view_pts, new_target_pts, down_sampled_model_pts, threshold=voxel_threshold)
|
||||
pred_cr, covered_pts_num = self.compute_coverage_rate(scanned_view_pts, new_target_pts, down_sampled_model_pts, threshold=voxel_threshold)
|
||||
end_time = time.time()
|
||||
print(f"Time taken for coverage rate computation: {end_time - start_time} seconds")
|
||||
print(pred_cr, last_pred_cr, " max: ", data["seq_max_coverage_rate"])
|
||||
print("new added pts num: ", new_added_pts_num)
|
||||
if pred_cr >= data["seq_max_coverage_rate"] - 1e-3:
|
||||
print("max coverage rate reached!: ", pred_cr)
|
||||
success += 1
|
||||
elif new_added_pts_num < 10:
|
||||
elif new_added_pts_num < 5:
|
||||
success += 1
|
||||
print("min added pts num reached!: ", new_added_pts_num)
|
||||
if pred_cr <= last_pred_cr + cr_increase_threshold:
|
||||
retry += 1
|
||||
@@ -180,6 +208,7 @@ class Inferencer(Runner):
|
||||
|
||||
input_data["scanned_n_to_world_pose_9d"] = input_data["scanned_n_to_world_pose_9d"][0].cpu().numpy().tolist()
|
||||
result = {
|
||||
"scanned_pts": scanned_pts,
|
||||
"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,
|
||||
@@ -189,6 +218,7 @@ class Inferencer(Runner):
|
||||
"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] = {
|
||||
@@ -216,7 +246,7 @@ class Inferencer(Runner):
|
||||
os.makedirs(dataset_dir)
|
||||
output_path = os.path.join(dataset_dir, f"{scene_name}.pkl")
|
||||
pickle.dump(output, open(output_path, "wb"))
|
||||
with open(os.path.join(dataset_dir, "stat.json"), "w") as f:
|
||||
with open(self.stat_result_path, "w") as f:
|
||||
json.dump(self.stat_result, f)
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user