diff --git a/Readme.md b/Readme.md new file mode 100644 index 0000000..6ad7dfc --- /dev/null +++ b/Readme.md @@ -0,0 +1,192 @@ +# Next Best View for Reconstruction + +## 1. Setup Environment +### 1.1 Install Main Project +```bash +mkdir nbv_rec +cd nbv_rec +git clone https://git.hofee.top/hofee/nbv_reconstruction.git +``` +### 1.2 Install PytorchBoot +the environment is based on PytorchBoot, clone and install it from [PytorchBoot](https://git.hofee.top/hofee/PyTorchBoot.git) +```bash +git clone https://git.hofee.top/hofee/PyTorchBoot.git +cd PyTorchBoot +pip install . +cd .. +``` +### 1.3 Install Blender (Optional) +If you want to render your own dataset as described in [section 2. Render Datasets](#2-render-datasets), you'll need to install Blender version 4.0 from [Blender Release](https://download.blender.org/release/Blender4.0/). Here is an example of installing Blender on Ubuntu: +```bash +wget https://download.blender.org/release/Blender4.0/blender-4.0.2-linux-x64.tar.xz +tar -xvf blender-4.0.2-linux-x64.tar.xz +``` +If blender is not in your PATH, you can add it by: +```bash +export PATH=$PATH:/path/to/blender/blender-4.0.2-linux-x64 +``` +To run the blender script, you need to install the `pyyaml` and `scipy` package into your blender python environment. Run the following command to print the python path of your blender: +```bash +./blender -b --python-expr "import sys; print(sys.executable)" +``` +Then copy the python path `/path/to/blender_python` shown in the output and run the following command to install the packages: +```bash +/path/to/blender_python -m pip install pyyaml scipy +``` +### 1.4 Install Blender Render Script (Optional) +Clone the script from [nbv_rec_blender_render](https://git.hofee.top/hofee/nbv_rec_blender_render.git) and rename it to `blender`: +```bash +git clone https://git.hofee.top/hofee/nbv_rec_blender_render.git +mv nbv_rec_blender_render blender +``` + +### 1.5 Check Dependencies +Switch to the project root directory and run `pytorch-boot scan` or `ptb scan` to check if all dependencies are installed: +```bash +cd nbv_reconstruction +pytorch-boot scan +# or +ptb scan +``` +If you see project structure information in the output, it means all dependencies are correctly installed. Otherwise, you may need to run `pip install xxx` to install the missing packages. + +## 2. Render Datasets (Optional) +### 2.1 Download Object Mesh Models +Download the mesh models divided into three parts from: + - [object_meshes_part1.zip](None) + - [object_meshes_part2.zip](https://pan.baidu.com/s/1pBPhrFtBwEGp1g4vwsLIxA?pwd=1234) + - [object_meshes_part3.zip](https://pan.baidu.com/s/1peE8HqFFL0qNFhM5OC69gA?pwd=1234) + +or download the whole dataset from [object_meshes.zip](https://pan.baidu.com/s/1ilWWgzg_l7_pPBv64eSgzA?pwd=1234) + +Download the table model from [table.obj](https://pan.baidu.com/s/1sjjiID25Es_kmcdUIjU_Dw?pwd=1234) + +### 2.2 Set Render Configurations +Open file `configs/local/view_generate_config.yaml` and modify the parameters to fit your needs. You are required to at least set the following parameters in `runner-generate`: + - `object_dir`: the directory of the downloaded object mesh models + - `output_dir`: the directory to save the rendered dataset + - `table_model_path`: the path of the downloaded table model + +### 2.3 Render Dataset + +There are two ways to render the dataset: + +#### 2.3.1 Render with Visual Monitoring + +If you want to visually monitor the rendering progress and machine resource usage: + +1. In the terminal, run: + ``` + ptb ui + ``` +2. Open your browser and visit http://localhost:5000 +3. Navigate to `Project Dashboard - Project Structure - Applications - generate_view` +4. Click the `Run` button to execute the rendering script + +#### 2.3.2 Render in Terminal + +If you don't need visual monitoring and prefer to run the rendering process directly in the terminal, simply run: + +``` +ptb run generate_view +``` + +This command will start the rendering process without launching the UI. + +## 3. Preprocess + +⚠️ The preprocessing code is currently not managed by `PytorchBoot`. To run the preprocessing: + +1. Open the `./preprocess/preprocessor.py` file. +2. Locate the `if __name__ == "__main__":` block at the bottom of the file. +3. Specify the dataset folder by setting `root = "path/to/your/dataset"`. +4. Run the preprocessing script directly: + + ``` + python ./preprocess/preprocessor.py + ``` + +This will preprocess the data in the specified dataset folder. + +## 4. Generate Strategy Label + +### 4.1 Set Configuration + +Open the file `configs/local/strategy_generate_config.yaml` and modify the parameters to fit your needs. You are required to at least set the following parameter: + +- `datasets.OmniObject3d.root_dir`: the directory of your dataset + +### 4.2 Generate Strategy Label + +There are two ways to generate the strategy label: + +#### 4.2.1 Generate with Visual Monitoring + +If you want to visually monitor the generation progress and machine resource usage: + +1. In the terminal, run: + ``` + ptb ui + ``` +2. Open your browser and visit http://localhost:5000 +3. Navigate to Project Dashboard - Project Structure - Applications - generate_strategy +4. Click the `Run` button to execute the generation script + +#### 4.2.2 Generate in Terminal + +If you don't need visual monitoring and prefer to run the generation process directly in the terminal, simply run: + +``` +ptb run generate_strategy +``` + +This command will start the strategy label generation process without launching the UI. + +## 5. Train + +### 5.1 Set Configuration + +Open the file `configs/local/train_config.yaml` and modify the parameters to fit your needs. You are required to at least set the following parameters in the `experiment` section: + +```yaml +experiment: + name: your_experiment_name + root_dir: path/to/your/experiment_dir + use_checkpoint: False # if True, the checkpoint will be loaded + epoch: 600 # specific epoch to load, -1 stands for last epoch + max_epochs: 5000 # maximum epochs to train + save_checkpoint_interval: 1 # save checkpoint interval + test_first: True # if True, test process will be performed before training at each epoch +``` + +Adjust these parameters according to your training requirements. + + +### 5.2 Start Training + +There are two ways to start the training process: + +#### 5.2.1 Train with Visual Monitoring + +If you want to visually monitor the training progress and machine resource usage: + +1. In the terminal, run: + ``` + ptb ui + ``` +2. Open your browser and visit http://localhost:5000 +3. Navigate to Project Dashboard - Project Structure - Applications - train +4. Click the `Run` button to start the training process + +#### 5.2.2 Train in Terminal + +If you don't need visual monitoring and prefer to run the training process directly in the terminal, simply run: + +``` +ptb run train +``` + +This command will start the training process without launching the UI. + +## 6. Evaluation +... diff --git a/TODO.md b/TODO.md deleted file mode 100644 index bc87a3d..0000000 --- a/TODO.md +++ /dev/null @@ -1,22 +0,0 @@ -# TODO -## 预处理数据 -### 1. 生成view阶段 -**input**: 物体mesh - -### 2. 生成label阶段 -**input**: 目标物体点云、目标物体点云法线、桌面扫描点、被拍到的桌面扫描点 - -**可以删掉的数据**: mask、normal - -### 3. 训练阶段 -**input**: 完整点云、pose、label - -**可以删掉的数据**：depth - -### view生成后 -预处理目标物体点云、目标物体点云法线、桌面扫描点、被拍到的桌面扫描点、完整点云 - -删除depth、mask、normal - -### label生成后 -只上传：完整点云、pose、label \ No newline at end of file diff --git a/app_generate_strategy.py b/app_generate_strategy.py index 9625297..28905e5 100644 --- a/app_generate_strategy.py +++ b/app_generate_strategy.py @@ -5,5 +5,5 @@ from runners.strategy_generator import StrategyGenerator class DataGenerateApp: @staticmethod def start(): - StrategyGenerator("configs/server/server_strategy_generate_config.yaml").run() + StrategyGenerator("configs/local/strategy_generate_config.yaml").run() \ No newline at end of file diff --git a/configs/local/strategy_generate_config.yaml b/configs/local/strategy_generate_config.yaml index 7b88d9c..cf444b0 100644 --- a/configs/local/strategy_generate_config.yaml +++ b/configs/local/strategy_generate_config.yaml @@ -12,14 +12,9 @@ runner: generate: voxel_threshold: 0.003 - soft_overlap_threshold: 0.3 - hard_overlap_threshold: 0.6 - filter_degree: 75 - to_specified_dir: True # if True, output_dir is used, otherwise, root_dir is used - save_points: True - load_points: True - save_best_combined_points: False - save_mesh: True + overlap_area_threshold: 25 + compute_with_normal: False + scan_points_threshold: 10 overwrite: False seq_num: 15 dataset_list: @@ -27,11 +22,8 @@ runner: datasets: OmniObject3d: - #"/media/hofee/data/data/temp_output" - root_dir: /media/hofee/repository/full_data_output - model_dir: /media/hofee/data/data/scaled_object_meshes + root_dir: C:\\Document\\Local Project\\nbv_rec\\nbv_reconstruction\\temp from: 0 - to: -1 # -1 means end - #output_dir: "/media/hofee/data/data/label_output" + to: 1 # -1 means end diff --git a/configs/local/view_generate_config.yaml b/configs/local/view_generate_config.yaml index e2f8779..b58d672 100644 --- a/configs/local/view_generate_config.yaml +++ b/configs/local/view_generate_config.yaml @@ -7,12 +7,12 @@ runner: name: debug root_dir: experiments generate: - port: 5004 - from: 0 - to: 1 # -1 means all - object_dir: H:\\AI\\Datasets\\scaled_object_box_meshes - table_model_path: "H:\\AI\\Datasets\\table.obj" - output_dir: C:\\Document\\Local Project\\nbv_rec\\nbv_reconstruction\\temp + port: 5002 + from: 600 + to: -1 # -1 means all + object_dir: /media/hofee/data/data/object_meshes_part1 + table_model_path: "/media/hofee/data/data/others/table.obj" + output_dir: /media/hofee/repository/data_part_1 binocular_vision: true plane_size: 10 max_views: 512 diff --git a/configs/server/server_strategy_generate_config.yaml b/configs/server/server_strategy_generate_config.yaml deleted file mode 100644 index ad8e161..0000000 --- a/configs/server/server_strategy_generate_config.yaml +++ /dev/null @@ -1,37 +0,0 @@ - -runner: - general: - seed: 0 - device: cpu - cuda_visible_devices: "0,1,2,3,4,5,6,7" - - - experiment: - name: debug - root_dir: "experiments" - - generate: - voxel_threshold: 0.003 - soft_overlap_threshold: 0.3 - hard_overlap_threshold: 0.6 - filter_degree: 75 - to_specified_dir: True # if True, output_dir is used, otherwise, root_dir is used - save_points: True - load_points: True - save_best_combined_points: False - save_mesh: True - overwrite: False - seq_num: 15 - dataset_list: - - OmniObject3d - -datasets: - OmniObject3d: - #"/media/hofee/data/data/temp_output" - root_dir: /data/hofee/data/packed_preprocessed_data - model_dir: /media/hofee/data/data/scaled_object_meshes - from: 0 - to: -1 # -1 means end - #output_dir: "/media/hofee/data/data/label_output" - - diff --git a/preprocess/clean_preprocessed_data.py b/preprocess/clean_preprocessed_data.py new file mode 100644 index 0000000..fb9b7aa --- /dev/null +++ b/preprocess/clean_preprocessed_data.py @@ -0,0 +1,43 @@ +import os +import shutil + +def clean_scene_data(root, scene): + # 清理目标点云数据 + pts_dir = os.path.join(root, scene, "pts") + if os.path.exists(pts_dir): + shutil.rmtree(pts_dir) + print(f"已删除 {pts_dir}") + + # 清理法线数据 + nrm_dir = os.path.join(root, scene, "nrm") + if os.path.exists(nrm_dir): + shutil.rmtree(nrm_dir) + print(f"已删除 {nrm_dir}") + + # 清理扫描点索引数据 + scan_points_indices_dir = os.path.join(root, scene, "scan_points_indices") + if os.path.exists(scan_points_indices_dir): + shutil.rmtree(scan_points_indices_dir) + print(f"已删除 {scan_points_indices_dir}") + + # 删除扫描点数据文件 + scan_points_file = os.path.join(root, scene, "scan_points.txt") + if os.path.exists(scan_points_file): + os.remove(scan_points_file) + print(f"已删除 {scan_points_file}") + +def clean_all_scenes(root, scene_list): + for idx, scene in enumerate(scene_list): + print(f"正在清理场景 {scene} ({idx+1}/{len(scene_list)})") + clean_scene_data(root, scene) + +if __name__ == "__main__": + root = r"c:\Document\Local Project\nbv_rec\nbv_reconstruction\temp" + scene_list = os.listdir(root) + from_idx = 0 + to_idx = len(scene_list) + print(f"正在清理场景 {scene_list[from_idx:to_idx]}") + + clean_all_scenes(root, scene_list[from_idx:to_idx]) + print("清理完成") + diff --git a/preprocess/preprocessor.py b/preprocess/preprocessor.py index ad8b166..e4ad251 100644 --- a/preprocess/preprocessor.py +++ b/preprocess/preprocessor.py @@ -9,8 +9,6 @@ from utils.reconstruction import ReconstructionUtil from utils.data_load import DataLoadUtil from utils.pts import PtsUtil -# scan shoe 536 - def save_np_pts(path, pts: np.ndarray, file_type="txt"): if file_type == "txt": np.savetxt(path, pts) @@ -23,6 +21,12 @@ def save_target_points(root, scene, frame_idx, target_points: np.ndarray, file_t if not os.path.exists(os.path.join(root,scene, "pts")): os.makedirs(os.path.join(root,scene, "pts")) save_np_pts(pts_path, target_points, file_type) + +def save_target_normals(root, scene, frame_idx, target_normals: np.ndarray, file_type="txt"): + pts_path = os.path.join(root,scene, "nrm", f"{frame_idx}.{file_type}") + if not os.path.exists(os.path.join(root,scene, "nrm")): + os.makedirs(os.path.join(root,scene, "nrm")) + save_np_pts(pts_path, target_normals, file_type) def save_scan_points_indices(root, scene, frame_idx, scan_points_indices: np.ndarray, file_type="txt"): indices_path = os.path.join(root,scene, "scan_points_indices", f"{frame_idx}.{file_type}") @@ -87,8 +91,8 @@ def get_scan_points_indices(scan_points, mask, display_table_mask_label, cam_int def save_scene_data(root, scene, scene_idx=0, scene_total=1,file_type="txt"): ''' configuration ''' - target_mask_label = (0, 255, 0, 255) - display_table_mask_label=(0, 0, 255, 255) + target_mask_label = (0, 255, 0) + display_table_mask_label=(0, 0, 255) random_downsample_N = 32768 voxel_size=0.003 filter_degree = 75 @@ -137,7 +141,7 @@ def save_scene_data(root, scene, scene_idx=0, scene_total=1,file_type="txt"): has_points = target_points.shape[0] > 0 if has_points: - target_points = PtsUtil.filter_points( + target_points, target_normals = PtsUtil.filter_points( target_points, sampled_target_normal_L, cam_info["cam_to_world"], theta_limit = filter_degree, z_range=(min_z, max_z) ) @@ -149,8 +153,10 @@ def save_scene_data(root, scene, scene_idx=0, scene_total=1,file_type="txt"): if not has_points: target_points = np.zeros((0, 3)) + target_normals = np.zeros((0, 3)) save_target_points(root, scene, frame_id, target_points, file_type=file_type) + save_target_normals(root, scene, frame_id, target_normals, file_type=file_type) save_scan_points_indices(root, scene, frame_id, scan_points_indices, file_type=file_type) save_scan_points(root, scene, scan_points) # The "done" flag of scene preprocess @@ -158,17 +164,10 @@ def save_scene_data(root, scene, scene_idx=0, scene_total=1,file_type="txt"): if __name__ == "__main__": #root = "/media/hofee/repository/new_data_with_normal" - root = r"C:\\Document\\Local Project\\nbv_rec\\nbv_reconstruction\\temp" - # list_path = r"/media/hofee/repository/full_list.txt" - # scene_list = [] - - # with open(list_path, "r") as f: - # for line in f: - # scene_list.append(line.strip()) + root = r"C:\Document\Datasets\nbv_rec_part2" scene_list = os.listdir(root) - from_idx = 0 # 1000 - to_idx = 1 # 1500 - print(scene_list) + from_idx = 600 # 1000 + to_idx = len(scene_list) # 1500 cnt = 0 @@ -176,6 +175,10 @@ if __name__ == "__main__": total = to_idx - from_idx for scene in scene_list[from_idx:to_idx]: start = time.time() + if os.path.exists(os.path.join(root, scene, "scan_points.txt")): + print(f"Scene {scene} has been processed") + cnt+=1 + continue save_scene_data(root, scene, cnt, total, file_type="npy") cnt+=1 end = time.time() diff --git a/runners/strategy_generator.py b/runners/strategy_generator.py index ca40f99..60a5cfc 100644 --- a/runners/strategy_generator.py +++ b/runners/strategy_generator.py @@ -22,20 +22,17 @@ class StrategyGenerator(Runner): "app_name": "generate_strategy", "runner_name": "strategy_generator" } - self.to_specified_dir = ConfigManager.get("runner", "generate", "to_specified_dir") - self.save_best_combined_pts = ConfigManager.get("runner", "generate", "save_best_combined_points") - self.save_mesh = ConfigManager.get("runner", "generate", "save_mesh") - self.load_pts = ConfigManager.get("runner", "generate", "load_points") - self.filter_degree = ConfigManager.get("runner", "generate", "filter_degree") self.overwrite = ConfigManager.get("runner", "generate", "overwrite") - self.save_pts = ConfigManager.get("runner","generate","save_points") self.seq_num = ConfigManager.get("runner","generate","seq_num") + self.overlap_area_threshold = ConfigManager.get("runner","generate","overlap_area_threshold") + self.compute_with_normal = ConfigManager.get("runner","generate","compute_with_normal") + self.scan_points_threshold = ConfigManager.get("runner","generate","scan_points_threshold") def run(self): dataset_name_list = ConfigManager.get("runner", "generate", "dataset_list") - voxel_threshold, soft_overlap_threshold, hard_overlap_threshold = ConfigManager.get("runner","generate","voxel_threshold"), ConfigManager.get("runner","generate","soft_overlap_threshold"), ConfigManager.get("runner","generate","hard_overlap_threshold") + voxel_threshold = ConfigManager.get("runner","generate","voxel_threshold") for dataset_idx in range(len(dataset_name_list)): dataset_name = dataset_name_list[dataset_idx] status_manager.set_progress("generate_strategy", "strategy_generator", "dataset", dataset_idx, len(dataset_name_list)) @@ -57,7 +54,7 @@ class StrategyGenerator(Runner): cnt += 1 continue - self.generate_sequence(root_dir, scene_name,voxel_threshold, soft_overlap_threshold, hard_overlap_threshold) + self.generate_sequence(root_dir, scene_name,voxel_threshold) cnt += 1 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)) @@ -70,28 +67,34 @@ class StrategyGenerator(Runner): def load_experiment(self, backup_name=None): super().load_experiment(backup_name) - def generate_sequence(self, root, scene_name, voxel_threshold, soft_overlap_threshold, hard_overlap_threshold): + def generate_sequence(self, root, scene_name, voxel_threshold): 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] - down_sampled_model_pts = PtsUtil.voxel_downsample_point_cloud(model_pts, voxel_threshold) + down_sampled_model_pts, idx = PtsUtil.voxel_downsample_point_cloud(model_pts, voxel_threshold, require_idx=True) + down_sampled_model_nrm = model_points_normals[idx, 3:] pts_list = [] + nrm_list = [] scan_points_indices_list = [] non_zero_cnt = 0 for frame_idx in range(frame_num): status_manager.set_progress("generate_strategy", "strategy_generator", "loading frame", frame_idx, frame_num) pts_path = os.path.join(root,scene_name, "pts", f"{frame_idx}.npy") + nrm_path = os.path.join(root,scene_name, "nrm", f"{frame_idx}.npy") idx_path = os.path.join(root,scene_name, "scan_points_indices", f"{frame_idx}.npy") - point_cloud = np.load(pts_path) - sampled_point_cloud = PtsUtil.voxel_downsample_point_cloud(point_cloud, voxel_threshold) + pts = np.load(pts_path) + if pts.shape[0] == 0: + nrm = np.zeros((0,3)) + else: + nrm = np.load(nrm_path) indices = np.load(idx_path) - pts_list.append(sampled_point_cloud) - + pts_list.append(pts) + nrm_list.append(nrm) scan_points_indices_list.append(indices) - if sampled_point_cloud.shape[0] > 0: + if pts.shape[0] > 0: non_zero_cnt += 1 status_manager.set_progress("generate_strategy", "strategy_generator", "loading frame", frame_num, frame_num) @@ -99,7 +102,7 @@ class StrategyGenerator(Runner): init_view_list = [] idx = 0 while len(init_view_list) < seq_num and idx < len(pts_list): - if pts_list[idx].shape[0] > 100: + if pts_list[idx].shape[0] > 50: init_view_list.append(idx) idx += 1 @@ -108,8 +111,13 @@ class StrategyGenerator(Runner): for init_view in init_view_list: status_manager.set_progress("generate_strategy", "strategy_generator", "computing sequence", seq_idx, len(init_view_list)) start = time.time() - limited_useful_view, _, _ = ReconstructionUtil.compute_next_best_view_sequence_with_overlap(down_sampled_model_pts, pts_list, scan_points_indices_list = scan_points_indices_list,init_view=init_view, - threshold=voxel_threshold, soft_overlap_threshold=soft_overlap_threshold, hard_overlap_threshold= hard_overlap_threshold, scan_points_threshold=10, status_info=self.status_info) + + if not self.compute_with_normal: + limited_useful_view, _, _ = ReconstructionUtil.compute_next_best_view_sequence(down_sampled_model_pts, pts_list, scan_points_indices_list = scan_points_indices_list,init_view=init_view, + threshold=voxel_threshold, scan_points_threshold=self.scan_points_threshold, overlap_area_threshold=self.overlap_area_threshold, status_info=self.status_info) + else: + limited_useful_view, _, _ = ReconstructionUtil.compute_next_best_view_sequence_with_normal(down_sampled_model_pts, down_sampled_model_nrm, pts_list, nrm_list, scan_points_indices_list = scan_points_indices_list,init_view=init_view, + threshold=voxel_threshold, scan_points_threshold=self.scan_points_threshold, overlap_area_threshold=self.overlap_area_threshold, status_info=self.status_info) end = time.time() print(f"Time: {end-start}") data_pairs = self.generate_data_pairs(limited_useful_view) diff --git a/runners/view_generator.py b/runners/view_generator.py index 2e28c44..634ccbf 100644 --- a/runners/view_generator.py +++ b/runners/view_generator.py @@ -9,7 +9,7 @@ class ViewGenerator(Runner): self.config_path = config_path def run(self): - result = subprocess.run(['blender', '-b', '-P', '../blender/run_blender.py', '--', self.config_path]) + result = subprocess.run(['/home/hofee/blender-4.0.2-linux-x64/blender', '-b', '-P', '../blender/run_blender.py', '--', self.config_path]) print() def create_experiment(self, backup_name=None): diff --git a/utils/data_load.py b/utils/data_load.py index b89a60d..8d6e8e3 100644 --- a/utils/data_load.py +++ b/utils/data_load.py @@ -14,23 +14,16 @@ class DataLoadUtil: @staticmethod def load_exr_image(file_path): - # 打开 EXR 文件 exr_file = OpenEXR.InputFile(file_path) - - # 获取 EXR 文件的头部信息，包括尺寸 header = exr_file.header() dw = header['dataWindow'] width = dw.max.x - dw.min.x + 1 height = dw.max.y - dw.min.y + 1 - - # 定义通道，通常法线图像是 RGB float_channels = ['R', 'G', 'B'] - - # 读取 EXR 文件中的每个通道并转化为浮点数数组 img_data = [] for channel in float_channels: - channel_data = exr_file.channel(channel, Imath.PixelType(Imath.PixelType.FLOAT)) - img_data.append(np.frombuffer(channel_data, dtype=np.float32).reshape((height, width))) + channel_data = exr_file.channel(channel) + img_data.append(np.frombuffer(channel_data, dtype=np.float16).reshape((height, width))) # 将各通道组合成一个 (height, width, 3) 的 RGB 图像 img = np.stack(img_data, axis=-1) @@ -143,8 +136,8 @@ class DataLoadUtil: if binocular and not left_only: def clean_mask(mask_image): - green = [0, 255, 0, 255] - red = [255, 0, 0, 255] + green = [0, 255, 0] + red = [255, 0, 0] threshold = 2 mask_image = np.where( np.abs(mask_image - green) <= threshold, green, mask_image diff --git a/utils/pts.py b/utils/pts.py index 30f8860..7b640df 100644 --- a/utils/pts.py +++ b/utils/pts.py @@ -5,10 +5,17 @@ import torch class PtsUtil: @staticmethod - def voxel_downsample_point_cloud(point_cloud, voxel_size=0.005): + def voxel_downsample_point_cloud(point_cloud, voxel_size=0.005, require_idx=False): voxel_indices = np.floor(point_cloud / voxel_size).astype(np.int32) - unique_voxels = np.unique(voxel_indices, axis=0, return_inverse=True) - return unique_voxels[0]*voxel_size + if require_idx: + _, 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, idx_unique + else: + unique_voxels = np.unique(voxel_indices, axis=0, return_inverse=True) + return unique_voxels[0]*voxel_size @staticmethod def random_downsample_point_cloud(point_cloud, num_points, require_idx=False): @@ -84,14 +91,14 @@ class PtsUtil: theta = np.arccos(cos_theta) * 180 / np.pi idx = theta < theta_limit filtered_sampled_points = points[idx] - + filtered_normals = normals[idx] """ filter with z range """ points_cam = PtsUtil.transform_point_cloud(filtered_sampled_points, np.linalg.inv(cam_pose)) idx = (points_cam[:, 2] > z_range[0]) & (points_cam[:, 2] < z_range[1]) z_filtered_points = filtered_sampled_points[idx] - - return z_filtered_points[:, :3] + z_filtered_normals = filtered_normals[idx] + return z_filtered_points[:, :3], z_filtered_normals @staticmethod def point_to_hash(point, voxel_size): diff --git a/utils/reconstruction.py b/utils/reconstruction.py index fed95e1..ac359aa 100644 --- a/utils/reconstruction.py +++ b/utils/reconstruction.py @@ -8,16 +8,23 @@ class ReconstructionUtil: def compute_coverage_rate(target_point_cloud, combined_point_cloud, threshold=0.01): kdtree = cKDTree(combined_point_cloud) distances, _ = kdtree.query(target_point_cloud) - covered_points_num = np.sum(distances < threshold) + covered_points_num = np.sum(distances < threshold*2) coverage_rate = covered_points_num / target_point_cloud.shape[0] return coverage_rate, covered_points_num + @staticmethod def compute_coverage_rate_with_normal(target_point_cloud, combined_point_cloud, target_normal, combined_normal, threshold=0.01, normal_threshold=0.1): kdtree = cKDTree(combined_point_cloud) distances, indices = kdtree.query(target_point_cloud) - is_covered_by_distance = distances < threshold + is_covered_by_distance = distances < threshold*2 normal_dots = np.einsum('ij,ij->i', target_normal, combined_normal[indices]) is_covered_by_normal = normal_dots > normal_threshold + + pts_nrm_target = np.hstack([target_point_cloud, target_normal]) + np.savetxt("pts_nrm_target.txt", pts_nrm_target) + pts_nrm_combined = np.hstack([combined_point_cloud, combined_normal]) + np.savetxt("pts_nrm_combined.txt", pts_nrm_combined) + import ipdb; ipdb.set_trace() covered_points_num = np.sum(is_covered_by_distance & is_covered_by_normal) coverage_rate = covered_points_num / target_point_cloud.shape[0] @@ -25,15 +32,14 @@ class ReconstructionUtil: @staticmethod - def compute_overlap_rate(new_point_cloud, combined_point_cloud, threshold=0.01): + def check_overlap(new_point_cloud, combined_point_cloud, overlap_area_threshold=25, voxel_size=0.01): kdtree = cKDTree(combined_point_cloud) distances, _ = kdtree.query(new_point_cloud) - overlapping_points = np.sum(distances < threshold) - if new_point_cloud.shape[0] == 0: - overlap_rate = 0 - else: - overlap_rate = overlapping_points / new_point_cloud.shape[0] - return overlap_rate + overlapping_points = np.sum(distances < voxel_size*2) + cm = 0.01 + voxel_size_cm = voxel_size / cm + overlap_area = overlapping_points * voxel_size_cm * voxel_size_cm + return overlap_area > overlap_area_threshold @staticmethod @@ -49,7 +55,7 @@ class ReconstructionUtil: return new_added_points @staticmethod - def compute_next_best_view_sequence_with_overlap(target_point_cloud, point_cloud_list, scan_points_indices_list, threshold=0.01, soft_overlap_threshold=0.5, hard_overlap_threshold=0.7, init_view = 0, scan_points_threshold=5, status_info=None): + def compute_next_best_view_sequence(target_point_cloud, point_cloud_list, scan_points_indices_list, threshold=0.01, overlap_area_threshold=25, init_view = 0, scan_points_threshold=5, status_info=None): selected_views = [init_view] combined_point_cloud = point_cloud_list[init_view] history_indices = [scan_points_indices_list[init_view]] @@ -83,22 +89,16 @@ class ReconstructionUtil: if selected_views: new_scan_points_indices = scan_points_indices_list[view_index] if not ReconstructionUtil.check_scan_points_overlap(history_indices, new_scan_points_indices, scan_points_threshold): - overlap_threshold = hard_overlap_threshold + curr_overlap_area_threshold = overlap_area_threshold else: - overlap_threshold = soft_overlap_threshold - start = time.time() - overlap_rate = ReconstructionUtil.compute_overlap_rate(point_cloud_list[view_index],combined_point_cloud, threshold) - end = time.time() - # print(f"overlap_rate Time: {end-start}") - if overlap_rate < overlap_threshold: + curr_overlap_area_threshold = overlap_area_threshold * 0.5 + + if not ReconstructionUtil.check_overlap(point_cloud_list[view_index], combined_point_cloud, overlap_area_threshold = curr_overlap_area_threshold, voxel_size=threshold): continue - start = time.time() new_combined_point_cloud = np.vstack([combined_point_cloud, point_cloud_list[view_index]]) new_downsampled_combined_point_cloud = PtsUtil.voxel_downsample_point_cloud(new_combined_point_cloud,threshold) new_coverage, new_covered_num = ReconstructionUtil.compute_coverage_rate(downsampled_max_rec_pts, new_downsampled_combined_point_cloud, threshold) - end = time.time() - #print(f"compute_coverage_rate Time: {end-start}") coverage_increase = new_coverage - current_coverage if coverage_increase > best_coverage_increase: best_coverage_increase = coverage_increase @@ -107,6 +107,100 @@ class ReconstructionUtil: best_combined_point_cloud = new_downsampled_combined_point_cloud + if best_view is not None: + if best_coverage_increase <=1e-3 or best_covered_num - current_covered_num <= 5: + break + + selected_views.append(best_view) + best_rec_pts_num = best_combined_point_cloud.shape[0] + print(f"Current rec pts num: {curr_rec_pts_num}, Best rec pts num: {best_rec_pts_num}, Best cover pts: {best_covered_num}, Max rec pts num: {max_rec_pts_num}") + print(f"Current coverage: {current_coverage+best_coverage_increase}, Best coverage increase: {best_coverage_increase}, Max Real coverage: {max_real_rec_pts_coverage}") + current_covered_num = best_covered_num + curr_rec_pts_num = best_rec_pts_num + combined_point_cloud = best_combined_point_cloud + remaining_views.remove(best_view) + history_indices.append(scan_points_indices_list[best_view]) + current_coverage += best_coverage_increase + cnt_processed_view += 1 + if status_info is not None: + sm = status_info["status_manager"] + app_name = status_info["app_name"] + runner_name = status_info["runner_name"] + sm.set_status(app_name, runner_name, "current coverage", current_coverage) + sm.set_progress(app_name, runner_name, "processed view", cnt_processed_view, len(point_cloud_list)) + + view_sequence.append((best_view, current_coverage)) + + else: + break + if status_info is not None: + sm = status_info["status_manager"] + app_name = status_info["app_name"] + runner_name = status_info["runner_name"] + sm.set_progress(app_name, runner_name, "processed view", len(point_cloud_list), len(point_cloud_list)) + return view_sequence, remaining_views, combined_point_cloud + + @staticmethod + def compute_next_best_view_sequence_with_normal(target_point_cloud, target_normal, point_cloud_list, normal_list, scan_points_indices_list, threshold=0.01, overlap_area_threshold=25, init_view = 0, scan_points_threshold=5, status_info=None): + selected_views = [init_view] + combined_point_cloud = point_cloud_list[init_view] + combined_normal = normal_list[init_view] + history_indices = [scan_points_indices_list[init_view]] + + max_rec_pts = np.vstack(point_cloud_list) + max_rec_nrm = np.vstack(normal_list) + downsampled_max_rec_pts, idx = PtsUtil.voxel_downsample_point_cloud(max_rec_pts, threshold, require_idx=True) + downsampled_max_rec_nrm = max_rec_nrm[idx] + max_rec_pts_num = downsampled_max_rec_pts.shape[0] + try: + max_real_rec_pts_coverage, _ = ReconstructionUtil.compute_coverage_rate_with_normal(target_point_cloud, downsampled_max_rec_pts, target_normal, downsampled_max_rec_nrm, threshold) + except: + import ipdb; ipdb.set_trace() + + new_coverage, new_covered_num = ReconstructionUtil.compute_coverage_rate_with_normal(downsampled_max_rec_pts, combined_point_cloud, downsampled_max_rec_nrm, combined_normal, threshold) + current_coverage = new_coverage + current_covered_num = new_covered_num + + remaining_views = list(range(len(point_cloud_list))) + view_sequence = [(init_view, current_coverage)] + cnt_processed_view = 0 + remaining_views.remove(init_view) + curr_rec_pts_num = combined_point_cloud.shape[0] + + while remaining_views: + best_view = None + best_coverage_increase = -1 + best_combined_point_cloud = None + best_combined_normal = None + best_covered_num = 0 + + for view_index in remaining_views: + if point_cloud_list[view_index].shape[0] == 0: + continue + if selected_views: + new_scan_points_indices = scan_points_indices_list[view_index] + 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 + + if not ReconstructionUtil.check_overlap(point_cloud_list[view_index], combined_point_cloud, overlap_area_threshold = curr_overlap_area_threshold, voxel_size=threshold): + continue + + new_combined_point_cloud = np.vstack([combined_point_cloud, point_cloud_list[view_index]]) + new_combined_normal = np.vstack([combined_normal, normal_list[view_index]]) + new_downsampled_combined_point_cloud, idx = PtsUtil.voxel_downsample_point_cloud(new_combined_point_cloud,threshold, require_idx=True) + new_downsampled_combined_normal = new_combined_normal[idx] + new_coverage, new_covered_num = ReconstructionUtil.compute_coverage_rate_with_normal(downsampled_max_rec_pts, new_downsampled_combined_point_cloud, downsampled_max_rec_nrm, new_downsampled_combined_normal, threshold) + coverage_increase = new_coverage - current_coverage + if coverage_increase > best_coverage_increase: + best_coverage_increase = coverage_increase + best_view = view_index + best_covered_num = new_covered_num + best_combined_point_cloud = new_downsampled_combined_point_cloud + best_combined_normal = new_downsampled_combined_normal + + if best_view is not None: if best_coverage_increase <=1e-3 or best_covered_num - current_covered_num <= 5: break @@ -118,6 +212,7 @@ class ReconstructionUtil: current_covered_num = best_covered_num curr_rec_pts_num = best_rec_pts_num combined_point_cloud = best_combined_point_cloud + combined_normal = best_combined_normal remaining_views.remove(best_view) history_indices.append(scan_points_indices_list[best_view]) current_coverage += best_coverage_increase diff --git a/utils/vis.py b/utils/vis.py index e8df420..532a21a 100644 --- a/utils/vis.py +++ b/utils/vis.py @@ -47,6 +47,42 @@ class visualizeUtil: all_combined_pts = np.vstack(all_combined_pts) downsampled_all_pts = PtsUtil.voxel_downsample_point_cloud(all_combined_pts, 0.001) np.savetxt(os.path.join(output_dir, "all_combined_pts.txt"), downsampled_all_pts) + + @staticmethod + def save_seq_cam_pos_and_cam_axis(root, scene, frame_idx_list, output_dir): + all_cam_pos = [] + all_cam_axis = [] + for i in frame_idx_list: + path = DataLoadUtil.get_path(root, scene, i) + cam_info = DataLoadUtil.load_cam_info(path, binocular=True) + cam_pose = cam_info["cam_to_world"] + cam_pos = cam_pose[:3, 3] + cam_axis = cam_pose[:3, 2] + + num_samples = 10 + sample_points = [cam_pos + 0.02*t * cam_axis for t in range(num_samples)] + sample_points = np.array(sample_points) + + all_cam_pos.append(cam_pos) + all_cam_axis.append(sample_points) + + all_cam_pos = np.array(all_cam_pos) + all_cam_axis = np.array(all_cam_axis).reshape(-1, 3) + np.savetxt(os.path.join(output_dir, "seq_cam_pos.txt"), all_cam_pos) + np.savetxt(os.path.join(output_dir, "seq_cam_axis.txt"), all_cam_axis) + + @staticmethod + def save_seq_combined_pts(root, scene, frame_idx_list, output_dir): + all_combined_pts = [] + for i in frame_idx_list: + path = DataLoadUtil.get_path(root, scene, i) + pts = DataLoadUtil.load_from_preprocessed_pts(path,"npy") + if pts.shape[0] == 0: + continue + all_combined_pts.append(pts) + all_combined_pts = np.vstack(all_combined_pts) + downsampled_all_pts = PtsUtil.voxel_downsample_point_cloud(all_combined_pts, 0.001) + np.savetxt(os.path.join(output_dir, "seq_combined_pts.txt"), downsampled_all_pts) @staticmethod def save_target_mesh_at_world_space( @@ -120,18 +156,34 @@ class visualizeUtil: sampled_visualized_normal = np.array(sampled_visualized_normal).reshape(-1, 3) np.savetxt(os.path.join(output_dir, "target_pts.txt"), sampled_target_points) np.savetxt(os.path.join(output_dir, "target_normal.txt"), sampled_visualized_normal) - + + @staticmethod + def save_pts_nrm(pts_nrm, output_dir): + pts = pts_nrm[:, :3] + nrm = pts_nrm[:, 3:] + visualized_nrm = [] + num_samples = 10 + for i in range(len(pts)): + visualized_nrm.append(pts[i] + 0.02*t * nrm[i] for t in range(num_samples)) + visualized_nrm = np.array(visualized_nrm).reshape(-1, 3) + np.savetxt(os.path.join(output_dir, "nrm.txt"), visualized_nrm) + np.savetxt(os.path.join(output_dir, "pts.txt"), pts) # ------ Debug ------ if __name__ == "__main__": - root = r"/home/yan20/nbv_rec/project/franka_control/temp" + root = r"C:\Document\Local Project\nbv_rec\nbv_reconstruction\temp" model_dir = r"H:\\AI\\Datasets\\scaled_object_box_meshes" - scene = "cad_model_world" - output_dir = r"/home/yan20/nbv_rec/project/franka_control/temp/output" + scene = "box" + output_dir = r"C:\Document\Local Project\nbv_rec\nbv_reconstruction\test" - visualizeUtil.save_all_cam_pos_and_cam_axis(root, scene, output_dir) - visualizeUtil.save_all_combined_pts(root, scene, output_dir) - visualizeUtil.save_target_mesh_at_world_space(root, model_dir, scene) - #visualizeUtil.save_points_and_normals(root, scene,"10", output_dir, binocular=True) \ No newline at end of file + #visualizeUtil.save_all_cam_pos_and_cam_axis(root, scene, output_dir) + # visualizeUtil.save_all_combined_pts(root, scene, output_dir) + # visualizeUtil.save_seq_combined_pts(root, scene, [0, 121, 286, 175, 111,366,45,230,232,225,255,17,199,78,60], output_dir) + # visualizeUtil.save_seq_cam_pos_and_cam_axis(root, scene, [0, 121, 286, 175, 111,366,45,230,232,225,255,17,199,78,60], output_dir) + # visualizeUtil.save_target_mesh_at_world_space(root, model_dir, scene) + #visualizeUtil.save_points_and_normals(root, scene,"10", output_dir, binocular=True) + pts_nrm = np.loadtxt(r"C:\Document\Local Project\nbv_rec\nbv_reconstruction\pts_nrm_target.txt") + visualizeUtil.save_pts_nrm(pts_nrm, output_dir) +