233 lines
11 KiB
Python
233 lines
11 KiB
Python
import os
|
|
import json
|
|
from utils.render import RenderUtil
|
|
from utils.pose import PoseUtil
|
|
from utils.pts import PtsUtil
|
|
from utils.reconstruction import ReconstructionUtil
|
|
|
|
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
|
|
|
|
@stereotype.runner("inferencer", comment="not tested")
|
|
class Inferencer(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")
|
|
''' 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")
|
|
|
|
''' 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 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()
|
|
test_loader = test_set.get_loader()
|
|
|
|
if test_loader.batch_size > 1:
|
|
Log.error("Batch size should be 1 for inference, found {} in {}".format(test_loader.batch_size, test_set_name), terminate=True)
|
|
|
|
total=int(len(test_loader))
|
|
loop = tqdm(enumerate(test_loader), total=total)
|
|
for i, data in loop:
|
|
status_manager.set_progress("inference", "inferencer", f"Batch[{test_set_name}]", i+1, total)
|
|
test_set.process_batch(data, self.device)
|
|
output = self.predict_sequence(data)
|
|
self.save_inference_result(output, data)
|
|
|
|
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=100):
|
|
pred_cr_seq = []
|
|
scene_name = data["scene_name"][0]
|
|
Log.info(f"Processing scene: {scene_name}")
|
|
status_manager.set_status("inference", "inferencer", "scene", scene_name)
|
|
|
|
''' data for rendering '''
|
|
scene_path = data["scene_path"][0]
|
|
O_to_L_pose = data["O_to_L_pose"][0]
|
|
voxel_threshold = data["voxel_threshold"][0]
|
|
filter_degree = data["filter_degree"][0]
|
|
model_points_normals = data["model_points_normals"][0]
|
|
model_pts = model_points_normals[:,:3]
|
|
down_sampled_model_pts = PtsUtil.voxel_downsample_point_cloud(model_pts, voxel_threshold)
|
|
first_frame_to_world = data["first_frame_to_world"][0]
|
|
|
|
''' data for inference '''
|
|
input_data = {}
|
|
input_data["scanned_pts"] = [data["first_pts"][0].to(self.device)]
|
|
input_data["scanned_n_to_1_pose_9d"] = [data["first_to_first_9d"][0].to(self.device)]
|
|
input_data["mode"] = namespace.Mode.TEST
|
|
input_pts_N = input_data["scanned_pts"][0].shape[1]
|
|
|
|
|
|
first_frame_target_pts, _ = RenderUtil.render_pts(first_frame_to_world, scene_path, self.script_path, model_points_normals, voxel_threshold=voxel_threshold, filter_degree=filter_degree, nO_to_nL_pose=O_to_L_pose)
|
|
scanned_view_pts = [first_frame_target_pts]
|
|
last_pred_cr = self.compute_coverage_rate(scanned_view_pts, None, down_sampled_model_pts, threshold=voxel_threshold)
|
|
|
|
|
|
|
|
while len(pred_cr_seq) < max_iter:
|
|
|
|
output = self.pipeline(input_data)
|
|
next_pose_9d = output["pred_pose_9d"]
|
|
pred_pose = torch.eye(4, device=next_pose_9d.device)
|
|
|
|
pred_pose[:3,:3] = PoseUtil.rotation_6d_to_matrix_tensor_batch(next_pose_9d[:,:6])[0]
|
|
pred_pose[:3,3] = next_pose_9d[0,6:]
|
|
pred_n_to_world_pose_mat = torch.matmul(first_frame_to_world, pred_pose)
|
|
try:
|
|
new_target_pts_world, new_pts_world = RenderUtil.render_pts(pred_n_to_world_pose_mat, scene_path, self.script_path, model_points_normals, voxel_threshold=voxel_threshold, filter_degree=filter_degree, nO_to_nL_pose=O_to_L_pose, require_full_scene=True)
|
|
except Exception as e:
|
|
Log.warning(f"Error in scene {scene_path}, {e}")
|
|
print("current pose: ", pred_pose)
|
|
print("curr_pred_cr: ", last_pred_cr)
|
|
continue
|
|
|
|
|
|
pred_cr = self.compute_coverage_rate(scanned_view_pts, new_target_pts_world, down_sampled_model_pts, threshold=voxel_threshold)
|
|
pred_cr_seq.append(pred_cr)
|
|
if pred_cr >= data["max_coverage_rate"]:
|
|
break
|
|
if pred_cr < last_pred_cr + cr_increase_threshold:
|
|
break
|
|
scanned_view_pts.append(new_target_pts_world)
|
|
down_sampled_new_pts_world = PtsUtil.random_downsample_point_cloud(new_pts_world, input_pts_N)
|
|
new_pts_world_aug = np.hstack([down_sampled_new_pts_world, np.ones((down_sampled_new_pts_world.shape[0], 1))])
|
|
new_pts = np.dot(np.linalg.inv(first_frame_to_world.cpu()), new_pts_world_aug.T).T[:,:3]
|
|
|
|
new_pts_tensor = torch.tensor(new_pts, dtype=torch.float32).unsqueeze(0).to(self.device)
|
|
|
|
input_data["scanned_pts"] = [torch.cat([input_data["scanned_pts"][0] , new_pts_tensor], dim=0)]
|
|
input_data["scanned_n_to_1_pose_9d"] = [torch.cat([input_data["scanned_n_to_1_pose_9d"][0], next_pose_9d], dim=0)]
|
|
|
|
last_pred_cr = pred_cr
|
|
# ------ Debug Start ------
|
|
import ipdb;ipdb.set_trace()
|
|
# ------ Debug End ------
|
|
|
|
|
|
input_data["scanned_pts"] = input_data["scanned_pts"][0].cpu().numpy().tolist()
|
|
input_data["scanned_n_to_1_pose_9d"] = input_data["scanned_n_to_1_pose_9d"][0].cpu().numpy().tolist()
|
|
result = {
|
|
"pred_pose_9d_seq": input_data["scanned_n_to_1_pose_9d"],
|
|
"pts_seq": input_data["scanned_pts"],
|
|
"target_pts_seq": scanned_view_pts,
|
|
"coverage_rate_seq": pred_cr_seq,
|
|
"max_coverage_rate": data["max_coverage_rate"],
|
|
"pred_max_coverage_rate": max(pred_cr_seq)
|
|
}
|
|
return result
|
|
|
|
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 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)
|
|
pickle.dump(output, open(f"result_{scene_name}.pkl", "wb"))
|
|
|
|
|
|
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)}" + '+')
|
|
|