From 935069d68c35c6f842874c8649ca6f56c2c29446 Mon Sep 17 00:00:00 2001 From: hofee Date: Thu, 19 Sep 2024 00:14:26 +0800 Subject: [PATCH] add inference --- app_inference.py | 16 +++ configs/local/inference_config.yaml | 70 ++++++++++ core/dataset.py | 10 +- core/evaluation.py | 35 +---- core/pipeline.py | 11 +- core/seq_dataset.py | 10 +- runners/inferencer.py | 210 +++++++++++++++++----------- runners/strategy_generator.py | 22 +-- utils/pts.py | 6 + utils/render.py | 51 +++++++ 10 files changed, 302 insertions(+), 139 deletions(-) create mode 100644 app_inference.py create mode 100644 configs/local/inference_config.yaml create mode 100644 utils/render.py diff --git a/app_inference.py b/app_inference.py new file mode 100644 index 0000000..1c464fd --- /dev/null +++ b/app_inference.py @@ -0,0 +1,16 @@ +from PytorchBoot.application import PytorchBootApplication +from runners.inferencer import Inferencer + +@PytorchBootApplication("inference") +class InferenceApp: + @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() + ''' + Inferencer("./configs/local/inference_config.yaml").run() diff --git a/configs/local/inference_config.yaml b/configs/local/inference_config.yaml new file mode 100644 index 0000000..f8fed9a --- /dev/null +++ b/configs/local/inference_config.yaml @@ -0,0 +1,70 @@ + +runner: + general: + seed: 1 + device: cuda + cuda_visible_devices: "0,1,2,3,4,5,6,7" + + experiment: + name: local_eval + root_dir: "experiments" + epoch: 600 # -1 stands for last epoch + + test: + dataset_list: + - OmniObject3d_train + + blender_script_path: "/media/hofee/data/project/python/nbv_reconstruction/blender/data_renderer.py" + output_dir: "/media/hofee/data/project/python/nbv_reconstruction/sample_for_training/inference_result" + pipeline: nbv_reconstruction_pipeline + +dataset: + OmniObject3d_train: + root_dir: "/media/hofee/data/project/python/nbv_reconstruction/sample_for_training/scenes" + model_dir: "/media/hofee/data/data/scaled_object_meshes" + source: seq_nbv_reconstruction_dataset + split_file: "/media/hofee/data/project/python/nbv_reconstruction/sample_for_training/OmniObject3d_train.txt" + type: test + filter_degree: 75 + ratio: 1 + batch_size: 1 + num_workers: 12 + pts_num: 4096 + +pipeline: + nbv_reconstruction_pipeline: + pts_encoder: pointnet_encoder + seq_encoder: transformer_seq_encoder + pose_encoder: pose_encoder + view_finder: gf_view_finder + +module: + + pointnet_encoder: + in_dim: 3 + out_dim: 1024 + global_feat: True + feature_transform: False + + transformer_seq_encoder: + pts_embed_dim: 1024 + pose_embed_dim: 256 + num_heads: 4 + ffn_dim: 256 + num_layers: 3 + output_dim: 2048 + + 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 diff --git a/core/dataset.py b/core/dataset.py index 1fe3980..679c4a9 100644 --- a/core/dataset.py +++ b/core/dataset.py @@ -27,7 +27,7 @@ class NBVReconstructionDataset(BaseDataset): self.pts_num = config["pts_num"] self.type = config["type"] - self.cache = config["cache"] + self.cache = config.get("cache") if self.type == namespace.Mode.TEST: self.model_dir = config["model_dir"] self.filter_degree = config["filter_degree"] @@ -105,7 +105,10 @@ class NBVReconstructionDataset(BaseDataset): nR_to_world_pose = cam_info["cam_to_world_R"] n_to_1_pose = np.dot(np.linalg.inv(first_frame_to_world), n_to_world_pose) nR_to_1_pose = np.dot(np.linalg.inv(first_frame_to_world), nR_to_world_pose) - cached_data = self.load_from_cache(scene_name, first_frame_idx, frame_idx) + + cached_data = None + if self.cache: + cached_data = self.load_from_cache(scene_name, first_frame_idx, frame_idx) if cached_data is None: depth_L, depth_R = DataLoadUtil.load_depth(view_path, cam_info['near_plane'], cam_info['far_plane'], binocular=True) @@ -116,7 +119,8 @@ class NBVReconstructionDataset(BaseDataset): point_cloud_R = PtsUtil.random_downsample_point_cloud(point_cloud_R, 65536) overlap_points = DataLoadUtil.get_overlapping_points(point_cloud_L, point_cloud_R) downsampled_target_point_cloud = PtsUtil.random_downsample_point_cloud(overlap_points, self.pts_num) - self.save_to_cache(scene_name, first_frame_idx, frame_idx, downsampled_target_point_cloud) + if self.cache: + self.save_to_cache(scene_name, first_frame_idx, frame_idx, downsampled_target_point_cloud) else: downsampled_target_point_cloud = cached_data diff --git a/core/evaluation.py b/core/evaluation.py index 047df1c..97793fe 100644 --- a/core/evaluation.py +++ b/core/evaluation.py @@ -1,43 +1,14 @@ import torch -import os -import json import numpy as np -import subprocess -import tempfile -from utils.data_load import DataLoadUtil from utils.reconstruction import ReconstructionUtil from utils.pose import PoseUtil from utils.pts import PtsUtil +from utils.render import RenderUtil import PytorchBoot.stereotype as stereotype import PytorchBoot.namespace as namespace from PytorchBoot.utils.log_util import Log -def render_pts(cam_pose, scene_path,script_path, model_points_normals, voxel_threshold=0.005, filter_degree=75, nO_to_nL_pose=None): - nO_to_world_pose = cam_pose.cpu().numpy() @ nO_to_nL_pose - nO_to_world_pose = DataLoadUtil.cam_pose_transformation(nO_to_world_pose) - - - with tempfile.TemporaryDirectory() as temp_dir: - params = { - "cam_pose": nO_to_world_pose.tolist(), - "scene_path": scene_path - } - params_data_path = os.path.join(temp_dir, "params.json") - with open(params_data_path, 'w') as f: - json.dump(params, f) - result = subprocess.run([ - 'blender', '-b', '-P', script_path, '--', temp_dir - ], capture_output=True, text=True) - if result.returncode != 0: - print("Blender script failed:") - print(result.stderr) - return None - path = os.path.join(temp_dir, "tmp") - - point_cloud = DataLoadUtil.get_target_point_cloud_world_from_path(path, binocular=True) - cam_params = DataLoadUtil.load_cam_info(path, binocular=True) - sampled_point_cloud = ReconstructionUtil.filter_points(point_cloud, model_points_normals, cam_pose=cam_params["cam_to_world"], voxel_size=voxel_threshold, theta=filter_degree) - return sampled_point_cloud + @stereotype.evaluation_method("pose_diff") class PoseDiff: @@ -110,7 +81,7 @@ class ConverageRateIncrease: filter_degree = filter_degree_list[idx] nO_to_nL_pose = nO_to_nL_pose_batch[idx] try: - new_pts = render_pts(pred_pose, scene_path, self.renderer_path, model_points_normals, voxel_threshold=voxel_threshold, filter_degree=filter_degree, nO_to_nL_pose=nO_to_nL_pose) + new_pts, _ = RenderUtil.render_pts(pred_pose, scene_path, self.renderer_path, model_points_normals, voxel_threshold=voxel_threshold, filter_degree=filter_degree, nO_to_nL_pose=nO_to_nL_pose) pred_cr = self.compute_coverage_rate(scanned_view_pts, new_pts, down_sampled_model_pts, threshold=voxel_threshold) except Exception as e: Log.warning(f"Error in scene {scene_path}, {e}") diff --git a/core/pipeline.py b/core/pipeline.py index d1a9cd9..5066dcd 100644 --- a/core/pipeline.py +++ b/core/pipeline.py @@ -5,7 +5,7 @@ import PytorchBoot.stereotype as stereotype from PytorchBoot.factory.component_factory import ComponentFactory from PytorchBoot.utils import Log -@stereotype.pipeline("nbv_reconstruction_pipeline", comment="should be tested") +@stereotype.pipeline("nbv_reconstruction_pipeline") class NBVReconstructionPipeline(nn.Module): def __init__(self, config): super(NBVReconstructionPipeline, self).__init__() @@ -67,14 +67,13 @@ class NBVReconstructionPipeline(nn.Module): def get_seq_feat(self, data): scanned_pts_batch = data['scanned_pts'] scanned_n_to_1_pose_9d_batch = data['scanned_n_to_1_pose_9d'] - best_to_1_pose_9d_batch = data["best_to_1_pose_9d"] pts_feat_seq_list = [] pose_feat_seq_list = [] - + device = next(self.parameters()).device for scanned_pts,scanned_n_to_1_pose_9d in zip(scanned_pts_batch,scanned_n_to_1_pose_9d_batch): - - scanned_pts = scanned_pts.to(best_to_1_pose_9d_batch.device) - scanned_n_to_1_pose_9d = scanned_n_to_1_pose_9d.to(best_to_1_pose_9d_batch.device) + + scanned_pts = scanned_pts.to(device) + scanned_n_to_1_pose_9d = scanned_n_to_1_pose_9d.to(device) pts_feat_seq_list.append(self.pts_encoder.encode_points(scanned_pts)) pose_feat_seq_list.append(self.pose_encoder.encode_pose(scanned_n_to_1_pose_9d)) diff --git a/core/seq_dataset.py b/core/seq_dataset.py index aecc892..b34746b 100644 --- a/core/seq_dataset.py +++ b/core/seq_dataset.py @@ -51,7 +51,7 @@ class SeqNBVReconstructionDataset(BaseDataset): "first_frame": first_frame, "max_coverage_rate": max_coverage_rate }) - return datalist + return datalist[5:] def __getitem__(self, index): data_item_info = self.datalist[index] @@ -85,6 +85,7 @@ class SeqNBVReconstructionDataset(BaseDataset): voxel_threshold = diag*0.02 first_O_to_first_L_pose = np.dot(np.linalg.inv(first_left_cam_pose), first_center_cam_pose) scene_path = os.path.join(self.root_dir, scene_name) + model_points_normals = DataLoadUtil.load_points_normals(self.root_dir, scene_name) data_item = { "first_pts": np.asarray([first_downsampled_target_point_cloud],dtype=np.float32), "first_to_first_9d": np.asarray([first_to_first_9d],dtype=np.float32), @@ -92,10 +93,11 @@ class SeqNBVReconstructionDataset(BaseDataset): "max_coverage_rate": max_coverage_rate, "voxel_threshold": voxel_threshold, "filter_degree": self.filter_degree, - "first_frame_to_world": first_frame_to_world, - "first_O_to_first_L_pose": first_O_to_first_L_pose, + "first_frame_to_world": np.asarray(first_frame_to_world, dtype=np.float32), + "O_to_L_pose": first_O_to_first_L_pose, "first_frame_coverage": first_frame_coverage, - "scene_path": scene_path + "scene_path": scene_path, + "model_points_normals": model_points_normals, } return data_item diff --git a/runners/inferencer.py b/runners/inferencer.py index a4d49a8..bb17bde 100644 --- a/runners/inferencer.py +++ b/runners/inferencer.py @@ -1,33 +1,35 @@ import os import json -from datetime import datetime +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.factory import OptimizerFactory from PytorchBoot.dataset import BaseDataset from PytorchBoot.runners.runner import Runner -from PytorchBoot.stereotype import EXTERNAL_FRONZEN_MODULES from PytorchBoot.utils import Log from PytorchBoot.status import status_manager -@stereotype.runner("nbv_evaluator") -class NextBestViewEvaluator(Runner): +@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.parallel = self.config["general"]["parallel"] self.pipeline:torch.nn.Module = ComponentFactory.create(namespace.Stereotype.PIPELINE, self.pipeline_name) - if self.parallel and self.device == "cuda": - self.pipeline = torch.nn.DataParallel(self.pipeline) self.pipeline = self.pipeline.to(self.device) ''' Experiment ''' @@ -46,55 +48,135 @@ class NextBestViewEvaluator(Runner): 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.test() + self.inference() + Log.success("Inference finished.") - def test(self): + def inference(self): self.pipeline.eval() with torch.no_grad(): test_set: BaseDataset for dataset_idx, test_set in enumerate(self.test_set_list): - test_set_config = test_set.get_config() - eval_list = test_set_config["eval_list"] - ratio = test_set_config["ratio"] + status_manager.set_progress("inference", "inferencer", f"dataset", dataset_idx, len(self.test_set_list)) test_set_name = test_set.get_name() - - output_list = [] - data_list = [] 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("train", "default_trainer", f"(test) Batch[{test_set_name}]", i+1, total) + status_manager.set_progress("inference", "inferencer", f"Batch[{test_set_name}]", i+1, total) test_set.process_batch(data, self.device) - data["mode"] = namespace.Mode.TEST - output = self.pipeline(data) - output_list.append(output) - data_list.append(data) - loop.set_description( - f'Epoch [{self.current_epoch}/{self.max_epochs}] (Test: {test_set_name}, ratio={ratio})') - result_dict = self.eval_fn(output_list, data_list, eval_list) - - @staticmethod - def eval_fn(output_list, data_list, eval_list): - collected_result = {} - for eval_method_name in eval_list: - eval_method = ComponentFactory.create(namespace.Stereotype.EVALUATION_METHOD, eval_method_name) - eval_results:dict = eval_method.evaluate(output_list, data_list) - for data_type, eval_result in eval_results.items(): - if data_type not in collected_result: - collected_result[data_type] = {} - for name, value in eval_result.items(): - collected_result[data_type][name] = value - status_manager.set_status("train", "default_trainer", f"[eval]{name}", value) + 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)) - return collected_result + 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, @@ -116,54 +198,18 @@ class NextBestViewEvaluator(Runner): self.current_epoch = meta["last_epoch"] self.current_iter = meta["last_iter"] - def save_checkpoint(self, is_last=False): - self.save(self.get_checkpoint_path(is_last)) - if not is_last: - Log.success(f"Checkpoint at epoch {self.current_epoch} saved to {self.get_checkpoint_path(is_last)}") - else: - meta = { - "last_epoch": self.current_epoch, - "last_iter": self.current_iter, - "time": str(datetime.now()) - } - checkpoint_root = os.path.join(self.experiment_path, namespace.Direcotry.CHECKPOINT_DIR_NAME) - file_path = os.path.join(checkpoint_root, "meta.json") - with open(file_path, "w") as f: - json.dump(meta, f) - def load_experiment(self, backup_name=None): super().load_experiment(backup_name) - if self.experiments_config["use_checkpoint"]: - self.current_epoch = self.experiments_config["epoch"] - self.load_checkpoint(is_last=(self.current_epoch == -1)) + 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) - ckpt_dir = os.path.join(str(self.experiment_path), namespace.Direcotry.CHECKPOINT_DIR_NAME) - os.makedirs(ckpt_dir) - tensorboard_dir = os.path.join(str(self.experiment_path), namespace.Direcotry.TENSORBOARD_DIR_NAME) - os.makedirs(tensorboard_dir) + def load(self, path): state_dict = torch.load(path) - if self.parallel: - self.pipeline.module.load_state_dict(state_dict) - else: - self.pipeline.load_state_dict(state_dict) - - def save(self, path): - if self.parallel: - state_dict = self.pipeline.module.state_dict() - else: - state_dict = self.pipeline.state_dict() - - for name, module in self.pipeline.named_modules(): - if module.__class__ in EXTERNAL_FRONZEN_MODULES: - if name in state_dict: - del state_dict[name] - - torch.save(state_dict, path) - + self.pipeline.load_state_dict(state_dict) def print_info(self): def print_dataset(dataset: BaseDataset): @@ -178,8 +224,6 @@ class NextBestViewEvaluator(Runner): Log.blue(f"{'+' + '-' * (table_size // 2)} Pipeline {'-' * (table_size // 2)}" + '+') Log.blue(self.pipeline) Log.blue(f"{'+' + '-' * (table_size // 2)} Datasets {'-' * (table_size // 2)}" + '+') - Log.blue("train dataset: ") - print_dataset(self.train_set) for i, test_set in enumerate(self.test_set_list): Log.blue(f"test dataset {i}: ") print_dataset(test_set) diff --git a/runners/strategy_generator.py b/runners/strategy_generator.py index d9f3baa..4b27e88 100644 --- a/runners/strategy_generator.py +++ b/runners/strategy_generator.py @@ -16,10 +16,10 @@ from utils.pts import PtsUtil class StrategyGenerator(Runner): def __init__(self, config): super().__init__(config) - self.load_experiment("generate") + self.load_experiment("generate_strategy") self.status_info = { "status_manager": status_manager, - "app_name": "generate", + "app_name": "generate_strategy", "runner_name": "strategy_generator" } self.to_specified_dir = ConfigManager.get("runner", "generate", "to_specified_dir") @@ -36,7 +36,7 @@ class StrategyGenerator(Runner): self.save_pts = ConfigManager.get("runner","generate","save_points") for dataset_idx in range(len(dataset_name_list)): dataset_name = dataset_name_list[dataset_idx] - status_manager.set_progress("generate", "strategy_generator", "dataset", dataset_idx, len(dataset_name_list)) + status_manager.set_progress("generate_strategy", "strategy_generator", "dataset", dataset_idx, len(dataset_name_list)) root_dir = ConfigManager.get("datasets", dataset_name, "root_dir") model_dir = ConfigManager.get("datasets", dataset_name, "model_dir") scene_name_list = os.listdir(root_dir) @@ -44,10 +44,10 @@ class StrategyGenerator(Runner): total = len(scene_name_list) for scene_name in scene_name_list: Log.info(f"({dataset_name})Processing [{cnt}/{total}]: {scene_name}") - status_manager.set_progress("generate", "strategy_generator", "scene", cnt, total) + status_manager.set_progress("generate_strategy", "strategy_generator", "scene", cnt, total) diag = DataLoadUtil.get_bbox_diag(model_dir, scene_name) voxel_threshold = diag*0.02 - status_manager.set_status("generate", "strategy_generator", "voxel_threshold", voxel_threshold) + status_manager.set_status("generate_strategy", "strategy_generator", "voxel_threshold", voxel_threshold) output_label_path = DataLoadUtil.get_label_path(root_dir, scene_name) if os.path.exists(output_label_path) and not self.overwrite: Log.info(f"Scene <{scene_name}> Already Exists, Skip") @@ -58,8 +58,8 @@ class StrategyGenerator(Runner): except Exception as e: Log.error(f"Scene <{scene_name}> Failed, Error: {e}") cnt += 1 - status_manager.set_progress("generate", "strategy_generator", "scene", total, total) - status_manager.set_progress("generate", "strategy_generator", "dataset", len(dataset_name_list), len(dataset_name_list)) + status_manager.set_progress("generate_strategy", "strategy_generator", "scene", total, total) + status_manager.set_progress("generate_strategy", "strategy_generator", "dataset", len(dataset_name_list), len(dataset_name_list)) def create_experiment(self, backup_name=None): super().create_experiment(backup_name) @@ -70,7 +70,7 @@ class StrategyGenerator(Runner): super().load_experiment(backup_name) def generate_sequence(self, root, model_dir, scene_name, voxel_threshold, overlap_threshold): - status_manager.set_status("generate", "strategy_generator", "scene", scene_name) + status_manager.set_status("generate_strategy", "strategy_generator", "scene", scene_name) frame_num = DataLoadUtil.get_scene_seq_length(root, scene_name) model_points_normals = DataLoadUtil.load_points_normals(root, scene_name) model_pts = model_points_normals[:,:3] @@ -81,7 +81,7 @@ class StrategyGenerator(Runner): for frame_idx in range(frame_num): path = DataLoadUtil.get_path(root, scene_name, frame_idx) cam_params = DataLoadUtil.load_cam_info(path, binocular=True) - status_manager.set_progress("generate", "strategy_generator", "loading frame", frame_idx, frame_num) + status_manager.set_progress("generate_strategy", "strategy_generator", "loading frame", frame_idx, frame_num) point_cloud = DataLoadUtil.get_target_point_cloud_world_from_path(path, binocular=True) #display_table = None #DataLoadUtil.get_target_point_cloud_world_from_path(path, binocular=True, target_mask_label=()) #TODO sampled_point_cloud = ReconstructionUtil.filter_points(point_cloud, model_points_normals, cam_pose=cam_params["cam_to_world"], voxel_size=voxel_threshold, theta=self.filter_degree) @@ -92,7 +92,7 @@ class StrategyGenerator(Runner): os.makedirs(pts_dir) np.savetxt(os.path.join(pts_dir, f"{frame_idx}.txt"), sampled_point_cloud) pts_list.append(sampled_point_cloud) - status_manager.set_progress("generate", "strategy_generator", "loading frame", frame_num, frame_num) + status_manager.set_progress("generate_strategy", "strategy_generator", "loading frame", frame_num, frame_num) limited_useful_view, _, best_combined_pts = ReconstructionUtil.compute_next_best_view_sequence_with_overlap(down_sampled_model_pts, pts_list, threshold=voxel_threshold, overlap_threshold=overlap_threshold, status_info=self.status_info) data_pairs = self.generate_data_pairs(limited_useful_view) @@ -102,7 +102,7 @@ class StrategyGenerator(Runner): "max_coverage_rate": limited_useful_view[-1][1] } - status_manager.set_status("generate", "strategy_generator", "max_coverage_rate", limited_useful_view[-1][1]) + status_manager.set_status("generate_strategy", "strategy_generator", "max_coverage_rate", limited_useful_view[-1][1]) Log.success(f"Scene <{scene_name}> Finished, Max Coverage Rate: {limited_useful_view[-1][1]}, Best Sequence length: {len(limited_useful_view)}") output_label_path = DataLoadUtil.get_label_path(root, scene_name) diff --git a/utils/pts.py b/utils/pts.py index 19d6e2a..b701af4 100644 --- a/utils/pts.py +++ b/utils/pts.py @@ -1,5 +1,6 @@ import numpy as np import open3d as o3d +import torch class PtsUtil: @@ -19,4 +20,9 @@ class PtsUtil: @staticmethod def random_downsample_point_cloud(point_cloud, num_points): idx = np.random.choice(len(point_cloud), num_points, replace=True) + return point_cloud[idx] + + @staticmethod + def random_downsample_point_cloud_tensor(point_cloud, num_points): + idx = torch.randint(0, len(point_cloud), (num_points,)) return point_cloud[idx] \ No newline at end of file diff --git a/utils/render.py b/utils/render.py new file mode 100644 index 0000000..b059a97 --- /dev/null +++ b/utils/render.py @@ -0,0 +1,51 @@ + +import os +import json +import subprocess +import tempfile +from utils.data_load import DataLoadUtil +from utils.reconstruction import ReconstructionUtil +from utils.pts import PtsUtil +class RenderUtil: + + @staticmethod + def render_pts(cam_pose, scene_path,script_path, model_points_normals, voxel_threshold=0.005, filter_degree=75, nO_to_nL_pose=None, require_full_scene=False): + nO_to_world_pose = cam_pose.cpu().numpy() @ nO_to_nL_pose + nO_to_world_pose = DataLoadUtil.cam_pose_transformation(nO_to_world_pose) + + + with tempfile.TemporaryDirectory() as temp_dir: + params = { + "cam_pose": nO_to_world_pose.tolist(), + "scene_path": scene_path + } + params_data_path = os.path.join(temp_dir, "params.json") + with open(params_data_path, 'w') as f: + json.dump(params, f) + result = subprocess.run([ + 'blender', '-b', '-P', script_path, '--', temp_dir + ], capture_output=True, text=True) + if result.returncode != 0: + print("Blender script failed:") + print(result.stderr) + return None + path = os.path.join(temp_dir, "tmp") + # ------ Debug Start ------ + import ipdb;ipdb.set_trace() + # ------ Debug End ------ + + point_cloud = DataLoadUtil.get_target_point_cloud_world_from_path(path, binocular=True) + cam_params = DataLoadUtil.load_cam_info(path, binocular=True) + filtered_point_cloud = ReconstructionUtil.filter_points(point_cloud, model_points_normals, cam_pose=cam_params["cam_to_world"], voxel_size=voxel_threshold, theta=filter_degree) + + full_scene_point_cloud = None + if require_full_scene: + depth_L, depth_R = DataLoadUtil.load_depth(path, cam_params['near_plane'], cam_params['far_plane'], binocular=True) + point_cloud_L = DataLoadUtil.get_point_cloud(depth_L, cam_params['cam_intrinsic'], cam_params['cam_to_world'])['points_world'] + point_cloud_R = DataLoadUtil.get_point_cloud(depth_R, cam_params['cam_intrinsic'], cam_params['cam_to_world_R'])['points_world'] + + point_cloud_L = PtsUtil.random_downsample_point_cloud(point_cloud_L, 65536) + point_cloud_R = PtsUtil.random_downsample_point_cloud(point_cloud_R, 65536) + full_scene_point_cloud = DataLoadUtil.get_overlapping_points(point_cloud_L, point_cloud_R) + + return filtered_point_cloud, full_scene_point_cloud \ No newline at end of file