diff --git a/app_inference.py b/app_inference.py index 1c464fd..7bd97b9 100644 --- a/app_inference.py +++ b/app_inference.py @@ -1,5 +1,6 @@ from PytorchBoot.application import PytorchBootApplication from runners.inferencer import Inferencer +from runners.inference_server import InferencerServer @PytorchBootApplication("inference") class InferenceApp: @@ -14,3 +15,17 @@ class InferenceApp: Evaluator("path_to_your_eval_config").run() ''' Inferencer("./configs/local/inference_config.yaml").run() + +@PytorchBootApplication("server") +class InferenceServerApp: + @staticmethod + def start(): + ''' + call default or your custom runners here, code will be executed + automatically when type "pytorch-boot run" or "ptb run" in terminal + + example: + Trainer("path_to_your_train_config").run() + Evaluator("path_to_your_eval_config").run() + ''' + InferencerServer("./configs/server/server_inference_server_config.yaml").run() \ No newline at end of file diff --git a/configs/server/server_inference_server_config.yaml b/configs/server/server_inference_server_config.yaml new file mode 100644 index 0000000..42bdbd4 --- /dev/null +++ b/configs/server/server_inference_server_config.yaml @@ -0,0 +1,53 @@ + +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 + +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 + +module: + pointnet_encoder: + in_dim: 3 + out_dim: 1024 + global_feat: True + feature_transform: False + transformer_seq_encoder: + embed_dim: 256 + num_heads: 4 + ffn_dim: 256 + num_layers: 3 + output_dim: 1024 + + gf_view_finder: + t_feat_dim: 128 + pose_feat_dim: 256 + main_feat_dim: 2048 + regression_head: Rx_Ry_and_T + pose_mode: rot_matrix + per_point_feature: False + sample_mode: ode + sampling_steps: 500 + sde_mode: ve + pose_encoder: + pose_dim: 9 + out_dim: 256 + pts_num_encoder: + out_dim: 64 \ No newline at end of file diff --git a/runners/inferece_server.py b/runners/inference_server.py similarity index 69% rename from runners/inferece_server.py rename to runners/inference_server.py index 239a0e0..e71b28b 100644 --- a/runners/inferece_server.py +++ b/runners/inference_server.py @@ -13,7 +13,7 @@ from PytorchBoot.utils import Log from utils.pts import PtsUtil -@stereotype.runner("inferencer") +@stereotype.runner("inferencer_server") class InferencerServer(Runner): def __init__(self, config_path): super().__init__(config_path) @@ -24,9 +24,10 @@ class InferencerServer(Runner): 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 ''' Experiment ''' - self.load_experiment("nbv_evaluator") + self.load_experiment("inferencer_server") def get_input_data(self, data): input_data = {} @@ -36,28 +37,36 @@ class InferencerServer(Runner): fps_downsampled_combined_scanned_pts, fps_idx = PtsUtil.fps_downsample_point_cloud( combined_scanned_views_pts, self.pts_num, require_idx=True ) - combined_scanned_views_pts_mask = np.zeros(len(scanned_pts), dtype=np.uint8) - start_idx = 0 - for i in range(len(scanned_pts)): - end_idx = start_idx + len(scanned_pts[i]) - combined_scanned_views_pts_mask[start_idx:end_idx] = i - start_idx = end_idx + # combined_scanned_views_pts_mask = np.zeros(len(scanned_pts), dtype=np.uint8) + # start_idx = 0 + # for i in range(len(scanned_pts)): + # end_idx = start_idx + len(scanned_pts[i]) + # combined_scanned_views_pts_mask[start_idx:end_idx] = i + # start_idx = end_idx - fps_downsampled_combined_scanned_pts_mask = combined_scanned_views_pts_mask[fps_idx] + # fps_downsampled_combined_scanned_pts_mask = combined_scanned_views_pts_mask[fps_idx] - input_data["scanned_pts_mask"] = np.asarray(fps_downsampled_combined_scanned_pts_mask, dtype=np.uint8) + input_data["scanned_pts"] = scanned_pts + # input_data["scanned_pts_mask"] = np.asarray(fps_downsampled_combined_scanned_pts_mask, dtype=np.uint8) 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): - estimated_delta_rot_9d = output_data["pred_pose_9d"] + pred_pose_9d = output_data["pred_pose_9d"] result = { - "estimated_delta_rot_9d": estimated_delta_rot_9d.tolist() + "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 run(self): Log.info("Loading from epoch {}.".format(self.current_epoch)) @@ -65,7 +74,8 @@ class InferencerServer(Runner): def inference(): data = request.json input_data = self.get_input_data(data) - output_data = self.pipeline.forward_test(input_data) + collated_input_data = self.collate_input(input_data) + output_data = self.pipeline.forward_test(collated_input_data) result = self.get_result(output_data) return jsonify(result) diff --git a/runners/inferencer.py b/runners/inferencer.py index 797f47e..50f838c 100644 --- a/runners/inferencer.py +++ b/runners/inferencer.py @@ -115,9 +115,12 @@ class Inferencer(Runner): retry = 0 pred_cr_seq = [last_pred_cr] success = 0 + import time while len(pred_cr_seq) < max_iter and retry < max_retry: - + start_time = time.time() output = self.pipeline(input_data) + end_time = time.time() + print(f"Time taken for inference: {end_time - start_time} seconds") pred_pose_9d = output["pred_pose_9d"] pred_pose = torch.eye(4, device=pred_pose_9d.device) @@ -125,7 +128,10 @@ class Inferencer(Runner): pred_pose[:3,3] = pred_pose_9d[0,6:] 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) + end_time = time.time() + print(f"Time taken for rendering: {end_time - start_time} seconds") except Exception as e: Log.warning(f"Error in scene {scene_path}, {e}") print("current pose: ", pred_pose) @@ -140,8 +146,10 @@ class Inferencer(Runner): retry += 1 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) + 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"]) if pred_cr >= data["seq_max_coverage_rate"] - 1e-3: print("max coverage rate reached!: ", pred_cr) diff --git a/utils/pts.py b/utils/pts.py index 7b640df..66d3d84 100644 --- a/utils/pts.py +++ b/utils/pts.py @@ -16,6 +16,17 @@ class PtsUtil: else: unique_voxels = np.unique(voxel_indices, axis=0, return_inverse=True) return unique_voxels[0]*voxel_size + + @staticmethod + def voxel_downsample_point_cloud_random(point_cloud, voxel_size=0.005, require_idx=False): + 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] + if require_idx: + return downsampled_points, inverse + return downsampled_points @staticmethod def random_downsample_point_cloud(point_cloud, num_points, require_idx=False): diff --git a/utils/render.py b/utils/render.py index 15f9c3e..5c1f5a2 100644 --- a/utils/render.py +++ b/utils/render.py @@ -1,6 +1,7 @@ import os import json +import time import subprocess import tempfile import shutil @@ -68,9 +69,13 @@ class RenderUtil: params_data_path = os.path.join(temp_dir, "params.json") with open(params_data_path, 'w') as f: json.dump(params, f) + start_time = time.time() result = subprocess.run([ 'blender', '-b', '-P', script_path, '--', temp_dir ], capture_output=True, text=True) + end_time = time.time() + print(result) + print(f"-- Time taken for blender: {end_time - start_time} seconds") if result.returncode != 0: print("Blender script failed:") print(result.stderr) @@ -82,6 +87,7 @@ class RenderUtil: cam_info["far_plane"], binocular=True ) + start_time = time.time() mask_L, mask_R = DataLoadUtil.load_seg(path, binocular=True) normal_L = DataLoadUtil.load_normal(path, binocular=True, left_only=True) ''' target points ''' @@ -114,6 +120,7 @@ class RenderUtil: if not has_points: target_points = np.zeros((0, 3)) target_normals = np.zeros((0, 3)) - + end_time = time.time() + print(f"-- Time taken for processing: {end_time - start_time} seconds") #import ipdb; ipdb.set_trace() return target_points, target_normals \ No newline at end of file