58 Commits

Author SHA1 Message Date
2cd811c1b7 upd 2024-11-07 19:33:18 +08:00
hofee
5bcd0fc6e3 upd 2024-11-04 23:49:12 +08:00
hofee
2b7243d1be upd infernce 2024-11-04 17:17:54 +08:00
04d3a359e1 upd 2024-11-02 21:54:46 +00:00
287983277a global: debug inference 2024-11-01 22:51:16 +00:00
982a3b9b60 global: inference debug 2024-11-01 21:58:44 +00:00
ecd4cfa806 global: debug inference 2024-11-01 15:47:11 +00:00
985a08d89c global: upd inference 2024-11-01 08:43:13 +00:00
b221036e8b global: upd 2024-10-31 16:02:26 +00:00
097712c0ea global_only: ratio2 2024-10-30 15:58:32 +00:00
a954ed0998 global_only: ratio2 2024-10-30 15:49:59 +00:00
f5f8e4266f global_only: ratio 2024-10-30 15:49:11 +00:00
8a05b7883d global_only: train 2024-10-30 15:46:15 +00:00
e23697eb87 global_only: debug 2024-10-29 16:21:30 +00:00
2487039445 global_only: config 2024-10-29 12:18:51 +00:00
f533104e4a global_only: pipeline 2024-10-29 12:04:54 +00:00
a21538c90a global_only: dataset 2024-10-29 11:41:44 +00:00
872405e239 remove fps 2024-10-29 11:23:28 +00:00
b13e45bafc solve merge 2024-10-29 08:14:43 +00:00
63a246c0c8 debug new training 2024-10-28 19:15:48 +00:00
9e39c6c6c9 solve merge 2024-10-28 18:27:16 +00:00
3c9e2c8d12 solve merge 2024-10-28 18:25:53 +00:00
a883a31968 solve merge 2024-10-28 17:03:03 +00:00
49bcf203a8 update 2024-10-28 16:48:34 +00:00
hofee
1c443e533d add inference_server 2024-10-27 04:17:08 -05:00
hofee
3b9c966fd9 Merge branch 'master' of https://git.hofee.top/hofee/nbv_reconstruction 2024-10-26 03:24:18 -05:00
hofee
a41571e79c update 2024-10-26 03:24:01 -05:00
bd27226f0f solve merge 2024-10-25 14:40:26 +00:00
5c56dae24f upd 2024-10-24 20:19:23 +08:00
ebb1ab3c61 udp 2024-10-24 20:18:47 +08:00
hofee
a1226eb294 update normal in computing strategy 2024-10-23 11:13:18 -05:00
hofee
9d0119549e Merge branch 'master' of https://git.hofee.top/hofee/nbv_reconstruction 2024-10-23 02:59:18 -05:00
hofee
64891ef189 update normal strategy 2024-10-23 02:58:58 -05:00
75c70a9e59 fix no normal case 2024-10-23 14:54:53 +08:00
hofee
7e68259f6d update clean preprocess 2024-10-23 01:03:40 -05:00
64b22fd0f4 solve merge 2024-10-23 13:59:12 +08:00
b18c1591b7 load 16bit float 2024-10-23 13:57:45 +08:00
hofee
c55a398b6d update nrm 2024-10-23 00:47:28 -05:00
hofee
e25f7b3334 add save preprocessed normals 2024-10-23 00:42:18 -05:00
hofee
cd56d9ea58 update readme 2024-10-22 16:42:10 +08:00
hofee
d58c7980ed update 2024-10-22 16:41:02 +08:00
hofee
41eddda8d4 solve merge 2024-10-22 16:01:56 +08:00
hofee
ccec9b8e8a add readme.md 2024-10-22 16:01:11 +08:00
0f61e1d64d Merge branch 'master' of https://git.hofee.top/hofee/nbv_reconstruction 2024-10-21 07:33:40 +00:00
9ca0851bf7 debug pipeline 2024-10-21 07:33:32 +00:00
be67be95e9 solve merge 2024-10-19 19:08:39 +08:00
c9d05f0c86 merge 2024-10-19 19:07:40 +08:00
hofee
ed569254dc Merge branch 'master' of https://git.hofee.top/hofee/nbv_reconstruction 2024-10-19 19:06:26 +08:00
hofee
be7ec1a433 update 2024-10-19 19:06:09 +08:00
d0fbb0f198 remove o3d voxel_downsample 2024-10-17 14:28:19 +00:00
5dae3c53db remove mesh from strategy generator 2024-10-17 11:23:08 +00:00
15d1903080 Merge branch 'master' of https://git.hofee.top/hofee/nbv_reconstruction 2024-10-17 11:15:04 +00:00
hofee
b3344626cf solve merge 2024-10-17 06:14:46 -05:00
hofee
0267aed6e5 add normal and visualize util 2024-10-17 06:13:18 -05:00
22e7a1aed4 Merge branch 'master' of https://git.hofee.top/hofee/nbv_reconstruction 2024-10-17 11:11:14 +00:00
8892b6ed05 sync 2024-10-17 11:07:29 +00:00
hofee
8d92676c34 Merge branch 'master' of https://git.hofee.top/hofee/nbv_reconstruction 2024-10-10 10:16:03 -05:00
hofee
1e4fd13a24 update yaml 2024-10-10 10:15:55 -05:00
32 changed files with 1767 additions and 710 deletions

3
.gitignore vendored
View File

@@ -11,4 +11,5 @@ test/
*.log *.log
/data_generation/data/* /data_generation/data/*
/data_generation/output/* /data_generation/output/*
test/ test/
temp*

192
Readme.md Normal file
View File

@@ -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
...

22
TODO.md
View File

@@ -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

View File

@@ -1,5 +1,6 @@
from PytorchBoot.application import PytorchBootApplication from PytorchBoot.application import PytorchBootApplication
from runners.inferencer import Inferencer from runners.inferencer import Inferencer
from runners.inference_server import InferencerServer
@PytorchBootApplication("inference") @PytorchBootApplication("inference")
class InferenceApp: class InferenceApp:
@@ -14,3 +15,17 @@ class InferenceApp:
Evaluator("path_to_your_eval_config").run() Evaluator("path_to_your_eval_config").run()
''' '''
Inferencer("./configs/local/inference_config.yaml").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()

View File

@@ -5,5 +5,5 @@ from runners.data_spliter import DataSpliter
class DataSplitApp: class DataSplitApp:
@staticmethod @staticmethod
def start(): def start():
DataSpliter("configs/server/split_dataset_config.yaml").run() DataSpliter("configs/server/server_split_dataset_config.yaml").run()

View File

@@ -1,76 +1,72 @@
runner: runner:
general: general:
seed: 1 seed: 0
device: cuda device: cuda
cuda_visible_devices: "0,1,2,3,4,5,6,7" cuda_visible_devices: "0,1,2,3,4,5,6,7"
experiment: experiment:
name: w_gf_wo_lf_full name: train_ab_global_only
root_dir: "experiments" root_dir: "experiments"
epoch: 1 # -1 stands for last epoch epoch: -1 # -1 stands for last epoch
test: test:
dataset_list: dataset_list:
- OmniObject3d_train - OmniObject3d_test
blender_script_path: "/media/hofee/data/project/python/nbv_reconstruction/blender/data_renderer.py" blender_script_path: "/media/hofee/data/project/python/nbv_reconstruction/blender/data_renderer.py"
output_dir: "/media/hofee/data/project/python/nbv_reconstruction/nbv_reconstruction/test/inference_global_full_on_testset" output_dir: "/media/hofee/data/data/new_inference_test_output"
pipeline: nbv_reconstruction_global_pts_pipeline pipeline: nbv_reconstruction_pipeline
voxel_size: 0.003
dataset: dataset:
OmniObject3d_train: # OmniObject3d_train:
root_dir: "/media/hofee/repository/nbv_reconstruction_data_512" # root_dir: "C:\\Document\\Datasets\\inference_test1"
# model_dir: "C:\\Document\\Datasets\\scaled_object_meshes"
# source: seq_reconstruction_dataset_preprocessed
# split_file: "C:\\Document\\Datasets\\data_list\\sample.txt"
# type: test
# filter_degree: 75
# ratio: 1
# batch_size: 1
# num_workers: 12
# pts_num: 8192
# load_from_preprocess: True
OmniObject3d_test:
root_dir: "/media/hofee/data/data/new_testset_output"
model_dir: "/media/hofee/data/data/scaled_object_meshes" model_dir: "/media/hofee/data/data/scaled_object_meshes"
source: seq_nbv_reconstruction_dataset source: seq_reconstruction_dataset_preprocessed
split_file: "/media/hofee/data/project/python/nbv_reconstruction/nbv_reconstruction/test/test_set_list.txt" # split_file: "C:\\Document\\Datasets\\data_list\\OmniObject3d_test.txt"
type: test type: test
filter_degree: 75 filter_degree: 75
ratio: 1 eval_list:
- pose_diff
- coverage_rate_increase
ratio: 0.1
batch_size: 1 batch_size: 1
num_workers: 12 num_workers: 12
pts_num: 4096 pts_num: 8192
load_from_preprocess: False load_from_preprocess: True
pipeline: pipeline:
nbv_reconstruction_local_pts_pipeline: nbv_reconstruction_pipeline:
modules: modules:
pts_encoder: pointnet_encoder pts_encoder: pointnet_encoder
seq_encoder: transformer_seq_encoder seq_encoder: transformer_seq_encoder
pose_encoder: pose_encoder pose_encoder: pose_encoder
view_finder: gf_view_finder view_finder: gf_view_finder
eps: 1e-5 eps: 1e-5
global_scanned_feat: False
nbv_reconstruction_global_pts_pipeline:
modules:
pts_encoder: pointnet_encoder
pose_seq_encoder: transformer_pose_seq_encoder
pose_encoder: pose_encoder
view_finder: gf_view_finder
eps: 1e-5
global_scanned_feat: True global_scanned_feat: True
module: module:
pointnet_encoder: pointnet_encoder:
in_dim: 3 in_dim: 3
out_dim: 1024 out_dim: 1024
global_feat: True global_feat: True
feature_transform: False feature_transform: False
transformer_seq_encoder: transformer_seq_encoder:
pts_embed_dim: 1024 embed_dim: 256
pose_embed_dim: 256
num_heads: 4
ffn_dim: 256
num_layers: 3
output_dim: 2048
transformer_pose_seq_encoder:
pose_embed_dim: 256
num_heads: 4 num_heads: 4
ffn_dim: 256 ffn_dim: 256
num_layers: 3 num_layers: 3
@@ -86,7 +82,8 @@ module:
sample_mode: ode sample_mode: ode
sampling_steps: 500 sampling_steps: 500
sde_mode: ve sde_mode: ve
pose_encoder: pose_encoder:
pose_dim: 9 pose_dim: 9
out_dim: 256 out_dim: 256
pts_num_encoder:
out_dim: 64

View File

@@ -12,26 +12,16 @@ runner:
generate: generate:
voxel_threshold: 0.003 voxel_threshold: 0.003
soft_overlap_threshold: 0.3 overlap_area_threshold: 30
hard_overlap_threshold: 0.6 compute_with_normal: False
filter_degree: 75 scan_points_threshold: 10
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 overwrite: False
seq_num: 15 seq_num: 10
dataset_list: dataset_list:
- OmniObject3d - OmniObject3d
datasets: datasets:
OmniObject3d: OmniObject3d:
#"/media/hofee/data/data/temp_output" root_dir: /data/hofee/nbv_rec_part2_preprocessed
root_dir: /media/hofee/repository/full_data_output from: 155
model_dir: /media/hofee/data/data/scaled_object_meshes to: 165 # ..-1 means end
from: 0
to: -1 # -1 means end
#output_dir: "/media/hofee/data/data/label_output"

View File

@@ -84,7 +84,7 @@ module:
gf_view_finder: gf_view_finder:
t_feat_dim: 128 t_feat_dim: 128
pose_feat_dim: 256 pose_feat_dim: 256
main_feat_dim: 2048 main_feat_dim: 3072
regression_head: Rx_Ry_and_T regression_head: Rx_Ry_and_T
pose_mode: rot_matrix pose_mode: rot_matrix
per_point_feature: False per_point_feature: False

View File

@@ -7,17 +7,19 @@ runner:
name: debug name: debug
root_dir: experiments root_dir: experiments
generate: generate:
port: 5004 port: 5000
from: 0 from: 0
to: 1 # -1 means all to: -1 # -1 means all
object_dir: /media/hofee/data/data/box_object_meshes object_dir: /media/hofee/data/data/scaled_object_meshes
table_model_path: /media/hofee/data/data/others/table.obj table_model_path: "/media/hofee/data/data/others/table.obj"
output_dir: /media/hofee/data/data/box_views output_dir: /media/hofee/data/data/new_testset
object_list_path: /media/hofee/data/data/OmniObject3d_test.txt
use_list: True
binocular_vision: true binocular_vision: true
plane_size: 10 plane_size: 10
max_views: 512 max_views: 512
min_views: 128 min_views: 128
random_view_ratio: 0.2 random_view_ratio: 0.01
min_cam_table_included_degree: 20 min_cam_table_included_degree: 20
max_diag: 0.7 max_diag: 0.7
min_diag: 0.01 min_diag: 0.01

View File

@@ -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

View File

@@ -10,13 +10,13 @@ runner:
root_dir: "experiments" root_dir: "experiments"
split: # split: #
root_dir: "/home/data/hofee/project/nbv_rec/data/nbv_rec_data_512_preproc_npy" root_dir: "/data/hofee/data/packed_preprocessed_data"
type: "unseen_instance" # "unseen_category" type: "unseen_instance" # "unseen_category"
datasets: datasets:
OmniObject3d_train: OmniObject3d_train:
path: "../data/sample_for_training_preprocessed/OmniObject3d_train.txt" path: "/data/hofee/data/OmniObject3d_train.txt"
ratio: 0.9 ratio: 0.9
OmniObject3d_test: OmniObject3d_test:
path: "../data/sample_for_training_preprocessed/OmniObject3d_test.txt" path: "/data/hofee/data/OmniObject3d_test.txt"
ratio: 0.1 ratio: 0.1

View File

@@ -1,32 +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.005
overlap_threshold: 0.5
filter_degree: 75
to_specified_dir: False # if True, output_dir is used, otherwise, root_dir is used
save_points: True
save_best_combined_points: True
save_mesh: True
overwrite: False
dataset_list:
- OmniObject3d
datasets:
OmniObject3d:
#"/media/hofee/data/data/temp_output"
root_dir: "../data/sample_for_training/scenes"
model_dir: "../data/scaled_object_meshes"
#output_dir: "/media/hofee/data/data/label_output"

View File

@@ -3,13 +3,13 @@ runner:
general: general:
seed: 0 seed: 0
device: cuda device: cuda
cuda_visible_devices: "1" cuda_visible_devices: "0"
parallel: False parallel: False
experiment: experiment:
name: full_w_global_feat_wo_local_pts_feat name: train_ab_global_only
root_dir: "experiments" root_dir: "experiments"
use_checkpoint: False use_checkpoint: True
epoch: -1 # -1 stands for last epoch epoch: -1 # -1 stands for last epoch
max_epochs: 5000 max_epochs: 5000
save_checkpoint_interval: 1 save_checkpoint_interval: 1
@@ -28,57 +28,57 @@ runner:
- OmniObject3d_test - OmniObject3d_test
- OmniObject3d_val - OmniObject3d_val
pipeline: nbv_reconstruction_global_pts_pipeline pipeline: nbv_reconstruction_pipeline
dataset: dataset:
OmniObject3d_train: OmniObject3d_train:
root_dir: "/home/data/hofee/project/nbv_rec/data/nbv_rec_data_512_preproc_npy" root_dir: "/data/hofee/data/new_full_data"
model_dir: "../data/scaled_object_meshes" model_dir: "../data/scaled_object_meshes"
source: nbv_reconstruction_dataset source: nbv_reconstruction_dataset
split_file: "/home/data/hofee/project/nbv_rec/data/OmniObject3d_train.txt" split_file: "/data/hofee/data/new_full_data_list/OmniObject3d_train.txt"
type: train type: train
cache: True cache: True
ratio: 1 ratio: 1
batch_size: 160 batch_size: 80
num_workers: 16 num_workers: 128
pts_num: 4096 pts_num: 8192
load_from_preprocess: True load_from_preprocess: True
OmniObject3d_test: OmniObject3d_test:
root_dir: "/home/data/hofee/project/nbv_rec/data/nbv_rec_data_512_preproc_npy" root_dir: "/data/hofee/data/new_full_data"
model_dir: "../data/scaled_object_meshes" model_dir: "../data/scaled_object_meshes"
source: nbv_reconstruction_dataset source: nbv_reconstruction_dataset
split_file: "/home/data/hofee/project/nbv_rec/data/OmniObject3d_test.txt" split_file: "/data/hofee/data/new_full_data_list/OmniObject3d_test.txt"
type: test type: test
cache: True cache: True
filter_degree: 75 filter_degree: 75
eval_list: eval_list:
- pose_diff - pose_diff
ratio: 0.05 ratio: 1
batch_size: 160 batch_size: 80
num_workers: 12 num_workers: 12
pts_num: 4096 pts_num: 8192
load_from_preprocess: True load_from_preprocess: True
OmniObject3d_val: OmniObject3d_val:
root_dir: "/home/data/hofee/project/nbv_rec/data/nbv_rec_data_512_preproc_npy" root_dir: "/data/hofee/data/new_full_data"
model_dir: "../data/scaled_object_meshes" model_dir: "../data/scaled_object_meshes"
source: nbv_reconstruction_dataset source: nbv_reconstruction_dataset
split_file: "/home/data/hofee/project/nbv_rec/data/OmniObject3d_train.txt" split_file: "/data/hofee/data/new_full_data_list/OmniObject3d_train.txt"
type: test type: test
cache: True cache: True
filter_degree: 75 filter_degree: 75
eval_list: eval_list:
- pose_diff - pose_diff
ratio: 0.005 ratio: 0.1
batch_size: 160 batch_size: 80
num_workers: 12 num_workers: 12
pts_num: 4096 pts_num: 8192
load_from_preprocess: True load_from_preprocess: True
pipeline: pipeline:
nbv_reconstruction_local_pts_pipeline: nbv_reconstruction_pipeline:
modules: modules:
pts_encoder: pointnet_encoder pts_encoder: pointnet_encoder
seq_encoder: transformer_seq_encoder seq_encoder: transformer_seq_encoder
@@ -87,16 +87,6 @@ pipeline:
eps: 1e-5 eps: 1e-5
global_scanned_feat: True global_scanned_feat: True
nbv_reconstruction_global_pts_pipeline:
modules:
pts_encoder: pointnet_encoder
pose_seq_encoder: transformer_seq_encoder
pose_encoder: pose_encoder
view_finder: gf_view_finder
eps: 1e-5
global_scanned_feat: True
module: module:
@@ -107,11 +97,11 @@ module:
feature_transform: False feature_transform: False
transformer_seq_encoder: transformer_seq_encoder:
embed_dim: 1344 embed_dim: 256
num_heads: 4 num_heads: 4
ffn_dim: 256 ffn_dim: 256
num_layers: 3 num_layers: 3
output_dim: 2048 output_dim: 1024
gf_view_finder: gf_view_finder:
t_feat_dim: 128 t_feat_dim: 128
@@ -128,6 +118,9 @@ module:
pose_dim: 9 pose_dim: 9
out_dim: 256 out_dim: 256
pts_num_encoder:
out_dim: 64
loss_function: loss_function:
gf_loss: gf_loss:

View File

@@ -1,53 +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:
object_dir: ../data/scaled_object_meshes
table_model_path: ../data/others/table.obj
output_dir: ../data/nbv_reconstruction_data_512
binocular_vision: true
plane_size: 10
max_views: 512
min_views: 64
max_diag: 0.7
min_diag: 0.1
random_config:
display_table:
min_height: 0.05
max_height: 0.15
min_radius: 0.3
max_radius: 0.5
min_R: 0.05
max_R: 0.3
min_G: 0.05
max_G: 0.3
min_B: 0.05
max_B: 0.3
display_object:
min_x: 0
max_x: 0.03
min_y: 0
max_y: 0.03
min_z: 0.01
max_z: 0.01
random_rotation_ratio: 0.3
random_objects:
num: 4
cluster: 0.9
light_and_camera_config:
Camera:
near_plane: 0.01
far_plane: 5
fov_vertical: 25
resolution: [1280,800]
eye_distance: 0.15
eye_angle: 25
Light:
location: [0,0,3.5]
orientation: [0,0,0]
power: 150

View File

@@ -8,7 +8,7 @@ import torch
import os import os
import sys import sys
sys.path.append(r"/home/data/hofee/project/nbv_rec/nbv_reconstruction") sys.path.append(r"/data/hofee/project/nbv_rec/nbv_reconstruction")
from utils.data_load import DataLoadUtil from utils.data_load import DataLoadUtil
from utils.pose import PoseUtil from utils.pose import PoseUtil
@@ -31,7 +31,7 @@ class NBVReconstructionDataset(BaseDataset):
self.load_from_preprocess = config.get("load_from_preprocess", False) self.load_from_preprocess = config.get("load_from_preprocess", False)
if self.type == namespace.Mode.TEST: if self.type == namespace.Mode.TEST:
self.model_dir = config["model_dir"] #self.model_dir = config["model_dir"]
self.filter_degree = config["filter_degree"] self.filter_degree = config["filter_degree"]
if self.type == namespace.Mode.TRAIN: if self.type == namespace.Mode.TRAIN:
scale_ratio = 1 scale_ratio = 1
@@ -66,7 +66,9 @@ class NBVReconstructionDataset(BaseDataset):
if max_coverage_rate > scene_max_coverage_rate: if max_coverage_rate > scene_max_coverage_rate:
scene_max_coverage_rate = max_coverage_rate scene_max_coverage_rate = max_coverage_rate
max_coverage_rate_list.append(max_coverage_rate) max_coverage_rate_list.append(max_coverage_rate)
mean_coverage_rate = np.mean(max_coverage_rate_list)
if max_coverage_rate_list:
mean_coverage_rate = np.mean(max_coverage_rate_list)
for seq_idx in range(seq_num): for seq_idx in range(seq_num):
label_path = DataLoadUtil.get_label_path( label_path = DataLoadUtil.get_label_path(
@@ -112,6 +114,10 @@ class NBVReconstructionDataset(BaseDataset):
except Exception as e: except Exception as e:
Log.error(f"Save cache failed: {e}") Log.error(f"Save cache failed: {e}")
def voxel_downsample_with_mask(self, pts, voxel_size):
pass
def __getitem__(self, index): def __getitem__(self, index):
data_item_info = self.datalist[index] data_item_info = self.datalist[index]
scanned_views = data_item_info["scanned_views"] scanned_views = data_item_info["scanned_views"]
@@ -122,7 +128,7 @@ class NBVReconstructionDataset(BaseDataset):
scanned_views_pts, scanned_views_pts,
scanned_coverages_rate, scanned_coverages_rate,
scanned_n_to_world_pose, scanned_n_to_world_pose,
) = ([], [], [], []) ) = ([], [], [])
for view in scanned_views: for view in scanned_views:
frame_idx = view[0] frame_idx = view[0]
coverage_rate = view[1] coverage_rate = view[1]
@@ -160,27 +166,12 @@ class NBVReconstructionDataset(BaseDataset):
) )
combined_scanned_views_pts = np.concatenate(scanned_views_pts, axis=0) combined_scanned_views_pts = np.concatenate(scanned_views_pts, axis=0)
fps_downsampled_combined_scanned_pts, fps_idx = PtsUtil.fps_downsample_point_cloud( voxel_downsampled_combined_scanned_pts_np = PtsUtil.voxel_downsample_point_cloud(combined_scanned_views_pts, 0.002)
combined_scanned_views_pts, self.pts_num, require_idx=True random_downsampled_combined_scanned_pts_np = PtsUtil.random_downsample_point_cloud(voxel_downsampled_combined_scanned_pts_np, self.pts_num)
)
combined_scanned_views_pts_mask = np.zeros(len(scanned_views_pts), dtype=np.uint8)
start_idx = 0
for i in range(len(scanned_views_pts)):
end_idx = start_idx + len(scanned_views_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]
data_item = { data_item = {
"scanned_pts": np.asarray(scanned_views_pts, dtype=np.float32), # Ndarray(S x Nv x 3) "scanned_pts": np.asarray(scanned_views_pts, dtype=np.float32), # Ndarray(S x Nv x 3)
"scanned_pts_mask": np.asarray(fps_downsampled_combined_scanned_pts_mask,dtype=np.uint8), # Ndarray(N), range(0, S) "combined_scanned_pts": np.asarray(random_downsampled_combined_scanned_pts_np, dtype=np.float32), # Ndarray(N x 3)
"combined_scanned_pts": np.asarray(fps_downsampled_combined_scanned_pts, dtype=np.float32), # Ndarray(N x 3)
"scanned_coverage_rate": scanned_coverages_rate, # List(S): Float, range(0, 1) "scanned_coverage_rate": scanned_coverages_rate, # List(S): Float, range(0, 1)
"scanned_n_to_world_pose_9d": np.asarray(scanned_n_to_world_pose, dtype=np.float32), # Ndarray(S x 9) "scanned_n_to_world_pose_9d": np.asarray(scanned_n_to_world_pose, dtype=np.float32), # Ndarray(S x 9)
"best_coverage_rate": nbv_coverage_rate, # Float, range(0, 1) "best_coverage_rate": nbv_coverage_rate, # Float, range(0, 1)
@@ -215,14 +206,9 @@ class NBVReconstructionDataset(BaseDataset):
collate_data["combined_scanned_pts"] = torch.stack( collate_data["combined_scanned_pts"] = torch.stack(
[torch.tensor(item["combined_scanned_pts"]) for item in batch] [torch.tensor(item["combined_scanned_pts"]) for item in batch]
) )
collate_data["scanned_pts_mask"] = torch.stack(
[torch.tensor(item["scanned_pts_mask"]) for item in batch]
)
for key in batch[0].keys(): for key in batch[0].keys():
if key not in [ if key not in [
"scanned_pts", "scanned_pts",
"scanned_pts_mask",
"scanned_n_to_world_pose_9d", "scanned_n_to_world_pose_9d",
"best_to_world_pose_9d", "best_to_world_pose_9d",
"combined_scanned_pts", "combined_scanned_pts",
@@ -241,10 +227,9 @@ if __name__ == "__main__":
torch.manual_seed(seed) torch.manual_seed(seed)
np.random.seed(seed) np.random.seed(seed)
config = { config = {
"root_dir": "/home/data/hofee/project/nbv_rec/data/nbv_rec_data_512_preproc_npy", "root_dir": "/data/hofee/data/packed_preprocessed_data",
"model_dir": "/home/data/hofee/project/nbv_rec/data/scaled_object_meshes",
"source": "nbv_reconstruction_dataset", "source": "nbv_reconstruction_dataset",
"split_file": "/home/data/hofee/project/nbv_rec/data/OmniObject3d_test.txt", "split_file": "/data/hofee/data/OmniObject3d_train.txt",
"load_from_preprocess": True, "load_from_preprocess": True,
"ratio": 0.5, "ratio": 0.5,
"batch_size": 2, "batch_size": 2,

154
core/old_seq_dataset.py Normal file
View File

@@ -0,0 +1,154 @@
import numpy as np
from PytorchBoot.dataset import BaseDataset
import PytorchBoot.namespace as namespace
import PytorchBoot.stereotype as stereotype
from PytorchBoot.utils.log_util import Log
import torch
import os
import sys
sys.path.append(r"/home/data/hofee/project/nbv_rec/nbv_reconstruction")
from utils.data_load import DataLoadUtil
from utils.pose import PoseUtil
from utils.pts import PtsUtil
@stereotype.dataset("old_seq_nbv_reconstruction_dataset")
class SeqNBVReconstructionDataset(BaseDataset):
def __init__(self, config):
super(SeqNBVReconstructionDataset, self).__init__(config)
self.type = config["type"]
if self.type != namespace.Mode.TEST:
Log.error("Dataset <seq_nbv_reconstruction_dataset> Only support test mode", terminate=True)
self.config = config
self.root_dir = config["root_dir"]
self.split_file_path = config["split_file"]
self.scene_name_list = self.load_scene_name_list()
self.datalist = self.get_datalist()
self.pts_num = config["pts_num"]
self.model_dir = config["model_dir"]
self.filter_degree = config["filter_degree"]
self.load_from_preprocess = config.get("load_from_preprocess", False)
def load_scene_name_list(self):
scene_name_list = []
with open(self.split_file_path, "r") as f:
for line in f:
scene_name = line.strip()
scene_name_list.append(scene_name)
return scene_name_list
def get_datalist(self):
datalist = []
for scene_name in self.scene_name_list:
seq_num = DataLoadUtil.get_label_num(self.root_dir, scene_name)
scene_max_coverage_rate = 0
scene_max_cr_idx = 0
for seq_idx in range(seq_num):
label_path = DataLoadUtil.get_label_path(self.root_dir, scene_name, seq_idx)
label_data = DataLoadUtil.load_label(label_path)
max_coverage_rate = label_data["max_coverage_rate"]
if max_coverage_rate > scene_max_coverage_rate:
scene_max_coverage_rate = max_coverage_rate
scene_max_cr_idx = seq_idx
label_path = DataLoadUtil.get_label_path(self.root_dir, scene_name, scene_max_cr_idx)
label_data = DataLoadUtil.load_label(label_path)
first_frame = label_data["best_sequence"][0]
best_seq_len = len(label_data["best_sequence"])
datalist.append({
"scene_name": scene_name,
"first_frame": first_frame,
"max_coverage_rate": scene_max_coverage_rate,
"best_seq_len": best_seq_len,
"label_idx": scene_max_cr_idx,
})
return datalist
def __getitem__(self, index):
data_item_info = self.datalist[index]
first_frame_idx = data_item_info["first_frame"][0]
first_frame_coverage = data_item_info["first_frame"][1]
max_coverage_rate = data_item_info["max_coverage_rate"]
scene_name = data_item_info["scene_name"]
first_cam_info = DataLoadUtil.load_cam_info(DataLoadUtil.get_path(self.root_dir, scene_name, first_frame_idx), binocular=True)
first_view_path = DataLoadUtil.get_path(self.root_dir, scene_name, first_frame_idx)
first_left_cam_pose = first_cam_info["cam_to_world"]
first_center_cam_pose = first_cam_info["cam_to_world_O"]
first_target_point_cloud = DataLoadUtil.load_from_preprocessed_pts(first_view_path)
first_pts_num = first_target_point_cloud.shape[0]
first_downsampled_target_point_cloud = PtsUtil.random_downsample_point_cloud(first_target_point_cloud, self.pts_num)
first_to_world_rot_6d = PoseUtil.matrix_to_rotation_6d_numpy(np.asarray(first_left_cam_pose[:3,:3]))
first_to_world_trans = first_left_cam_pose[:3,3]
first_to_world_9d = np.concatenate([first_to_world_rot_6d, first_to_world_trans], axis=0)
diag = DataLoadUtil.get_bbox_diag(self.model_dir, scene_name)
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_num": np.asarray(
first_pts_num, dtype=np.int32
),
"first_pts": np.asarray([first_downsampled_target_point_cloud],dtype=np.float32),
"combined_scanned_pts": np.asarray(first_downsampled_target_point_cloud,dtype=np.float32),
"first_to_world_9d": np.asarray([first_to_world_9d],dtype=np.float32),
"scene_name": scene_name,
"max_coverage_rate": max_coverage_rate,
"voxel_threshold": voxel_threshold,
"filter_degree": self.filter_degree,
"O_to_L_pose": first_O_to_first_L_pose,
"first_frame_coverage": first_frame_coverage,
"scene_path": scene_path,
"model_points_normals": model_points_normals,
"best_seq_len": data_item_info["best_seq_len"],
"first_frame_id": first_frame_idx,
}
return data_item
def __len__(self):
return len(self.datalist)
def get_collate_fn(self):
def collate_fn(batch):
collate_data = {}
collate_data["first_pts"] = [torch.tensor(item['first_pts']) for item in batch]
collate_data["first_to_world_9d"] = [torch.tensor(item['first_to_world_9d']) for item in batch]
collate_data["combined_scanned_pts"] = torch.stack([torch.tensor(item['combined_scanned_pts']) for item in batch])
for key in batch[0].keys():
if key not in ["first_pts", "first_to_world_9d", "combined_scanned_pts"]:
collate_data[key] = [item[key] for item in batch]
return collate_data
return collate_fn
# -------------- Debug ---------------- #
if __name__ == "__main__":
import torch
seed = 0
torch.manual_seed(seed)
np.random.seed(seed)
config = {
"root_dir": "/home/data/hofee/project/nbv_rec/data/nbv_rec_data_512_preproc_npy",
"split_file": "/home/data/hofee/project/nbv_rec/data/OmniObject3d_train.txt",
"model_dir": "/home/data/hofee/project/nbv_rec/data/scaled_object_meshes",
"ratio": 0.005,
"batch_size": 2,
"filter_degree": 75,
"num_workers": 0,
"pts_num": 32684,
"type": namespace.Mode.TEST,
"load_from_preprocess": True
}
ds = SeqNBVReconstructionDataset(config)
print(len(ds))
#ds.__getitem__(10)
dl = ds.get_loader(shuffle=True)
for idx, data in enumerate(dl):
data = ds.process_batch(data, "cuda:0")
print(data)
# ------ Debug Start ------
import ipdb;ipdb.set_trace()
# ------ Debug End ------+

View File

@@ -1,4 +1,5 @@
import torch import torch
import time
from torch import nn from torch import nn
import PytorchBoot.namespace as namespace import PytorchBoot.namespace as namespace
import PytorchBoot.stereotype as stereotype import PytorchBoot.stereotype as stereotype
@@ -6,10 +7,10 @@ from PytorchBoot.factory.component_factory import ComponentFactory
from PytorchBoot.utils import Log from PytorchBoot.utils import Log
@stereotype.pipeline("nbv_reconstruction_global_pts_n_num_pipeline") @stereotype.pipeline("nbv_reconstruction_pipeline")
class NBVReconstructionGlobalPointsPipeline(nn.Module): class NBVReconstructionPipeline(nn.Module):
def __init__(self, config): def __init__(self, config):
super(NBVReconstructionGlobalPointsPipeline, self).__init__() super(NBVReconstructionPipeline, self).__init__()
self.config = config self.config = config
self.module_config = config["modules"] self.module_config = config["modules"]
@@ -19,12 +20,8 @@ class NBVReconstructionGlobalPointsPipeline(nn.Module):
self.pose_encoder = ComponentFactory.create( self.pose_encoder = ComponentFactory.create(
namespace.Stereotype.MODULE, self.module_config["pose_encoder"] namespace.Stereotype.MODULE, self.module_config["pose_encoder"]
) )
self.pts_num_encoder = ComponentFactory.create( self.seq_encoder = ComponentFactory.create(
namespace.Stereotype.MODULE, self.module_config["pts_num_encoder"] namespace.Stereotype.MODULE, self.module_config["seq_encoder"]
)
self.transformer_seq_encoder = ComponentFactory.create(
namespace.Stereotype.MODULE, self.module_config["transformer_seq_encoder"]
) )
self.view_finder = ComponentFactory.create( self.view_finder = ComponentFactory.create(
namespace.Stereotype.MODULE, self.module_config["view_finder"] namespace.Stereotype.MODULE, self.module_config["view_finder"]
@@ -32,7 +29,6 @@ class NBVReconstructionGlobalPointsPipeline(nn.Module):
self.eps = float(self.config["eps"]) self.eps = float(self.config["eps"])
self.enable_global_scanned_feat = self.config["global_scanned_feat"]
def forward(self, data): def forward(self, data):
mode = data["mode"] mode = data["mode"]
@@ -92,47 +88,23 @@ class NBVReconstructionGlobalPointsPipeline(nn.Module):
scanned_n_to_world_pose_9d_batch = data[ scanned_n_to_world_pose_9d_batch = data[
"scanned_n_to_world_pose_9d" "scanned_n_to_world_pose_9d"
] # List(B): Tensor(S x 9) ] # List(B): Tensor(S x 9)
scanned_pts_mask_batch = data[
"scanned_pts_mask"
] # Tensor(B x N)
device = next(self.parameters()).device device = next(self.parameters()).device
embedding_list_batch = [] embedding_list_batch = []
combined_scanned_pts_batch = data["combined_scanned_pts"] # Tensor(B x N x 3) combined_scanned_pts_batch = data["combined_scanned_pts"] # Tensor(B x N x 3)
global_scanned_feat, perpoint_scanned_feat_batch = self.pts_encoder.encode_points( global_scanned_feat = self.pts_encoder.encode_points(
combined_scanned_pts_batch, require_per_point_feat=True combined_scanned_pts_batch, require_per_point_feat=False
) # global_scanned_feat: Tensor(B x Dg), perpoint_scanned_feat: Tensor(B x N x Dl) ) # global_scanned_feat: Tensor(B x Dg)
for scanned_n_to_world_pose_9d, scanned_mask, perpoint_scanned_feat in zip( for scanned_n_to_world_pose_9d in scanned_n_to_world_pose_9d_batch:
scanned_n_to_world_pose_9d_batch,
scanned_pts_mask_batch,
perpoint_scanned_feat_batch,
):
scanned_target_pts_num = [] # List(S): Int
partial_feat_seq = []
seq_len = len(scanned_n_to_world_pose_9d)
for seq_idx in range(seq_len):
partial_idx_in_combined_pts = scanned_mask == seq_idx # Ndarray(V), N->V idx mask
partial_perpoint_feat = perpoint_scanned_feat[partial_idx_in_combined_pts] # Ndarray(V x Dl)
partial_feat = torch.mean(partial_perpoint_feat, dim=0)[0] # Tensor(Dl)
partial_feat_seq.append(partial_feat)
scanned_target_pts_num.append(partial_perpoint_feat.shape[0])
scanned_target_pts_num = torch.tensor(scanned_target_pts_num, dtype=torch.int32).to(device) # Tensor(S)
scanned_n_to_world_pose_9d = scanned_n_to_world_pose_9d.to(device) # Tensor(S x 9) scanned_n_to_world_pose_9d = scanned_n_to_world_pose_9d.to(device) # Tensor(S x 9)
pose_feat_seq = self.pose_encoder.encode_pose(scanned_n_to_world_pose_9d) # Tensor(S x Dp)
pose_feat_seq = self.pose_encoder.encode_pose(scanned_n_to_world_pose_9d) # Tensor(S x Dp) seq_embedding = pose_feat_seq
pts_num_feat_seq = self.pts_num_encoder.encode_pts_num(scanned_target_pts_num) # Tensor(S x Dn) embedding_list_batch.append(seq_embedding) # List(B): Tensor(S x (Dp))
partial_feat_seq = torch.stack(partial_feat_seq) # Tensor(S x Dl)
seq_feat = self.seq_encoder.encode_sequence(embedding_list_batch) # Tensor(B x Ds)
seq_embedding = torch.cat([pose_feat_seq, pts_num_feat_seq, partial_feat_seq], dim=-1) # Tensor(S x (Dp+Dn+Dl))
embedding_list_batch.append(seq_embedding) # List(B): Tensor(S x (Dp+Dn+Dl))
seq_feat = self.transformer_seq_encoder.encode_sequence(embedding_list_batch) # Tensor(B x Ds)
main_feat = torch.cat([seq_feat, global_scanned_feat], dim=-1) # Tensor(B x (Ds+Dg)) main_feat = torch.cat([seq_feat, global_scanned_feat], dim=-1) # Tensor(B x (Ds+Dg))
if torch.isnan(main_feat).any(): if torch.isnan(main_feat).any():

View File

@@ -1,154 +1,204 @@
import numpy as np import numpy as np
from PytorchBoot.dataset import BaseDataset from PytorchBoot.dataset import BaseDataset
import PytorchBoot.namespace as namespace import PytorchBoot.namespace as namespace
import PytorchBoot.stereotype as stereotype import PytorchBoot.stereotype as stereotype
from PytorchBoot.utils.log_util import Log from PytorchBoot.config import ConfigManager
import torch from PytorchBoot.utils.log_util import Log
import os import torch
import sys import os
sys.path.append(r"/home/data/hofee/project/nbv_rec/nbv_reconstruction") import sys
from utils.data_load import DataLoadUtil sys.path.append(r"/media/hofee/data/project/python/nbv_reconstruction/nbv_reconstruction")
from utils.pose import PoseUtil
from utils.pts import PtsUtil from utils.data_load import DataLoadUtil
from utils.pose import PoseUtil
@stereotype.dataset("seq_nbv_reconstruction_dataset") from utils.pts import PtsUtil
class SeqNBVReconstructionDataset(BaseDataset):
def __init__(self, config):
super(SeqNBVReconstructionDataset, self).__init__(config) @stereotype.dataset("seq_reconstruction_dataset")
self.type = config["type"] class SeqReconstructionDataset(BaseDataset):
if self.type != namespace.Mode.TEST: def __init__(self, config):
Log.error("Dataset <seq_nbv_reconstruction_dataset> Only support test mode", terminate=True) super(SeqReconstructionDataset, self).__init__(config)
self.config = config self.config = config
self.root_dir = config["root_dir"] self.root_dir = config["root_dir"]
self.split_file_path = config["split_file"] self.split_file_path = config["split_file"]
self.scene_name_list = self.load_scene_name_list() self.scene_name_list = self.load_scene_name_list()
self.datalist = self.get_datalist() self.datalist = self.get_datalist()
self.pts_num = config["pts_num"]
self.pts_num = config["pts_num"]
self.model_dir = config["model_dir"] self.type = config["type"]
self.filter_degree = config["filter_degree"] self.cache = config.get("cache")
self.load_from_preprocess = config.get("load_from_preprocess", False) self.load_from_preprocess = config.get("load_from_preprocess", False)
if self.type == namespace.Mode.TEST:
def load_scene_name_list(self): #self.model_dir = config["model_dir"]
scene_name_list = [] self.filter_degree = config["filter_degree"]
with open(self.split_file_path, "r") as f: if self.type == namespace.Mode.TRAIN:
for line in f: scale_ratio = 1
scene_name = line.strip() self.datalist = self.datalist*scale_ratio
scene_name_list.append(scene_name) if self.cache:
return scene_name_list expr_root = ConfigManager.get("runner", "experiment", "root_dir")
expr_name = ConfigManager.get("runner", "experiment", "name")
def get_datalist(self): self.cache_dir = os.path.join(expr_root, expr_name, "cache")
datalist = [] # self.preprocess_cache()
for scene_name in self.scene_name_list:
seq_num = DataLoadUtil.get_label_num(self.root_dir, scene_name) def load_scene_name_list(self):
scene_max_coverage_rate = 0 scene_name_list = []
scene_max_cr_idx = 0 with open(self.split_file_path, "r") as f:
for line in f:
for seq_idx in range(seq_num): scene_name = line.strip()
label_path = DataLoadUtil.get_label_path(self.root_dir, scene_name, seq_idx) if os.path.exists(os.path.join(self.root_dir, scene_name)):
label_data = DataLoadUtil.load_label(label_path) scene_name_list.append(scene_name)
max_coverage_rate = label_data["max_coverage_rate"] return scene_name_list
if max_coverage_rate > scene_max_coverage_rate:
scene_max_coverage_rate = max_coverage_rate def get_scene_name_list(self):
scene_max_cr_idx = seq_idx return self.scene_name_list
label_path = DataLoadUtil.get_label_path(self.root_dir, scene_name, scene_max_cr_idx) def get_datalist(self):
label_data = DataLoadUtil.load_label(label_path) datalist = []
first_frame = label_data["best_sequence"][0] total = len(self.scene_name_list)
best_seq_len = len(label_data["best_sequence"]) for idx, scene_name in enumerate(self.scene_name_list):
datalist.append({ print(f"processing {scene_name} ({idx}/{total})")
"scene_name": scene_name, scene_max_cr_idx = 0
"first_frame": first_frame, frame_len = DataLoadUtil.get_scene_seq_length(self.root_dir, scene_name)
"max_coverage_rate": scene_max_coverage_rate,
"best_seq_len": best_seq_len, for i in range(frame_len):
"label_idx": scene_max_cr_idx, path = DataLoadUtil.get_path(self.root_dir, scene_name, i)
}) pts = DataLoadUtil.load_from_preprocessed_pts(path, "npy")
return datalist if pts.shape[0] == 0:
continue
def __getitem__(self, index): datalist.append({
data_item_info = self.datalist[index] "scene_name": scene_name,
first_frame_idx = data_item_info["first_frame"][0] "first_frame": i,
first_frame_coverage = data_item_info["first_frame"][1] "best_seq_len": -1,
max_coverage_rate = data_item_info["max_coverage_rate"] "max_coverage_rate": 1.0,
scene_name = data_item_info["scene_name"] "label_idx": scene_max_cr_idx,
first_cam_info = DataLoadUtil.load_cam_info(DataLoadUtil.get_path(self.root_dir, scene_name, first_frame_idx), binocular=True) })
first_view_path = DataLoadUtil.get_path(self.root_dir, scene_name, first_frame_idx) return datalist
first_left_cam_pose = first_cam_info["cam_to_world"]
first_center_cam_pose = first_cam_info["cam_to_world_O"] def preprocess_cache(self):
first_target_point_cloud = DataLoadUtil.load_from_preprocessed_pts(first_view_path) Log.info("preprocessing cache...")
first_pts_num = first_target_point_cloud.shape[0] for item_idx in range(len(self.datalist)):
first_downsampled_target_point_cloud = PtsUtil.random_downsample_point_cloud(first_target_point_cloud, self.pts_num) self.__getitem__(item_idx)
first_to_world_rot_6d = PoseUtil.matrix_to_rotation_6d_numpy(np.asarray(first_left_cam_pose[:3,:3])) Log.success("finish preprocessing cache.")
first_to_world_trans = first_left_cam_pose[:3,3]
first_to_world_9d = np.concatenate([first_to_world_rot_6d, first_to_world_trans], axis=0) def load_from_cache(self, scene_name, curr_frame_idx):
diag = DataLoadUtil.get_bbox_diag(self.model_dir, scene_name) cache_name = f"{scene_name}_{curr_frame_idx}.txt"
voxel_threshold = diag*0.02 cache_path = os.path.join(self.cache_dir, cache_name)
first_O_to_first_L_pose = np.dot(np.linalg.inv(first_left_cam_pose), first_center_cam_pose) if os.path.exists(cache_path):
scene_path = os.path.join(self.root_dir, scene_name) data = np.loadtxt(cache_path)
model_points_normals = DataLoadUtil.load_points_normals(self.root_dir, scene_name) return data
else:
data_item = { return None
"first_pts_num": np.asarray(
first_pts_num, dtype=np.int32 def save_to_cache(self, scene_name, curr_frame_idx, data):
), cache_name = f"{scene_name}_{curr_frame_idx}.txt"
"first_pts": np.asarray([first_downsampled_target_point_cloud],dtype=np.float32), cache_path = os.path.join(self.cache_dir, cache_name)
"combined_scanned_pts": np.asarray(first_downsampled_target_point_cloud,dtype=np.float32), try:
"first_to_world_9d": np.asarray([first_to_world_9d],dtype=np.float32), np.savetxt(cache_path, data)
"scene_name": scene_name, except Exception as e:
"max_coverage_rate": max_coverage_rate, Log.error(f"Save cache failed: {e}")
"voxel_threshold": voxel_threshold,
"filter_degree": self.filter_degree, def seq_combined_pts(self, scene, frame_idx_list):
"O_to_L_pose": first_O_to_first_L_pose, all_combined_pts = []
"first_frame_coverage": first_frame_coverage, for i in frame_idx_list:
"scene_path": scene_path, path = DataLoadUtil.get_path(self.root_dir, scene, i)
"model_points_normals": model_points_normals, pts = DataLoadUtil.load_from_preprocessed_pts(path,"npy")
"best_seq_len": data_item_info["best_seq_len"], if pts.shape[0] == 0:
"first_frame_id": first_frame_idx, continue
} all_combined_pts.append(pts)
return data_item all_combined_pts = np.vstack(all_combined_pts)
downsampled_all_pts = PtsUtil.voxel_downsample_point_cloud(all_combined_pts, 0.003)
def __len__(self): return downsampled_all_pts
return len(self.datalist)
def __getitem__(self, index):
def get_collate_fn(self): data_item_info = self.datalist[index]
def collate_fn(batch): max_coverage_rate = data_item_info["max_coverage_rate"]
collate_data = {} best_seq_len = data_item_info["best_seq_len"]
collate_data["first_pts"] = [torch.tensor(item['first_pts']) for item in batch] scene_name = data_item_info["scene_name"]
collate_data["first_to_world_9d"] = [torch.tensor(item['first_to_world_9d']) for item in batch] (
collate_data["combined_scanned_pts"] = torch.stack([torch.tensor(item['combined_scanned_pts']) for item in batch]) scanned_views_pts,
for key in batch[0].keys(): scanned_coverages_rate,
if key not in ["first_pts", "first_to_world_9d", "combined_scanned_pts"]: scanned_n_to_world_pose,
collate_data[key] = [item[key] for item in batch] ) = ([], [], [])
return collate_data view = data_item_info["first_frame"]
return collate_fn frame_idx = view
view_path = DataLoadUtil.get_path(self.root_dir, scene_name, frame_idx)
# -------------- Debug ---------------- # cam_info = DataLoadUtil.load_cam_info(view_path, binocular=True)
if __name__ == "__main__":
import torch n_to_world_pose = cam_info["cam_to_world"]
seed = 0 target_point_cloud = (
torch.manual_seed(seed) DataLoadUtil.load_from_preprocessed_pts(view_path)
np.random.seed(seed) )
config = { downsampled_target_point_cloud = PtsUtil.random_downsample_point_cloud(
"root_dir": "/home/data/hofee/project/nbv_rec/data/nbv_rec_data_512_preproc_npy", target_point_cloud, self.pts_num
"split_file": "/home/data/hofee/project/nbv_rec/data/OmniObject3d_train.txt", )
"model_dir": "/home/data/hofee/project/nbv_rec/data/scaled_object_meshes", scanned_views_pts.append(downsampled_target_point_cloud)
"ratio": 0.005,
"batch_size": 2, n_to_world_6d = PoseUtil.matrix_to_rotation_6d_numpy(
"filter_degree": 75, np.asarray(n_to_world_pose[:3, :3])
"num_workers": 0, )
"pts_num": 32684, first_left_cam_pose = cam_info["cam_to_world"]
"type": namespace.Mode.TEST, first_center_cam_pose = cam_info["cam_to_world_O"]
"load_from_preprocess": True first_O_to_first_L_pose = np.dot(np.linalg.inv(first_left_cam_pose), first_center_cam_pose)
} n_to_world_trans = n_to_world_pose[:3, 3]
ds = SeqNBVReconstructionDataset(config) n_to_world_9d = np.concatenate([n_to_world_6d, n_to_world_trans], axis=0)
print(len(ds)) scanned_n_to_world_pose.append(n_to_world_9d)
#ds.__getitem__(10)
dl = ds.get_loader(shuffle=True) frame_list = []
for idx, data in enumerate(dl): for i in range(DataLoadUtil.get_scene_seq_length(self.root_dir, scene_name)):
data = ds.process_batch(data, "cuda:0") frame_list.append(i)
print(data) gt_pts = self.seq_combined_pts(scene_name, frame_list)
# ------ Debug Start ------ data_item = {
import ipdb;ipdb.set_trace() "first_scanned_pts": np.asarray(scanned_views_pts, dtype=np.float32), # Ndarray(S x Nv x 3)
# ------ Debug End ------+ "first_scanned_n_to_world_pose_9d": np.asarray(scanned_n_to_world_pose, dtype=np.float32), # Ndarray(S x 9)
"seq_max_coverage_rate": max_coverage_rate, # Float, range(0, 1)
"best_seq_len": best_seq_len, # Int
"scene_name": scene_name, # String
"gt_pts": gt_pts, # Ndarray(N x 3)
"scene_path": os.path.join(self.root_dir, scene_name), # String
"O_to_L_pose": first_O_to_first_L_pose,
}
return data_item
def __len__(self):
return len(self.datalist)
# -------------- Debug ---------------- #
if __name__ == "__main__":
#import ipdb; ipdb.set_trace()
import torch
from tqdm import tqdm
import pickle
import os
seed = 0
torch.manual_seed(seed)
np.random.seed(seed)
config = {
"root_dir": "/media/hofee/data/data/new_testset",
"source": "seq_reconstruction_dataset",
"split_file": "/media/hofee/data/data/OmniObject3d_test.txt",
"load_from_preprocess": True,
"filter_degree": 75,
"num_workers": 0,
"pts_num": 8192,
"type": namespace.Mode.TEST,
}
output_dir = "/media/hofee/data/data/new_testset_output"
os.makedirs(output_dir, exist_ok=True)
ds = SeqReconstructionDataset(config)
for i in tqdm(range(len(ds)), desc="processing dataset"):
output_path = os.path.join(output_dir, f"item_{i}.pkl")
item = ds.__getitem__(i)
for key, value in item.items():
if isinstance(value, np.ndarray):
item[key] = value.tolist()
import ipdb; ipdb.set_trace()
with open(output_path, "wb") as f:
pickle.dump(item, f)

View File

@@ -0,0 +1,84 @@
import numpy as np
from PytorchBoot.dataset import BaseDataset
import PytorchBoot.namespace as namespace
import PytorchBoot.stereotype as stereotype
from PytorchBoot.config import ConfigManager
from PytorchBoot.utils.log_util import Log
import pickle
import torch
import os
import sys
sys.path.append(r"C:\Document\Local Project\nbv_rec\nbv_reconstruction")
from utils.data_load import DataLoadUtil
from utils.pose import PoseUtil
from utils.pts import PtsUtil
@stereotype.dataset("seq_reconstruction_dataset_preprocessed")
class SeqReconstructionDatasetPreprocessed(BaseDataset):
def __init__(self, config):
super(SeqReconstructionDatasetPreprocessed, self).__init__(config)
self.config = config
self.root_dir = config["root_dir"]
self.real_root_dir = r"/media/hofee/data/data/new_testset"
self.item_list = os.listdir(self.root_dir)
def __getitem__(self, index):
data = pickle.load(open(os.path.join(self.root_dir, self.item_list[index]), "rb"))
data_item = {
"first_scanned_pts": np.asarray(data["first_scanned_pts"], dtype=np.float32), # Ndarray(S x Nv x 3)
"first_scanned_n_to_world_pose_9d": np.asarray(data["first_scanned_n_to_world_pose_9d"], dtype=np.float32), # Ndarray(S x 9)
"seq_max_coverage_rate": data["seq_max_coverage_rate"], # Float, range(0, 1)
"best_seq_len": data["best_seq_len"], # Int
"scene_name": data["scene_name"], # String
"gt_pts": np.asarray(data["gt_pts"], dtype=np.float32), # Ndarray(N x 3)
"scene_path": os.path.join(self.real_root_dir, data["scene_name"]), # String
"O_to_L_pose": np.asarray(data["O_to_L_pose"], dtype=np.float32),
}
return data_item
def __len__(self):
return len(self.item_list)
# -------------- Debug ---------------- #
if __name__ == "__main__":
import torch
seed = 0
torch.manual_seed(seed)
np.random.seed(seed)
'''
OmniObject3d_test:
root_dir: "H:\\AI\\Datasets\\packed_test_data"
model_dir: "H:\\AI\\Datasets\\scaled_object_meshes"
source: seq_reconstruction_dataset
split_file: "H:\\AI\\Datasets\\data_list\\OmniObject3d_test.txt"
type: test
filter_degree: 75
eval_list:
- pose_diff
- coverage_rate_increase
ratio: 0.1
batch_size: 1
num_workers: 12
pts_num: 8192
load_from_preprocess: True
'''
config = {
"root_dir": "H:\\AI\\Datasets\\packed_test_data",
"source": "seq_reconstruction_dataset",
"split_file": "H:\\AI\\Datasets\\data_list\\OmniObject3d_test.txt",
"load_from_preprocess": True,
"ratio": 1,
"filter_degree": 75,
"num_workers": 0,
"pts_num": 8192,
"type": "test",
}
ds = SeqReconstructionDataset(config)
print(len(ds))
print(ds.__getitem__(10))

View File

@@ -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("清理完成")

View File

@@ -0,0 +1,48 @@
import os
import shutil
def pack_scene_data(root, scene, output_dir):
scene_dir = os.path.join(output_dir, scene)
if not os.path.exists(scene_dir):
os.makedirs(scene_dir)
pts_dir = os.path.join(root, scene, "pts")
if os.path.exists(pts_dir):
shutil.move(pts_dir, os.path.join(scene_dir, "pts"))
scan_points_indices_dir = os.path.join(root, scene, "scan_points_indices")
if os.path.exists(scan_points_indices_dir):
shutil.move(scan_points_indices_dir, os.path.join(scene_dir, "scan_points_indices"))
scan_points_file = os.path.join(root, scene, "scan_points.txt")
if os.path.exists(scan_points_file):
shutil.move(scan_points_file, os.path.join(scene_dir, "scan_points.txt"))
model_pts_nrm_file = os.path.join(root, scene, "points_and_normals.txt")
if os.path.exists(model_pts_nrm_file):
shutil.move(model_pts_nrm_file, os.path.join(scene_dir, "points_and_normals.txt"))
camera_dir = os.path.join(root, scene, "camera_params")
if os.path.exists(camera_dir):
shutil.move(camera_dir, os.path.join(scene_dir, "camera_params"))
scene_info_file = os.path.join(root, scene, "scene_info.json")
if os.path.exists(scene_info_file):
shutil.move(scene_info_file, os.path.join(scene_dir, "scene_info.json"))
def pack_all_scenes(root, scene_list, output_dir):
for idx, scene in enumerate(scene_list):
print(f"正在打包场景 {scene} ({idx+1}/{len(scene_list)})")
pack_scene_data(root, scene, output_dir)
if __name__ == "__main__":
root = r"H:\AI\Datasets\nbv_rec_part2"
output_dir = r"H:\AI\Datasets\scene_info_part2"
scene_list = os.listdir(root)
from_idx = 0
to_idx = len(scene_list)
print(f"正在打包场景 {scene_list[from_idx:to_idx]}")
pack_all_scenes(root, scene_list[from_idx:to_idx], output_dir)
print("打包完成")

View File

@@ -0,0 +1,41 @@
import os
import shutil
def pack_scene_data(root, scene, output_dir):
scene_dir = os.path.join(output_dir, scene)
if not os.path.exists(scene_dir):
os.makedirs(scene_dir)
pts_dir = os.path.join(root, scene, "pts")
if os.path.exists(pts_dir):
shutil.move(pts_dir, os.path.join(scene_dir, "pts"))
camera_dir = os.path.join(root, scene, "camera_params")
if os.path.exists(camera_dir):
shutil.move(camera_dir, os.path.join(scene_dir, "camera_params"))
scene_info_file = os.path.join(root, scene, "scene_info.json")
if os.path.exists(scene_info_file):
shutil.move(scene_info_file, os.path.join(scene_dir, "scene_info.json"))
label_dir = os.path.join(root, scene, "label")
if os.path.exists(label_dir):
shutil.move(label_dir, os.path.join(scene_dir, "label"))
def pack_all_scenes(root, scene_list, output_dir):
for idx, scene in enumerate(scene_list):
print(f"packing {scene} ({idx+1}/{len(scene_list)})")
pack_scene_data(root, scene, output_dir)
if __name__ == "__main__":
root = r"/media/hofee/repository/data_part_1"
output_dir = r"/media/hofee/repository/upload_part1"
scene_list = os.listdir(root)
from_idx = 0
to_idx = len(scene_list)
print(f"packing {scene_list[from_idx:to_idx]}")
pack_all_scenes(root, scene_list[from_idx:to_idx], output_dir)
print("packing done")

View File

@@ -9,8 +9,6 @@ from utils.reconstruction import ReconstructionUtil
from utils.data_load import DataLoadUtil from utils.data_load import DataLoadUtil
from utils.pts import PtsUtil from utils.pts import PtsUtil
# scan shoe 536
def save_np_pts(path, pts: np.ndarray, file_type="txt"): def save_np_pts(path, pts: np.ndarray, file_type="txt"):
if file_type == "txt": if file_type == "txt":
np.savetxt(path, pts) 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")): if not os.path.exists(os.path.join(root,scene, "pts")):
os.makedirs(os.path.join(root,scene, "pts")) os.makedirs(os.path.join(root,scene, "pts"))
save_np_pts(pts_path, target_points, file_type) 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"): 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}") indices_path = os.path.join(root,scene, "scan_points_indices", f"{frame_idx}.{file_type}")
@@ -49,6 +53,26 @@ def get_world_points(depth, mask, cam_intrinsic, cam_extrinsic, random_downsampl
return points_camera_world return points_camera_world
def get_world_points_and_normal(depth, mask, normal, cam_intrinsic, cam_extrinsic, random_downsample_N):
z = depth[mask]
i, j = np.nonzero(mask)
x = (j - cam_intrinsic[0, 2]) * z / cam_intrinsic[0, 0]
y = (i - cam_intrinsic[1, 2]) * z / cam_intrinsic[1, 1]
points_camera = np.stack((x, y, z), axis=-1).reshape(-1, 3)
normal_camera = normal[mask].reshape(-1, 3)
sampled_target_points, idx = PtsUtil.random_downsample_point_cloud(
points_camera, random_downsample_N, require_idx=True
)
if len(sampled_target_points) == 0:
return np.zeros((0, 3)), np.zeros((0, 3))
sampled_normal_camera = normal_camera[idx]
points_camera_aug = np.concatenate((sampled_target_points, np.ones((sampled_target_points.shape[0], 1))), axis=-1)
points_camera_world = np.dot(cam_extrinsic, points_camera_aug.T).T[:, :3]
return points_camera_world, sampled_normal_camera
def get_scan_points_indices(scan_points, mask, display_table_mask_label, cam_intrinsic, cam_extrinsic): def get_scan_points_indices(scan_points, mask, display_table_mask_label, cam_intrinsic, cam_extrinsic):
scan_points_homogeneous = np.hstack((scan_points, np.ones((scan_points.shape[0], 1)))) scan_points_homogeneous = np.hstack((scan_points, np.ones((scan_points.shape[0], 1))))
points_camera = np.dot(np.linalg.inv(cam_extrinsic), scan_points_homogeneous.T).T[:, :3] points_camera = np.dot(np.linalg.inv(cam_extrinsic), scan_points_homogeneous.T).T[:, :3]
@@ -67,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"): def save_scene_data(root, scene, scene_idx=0, scene_total=1,file_type="txt"):
''' configuration ''' ''' configuration '''
target_mask_label = (0, 255, 0, 255) target_mask_label = (0, 255, 0)
display_table_mask_label=(0, 0, 255, 255) display_table_mask_label=(0, 0, 255)
random_downsample_N = 32768 random_downsample_N = 32768
voxel_size=0.003 voxel_size=0.003
filter_degree = 75 filter_degree = 75
@@ -93,7 +117,7 @@ def save_scene_data(root, scene, scene_idx=0, scene_total=1,file_type="txt"):
binocular=True binocular=True
) )
mask_L, mask_R = DataLoadUtil.load_seg(path, binocular=True) mask_L, mask_R = DataLoadUtil.load_seg(path, binocular=True)
normal_L = DataLoadUtil.load_normal(path, binocular=True, left_only=True)
''' target points ''' ''' target points '''
mask_img_L = mask_L mask_img_L = mask_L
mask_img_R = mask_R mask_img_R = mask_R
@@ -102,23 +126,23 @@ def save_scene_data(root, scene, scene_idx=0, scene_total=1,file_type="txt"):
target_mask_img_R = (mask_R == target_mask_label).all(axis=-1) target_mask_img_R = (mask_R == target_mask_label).all(axis=-1)
sampled_target_points_L = get_world_points(depth_L, target_mask_img_L, cam_info["cam_intrinsic"], cam_info["cam_to_world"], random_downsample_N) sampled_target_points_L, sampled_target_normal_L = get_world_points_and_normal(depth_L,target_mask_img_L,normal_L, cam_info["cam_intrinsic"], cam_info["cam_to_world"], random_downsample_N)
sampled_target_points_R = get_world_points(depth_R, target_mask_img_R, cam_info["cam_intrinsic"], cam_info["cam_to_world_R"], random_downsample_N) sampled_target_points_R = get_world_points(depth_R, target_mask_img_R, cam_info["cam_intrinsic"], cam_info["cam_to_world_R"], random_downsample_N)
has_points = sampled_target_points_L.shape[0] > 0 and sampled_target_points_R.shape[0] > 0 has_points = sampled_target_points_L.shape[0] > 0 and sampled_target_points_R.shape[0] > 0
if has_points: if has_points:
target_points = PtsUtil.get_overlapping_points( target_points, overlap_idx = PtsUtil.get_overlapping_points(
sampled_target_points_L, sampled_target_points_R, voxel_size sampled_target_points_L, sampled_target_points_R, voxel_size, require_idx=True
) )
sampled_target_normal_L = sampled_target_normal_L[overlap_idx]
if has_points: if has_points:
has_points = target_points.shape[0] > 0 has_points = target_points.shape[0] > 0
if has_points: if has_points:
points_normals = DataLoadUtil.load_points_normals(root, scene, display_table_as_world_space_origin=True) target_points, target_normals = PtsUtil.filter_points(
target_points = PtsUtil.filter_points( target_points, sampled_target_normal_L, cam_info["cam_to_world"], theta_limit = filter_degree, z_range=(min_z, max_z)
target_points, points_normals, cam_info["cam_to_world"],voxel_size=0.002, theta = filter_degree, z_range=(min_z, max_z)
) )
@@ -129,8 +153,10 @@ def save_scene_data(root, scene, scene_idx=0, scene_total=1,file_type="txt"):
if not has_points: if not has_points:
target_points = np.zeros((0, 3)) 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_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_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 save_scan_points(root, scene, scan_points) # The "done" flag of scene preprocess
@@ -138,17 +164,10 @@ def save_scene_data(root, scene, scene_idx=0, scene_total=1,file_type="txt"):
if __name__ == "__main__": if __name__ == "__main__":
#root = "/media/hofee/repository/new_data_with_normal" #root = "/media/hofee/repository/new_data_with_normal"
root = r"/media/hofee/repository/full_data_output" root = "/media/hofee/data/data/new_testset"
# 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())
scene_list = os.listdir(root) scene_list = os.listdir(root)
from_idx = 0 # 1000 from_idx = 0 # 1000
to_idx = 700 # 1500 to_idx = len(scene_list) # 1500
print(scene_list)
cnt = 0 cnt = 0
@@ -156,7 +175,15 @@ if __name__ == "__main__":
total = to_idx - from_idx total = to_idx - from_idx
for scene in scene_list[from_idx:to_idx]: for scene in scene_list[from_idx:to_idx]:
start = time.time() start = time.time()
save_scene_data(root, scene, cnt, total, file_type="npy") if os.path.exists(os.path.join(root, scene, "scan_points.txt")):
print(f"Scene {scene} has been processed")
cnt+=1
continue
try:
save_scene_data(root, scene, cnt, total, file_type="npy")
except Exception as e:
print(f"Error occurred when processing scene {scene}")
print(e)
cnt+=1 cnt+=1
end = time.time() end = time.time()
print(f"Time cost: {end-start}") print(f"Time cost: {end-start}")

119
runners/inference_server.py Normal file
View File

@@ -0,0 +1,119 @@
import os
import json
import torch
import numpy as np
from flask import Flask, request, jsonify
import PytorchBoot.namespace as namespace
import PytorchBoot.stereotype as stereotype
from PytorchBoot.factory import ComponentFactory
from PytorchBoot.runners.runner import Runner
from PytorchBoot.utils import Log
from utils.pts import PtsUtil
@stereotype.runner("inferencer_server")
class InferencerServer(Runner):
def __init__(self, config_path):
super().__init__(config_path)
''' Web Server '''
self.app = Flask(__name__)
''' 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)
self.pts_num = 8192
''' Experiment '''
self.load_experiment("inferencer_server")
def get_input_data(self, data):
input_data = {}
scanned_pts = data["scanned_pts"]
scanned_n_to_world_pose_9d = data["scanned_n_to_world_pose_9d"]
combined_scanned_views_pts = np.concatenate(scanned_pts, axis=0)
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
# fps_downsampled_combined_scanned_pts_mask = combined_scanned_views_pts_mask[fps_idx]
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):
pred_pose_9d = output_data["pred_pose_9d"]
result = {
"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))
@self.app.route("/inference", methods=["POST"])
def inference():
data = request.json
input_data = self.get_input_data(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)
self.app.run(host="0.0.0.0", port=5000)
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)

View File

@@ -19,7 +19,7 @@ from PytorchBoot.dataset import BaseDataset
from PytorchBoot.runners.runner import Runner from PytorchBoot.runners.runner import Runner
from PytorchBoot.utils import Log from PytorchBoot.utils import Log
from PytorchBoot.status import status_manager from PytorchBoot.status import status_manager
from utils.data_load import DataLoadUtil
@stereotype.runner("inferencer") @stereotype.runner("inferencer")
class Inferencer(Runner): class Inferencer(Runner):
def __init__(self, config_path): def __init__(self, config_path):
@@ -27,6 +27,7 @@ class Inferencer(Runner):
self.script_path = ConfigManager.get(namespace.Stereotype.RUNNER, "blender_script_path") self.script_path = ConfigManager.get(namespace.Stereotype.RUNNER, "blender_script_path")
self.output_dir = ConfigManager.get(namespace.Stereotype.RUNNER, "output_dir") self.output_dir = ConfigManager.get(namespace.Stereotype.RUNNER, "output_dir")
self.voxel_size = ConfigManager.get(namespace.Stereotype.RUNNER, "voxel_size")
''' Pipeline ''' ''' Pipeline '''
self.pipeline_name = self.config[namespace.Stereotype.PIPELINE] self.pipeline_name = self.config[namespace.Stereotype.PIPELINE]
self.pipeline:torch.nn.Module = ComponentFactory.create(namespace.Stereotype.PIPELINE, self.pipeline_name) self.pipeline:torch.nn.Module = ComponentFactory.create(namespace.Stereotype.PIPELINE, self.pipeline_name)
@@ -34,7 +35,12 @@ class Inferencer(Runner):
''' Experiment ''' ''' Experiment '''
self.load_experiment("nbv_evaluator") 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 ''' ''' Test '''
self.test_config = ConfigManager.get(namespace.Stereotype.RUNNER, namespace.Mode.TEST) self.test_config = ConfigManager.get(namespace.Stereotype.RUNNER, namespace.Mode.TEST)
@@ -65,59 +71,71 @@ class Inferencer(Runner):
for dataset_idx, test_set in enumerate(self.test_set_list): 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)) status_manager.set_progress("inference", "inferencer", f"dataset", dataset_idx, len(self.test_set_list))
test_set_name = test_set.get_name() test_set_name = test_set.get_name()
test_loader = test_set.get_loader()
if test_loader.batch_size > 1: total=int(len(test_set))
Log.error("Batch size should be 1 for inference, found {} in {}".format(test_loader.batch_size, test_set_name), terminate=True) for i in tqdm(range(total), desc=f"Processing {test_set_name}", ncols=100):
data = test_set.__getitem__(i)
total=int(len(test_loader)) scene_name = data["scene_name"]
loop = tqdm(enumerate(test_loader), total=total) if scene_name != "omniobject3d-book_004":
for i, data in loop: continue
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
status_manager.set_progress("inference", "inferencer", f"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)
output = self.predict_sequence(data) output = self.predict_sequence(data)
self.save_inference_result(test_set_name, data["scene_name"][0], output) 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)) 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, overlap_area_threshold=25, scan_points_threshold=10, max_iter=50, max_retry = 5):
scene_name = data["scene_name"][0] scene_name = data["scene_name"]
Log.info(f"Processing scene: {scene_name}") Log.info(f"Processing scene: {scene_name}")
status_manager.set_status("inference", "inferencer", "scene", scene_name) status_manager.set_status("inference", "inferencer", "scene", scene_name)
''' data for rendering ''' ''' data for rendering '''
scene_path = data["scene_path"][0] scene_path = data["scene_path"]
O_to_L_pose = data["O_to_L_pose"][0] O_to_L_pose = data["O_to_L_pose"]
voxel_threshold = data["voxel_threshold"][0] voxel_threshold = self.voxel_size
filter_degree = data["filter_degree"][0] filter_degree = 75
model_points_normals = data["model_points_normals"][0] down_sampled_model_pts = data["gt_pts"]
model_pts = model_points_normals[:,:3]
down_sampled_model_pts = PtsUtil.voxel_downsample_point_cloud(model_pts, voxel_threshold) first_frame_to_world_9d = data["first_scanned_n_to_world_pose_9d"][0]
first_frame_to_world_9d = data["first_to_world_9d"][0] first_frame_to_world = np.eye(4)
first_frame_to_world = torch.eye(4, device=first_frame_to_world_9d.device) first_frame_to_world[:3,:3] = PoseUtil.rotation_6d_to_matrix_numpy(first_frame_to_world_9d[:6])
first_frame_to_world[:3,:3] = PoseUtil.rotation_6d_to_matrix_tensor_batch(first_frame_to_world_9d[:,:6])[0] first_frame_to_world[:3,3] = first_frame_to_world_9d[6:]
first_frame_to_world[:3,3] = first_frame_to_world_9d[0,6:]
first_frame_to_world = first_frame_to_world.to(self.device)
''' data for inference ''' ''' data for inference '''
input_data = {} input_data = {}
input_data["scanned_pts"] = [data["first_pts"][0].to(self.device)] 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"] = [data["first_to_world_9d"][0].to(self.device)] 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_data["mode"] = namespace.Mode.TEST
input_data["combined_scanned_pts"] = data["combined_scanned_pts"] input_pts_N = input_data["combined_scanned_pts"].shape[1]
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) root = os.path.dirname(scene_path)
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) 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)
retry_duplication_pose = [] retry_duplication_pose = []
retry_no_pts_pose = [] retry_no_pts_pose = []
retry_overlap_pose = []
retry = 0 retry = 0
pred_cr_seq = [last_pred_cr] pred_cr_seq = [last_pred_cr]
success = 0
last_pts_num = PtsUtil.voxel_downsample_point_cloud(data["first_scanned_pts"][0], 0.002).shape[0]
import time
while len(pred_cr_seq) < max_iter and retry < max_retry: while len(pred_cr_seq) < max_iter and retry < max_retry:
start_time = time.time()
output = self.pipeline(input_data) 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_9d = output["pred_pose_9d"]
pred_pose = torch.eye(4, device=pred_pose_9d.device) pred_pose = torch.eye(4, device=pred_pose_9d.device)
@@ -125,7 +143,24 @@ class Inferencer(Runner):
pred_pose[:3,3] = pred_pose_9d[0,6:] pred_pose[:3,3] = pred_pose_9d[0,6:]
try: try:
new_target_pts_world, new_pts_world = RenderUtil.render_pts(pred_pose, 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) start_time = time.time()
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, _ = 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
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: except Exception as e:
Log.warning(f"Error in scene {scene_path}, {e}") Log.warning(f"Error in scene {scene_path}, {e}")
print("current pose: ", pred_pose) print("current pose: ", pred_pose)
@@ -134,61 +169,42 @@ class Inferencer(Runner):
retry += 1 retry += 1
continue continue
if new_target_pts.shape[0] == 0:
pred_cr = self.compute_coverage_rate(scanned_view_pts, new_target_pts_world, down_sampled_model_pts, threshold=voxel_threshold) print("no pts in new target")
retry_no_pts_pose.append(pred_pose.cpu().numpy().tolist())
print(pred_cr, last_pred_cr, " max: ", data["max_coverage_rate"])
if pred_cr >= data["max_coverage_rate"]:
print("max coverage rate reached!")
if pred_cr <= last_pred_cr + cr_increase_threshold:
retry += 1 retry += 1
retry_duplication_pose.append(pred_pose.cpu().numpy().tolist())
continue continue
start_time = time.time()
pred_cr, _ = 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)
success += 1
retry = 0 retry = 0
pred_cr_seq.append(pred_cr) pred_cr_seq.append(pred_cr)
scanned_view_pts.append(new_target_pts_world) scanned_view_pts.append(new_target_pts)
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_world_pose_9d"] = [torch.cat([input_data["scanned_n_to_world_pose_9d"][0], pred_pose_9d], dim=0)] input_data["scanned_n_to_world_pose_9d"] = [torch.cat([input_data["scanned_n_to_world_pose_9d"][0], pred_pose_9d], dim=0)]
combined_scanned_views_pts = np.concatenate(input_data["scanned_pts"][0].tolist(), axis=0)
voxel_downsampled_combined_scanned_pts_np = PtsUtil.voxel_downsample_point_cloud(combined_scanned_views_pts, 0.002) combined_scanned_pts = np.vstack(scanned_view_pts)
voxel_downsampled_combined_scanned_pts_np = PtsUtil.voxel_downsample_point_cloud(combined_scanned_pts, 0.002)
random_downsampled_combined_scanned_pts_np = PtsUtil.random_downsample_point_cloud(voxel_downsampled_combined_scanned_pts_np, input_pts_N) random_downsampled_combined_scanned_pts_np = PtsUtil.random_downsample_point_cloud(voxel_downsampled_combined_scanned_pts_np, input_pts_N)
input_data["combined_scanned_pts"] = torch.tensor(random_downsampled_combined_scanned_pts_np, dtype=torch.float32).unsqueeze(0).to(self.device) input_data["combined_scanned_pts"] = torch.tensor(random_downsampled_combined_scanned_pts_np, dtype=torch.float32).unsqueeze(0).to(self.device)
if success > 3:
break
last_pred_cr = pred_cr last_pred_cr = pred_cr
pts_num = voxel_downsampled_combined_scanned_pts_np.shape[0]
if pts_num - last_pts_num < 10 and pred_cr < data["seq_max_coverage_rate"] - 1e-3:
input_data["scanned_pts"] = input_data["scanned_pts"][0].cpu().numpy().tolist() retry += 1
input_data["scanned_n_to_world_pose_9d"] = input_data["scanned_n_to_world_pose_9d"][0].cpu().numpy().tolist() retry_duplication_pose.append(pred_pose.cpu().numpy().tolist())
result = { print("delta pts num < 10:", pts_num, last_pts_num)
"pred_pose_9d_seq": input_data["scanned_n_to_world_pose_9d"], last_pts_num = pts_num
"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"][0],
"pred_max_coverage_rate": max(pred_cr_seq),
"scene_name": scene_name,
"retry_no_pts_pose": retry_no_pts_pose,
"retry_duplication_pose": retry_duplication_pose,
"best_seq_len": data["best_seq_len"][0],
}
self.stat_result[scene_name] = {
"max_coverage_rate": data["max_coverage_rate"][0],
"success_rate": max(pred_cr_seq)/ data["max_coverage_rate"][0],
"coverage_rate_seq": pred_cr_seq,
"pred_max_coverage_rate": max(pred_cr_seq),
"pred_seq_len": len(pred_cr_seq),
}
print('success rate: ', max(pred_cr_seq) / data["max_coverage_rate"][0])
return result
def compute_coverage_rate(self, scanned_view_pts, new_pts, model_pts, threshold=0.005): def compute_coverage_rate(self, scanned_view_pts, new_pts, model_pts, threshold=0.005):
if new_pts is not None: if new_pts is not None:
@@ -206,7 +222,7 @@ class Inferencer(Runner):
os.makedirs(dataset_dir) os.makedirs(dataset_dir)
output_path = os.path.join(dataset_dir, f"{scene_name}.pkl") output_path = os.path.join(dataset_dir, f"{scene_name}.pkl")
pickle.dump(output, open(output_path, "wb")) 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) json.dump(self.stat_result, f)

View File

@@ -22,25 +22,21 @@ class StrategyGenerator(Runner):
"app_name": "generate_strategy", "app_name": "generate_strategy",
"runner_name": "strategy_generator" "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.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.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): def run(self):
dataset_name_list = ConfigManager.get("runner", "generate", "dataset_list") 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)): for dataset_idx in range(len(dataset_name_list)):
dataset_name = dataset_name_list[dataset_idx] dataset_name = dataset_name_list[dataset_idx]
status_manager.set_progress("generate_strategy", "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") root_dir = ConfigManager.get("datasets", dataset_name, "root_dir")
model_dir = ConfigManager.get("datasets", dataset_name, "model_dir")
from_idx = ConfigManager.get("datasets",dataset_name,"from") from_idx = ConfigManager.get("datasets",dataset_name,"from")
to_idx = ConfigManager.get("datasets",dataset_name,"to") to_idx = ConfigManager.get("datasets",dataset_name,"to")
scene_name_list = os.listdir(root_dir) scene_name_list = os.listdir(root_dir)
@@ -52,17 +48,13 @@ class StrategyGenerator(Runner):
for scene_name in scene_name_list[from_idx:to_idx]: for scene_name in scene_name_list[from_idx:to_idx]:
Log.info(f"({dataset_name})Processing [{cnt}/{total}]: {scene_name}") Log.info(f"({dataset_name})Processing [{cnt}/{total}]: {scene_name}")
status_manager.set_progress("generate_strategy", "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)
status_manager.set_status("generate_strategy", "strategy_generator", "diagonal", diag)
output_label_path = DataLoadUtil.get_label_path(root_dir, scene_name,0) output_label_path = DataLoadUtil.get_label_path(root_dir, scene_name,0)
if os.path.exists(output_label_path) and not self.overwrite: if os.path.exists(output_label_path) and not self.overwrite:
Log.info(f"Scene <{scene_name}> Already Exists, Skip") Log.info(f"Scene <{scene_name}> Already Exists, Skip")
cnt += 1 cnt += 1
continue continue
self.generate_sequence(root_dir, model_dir, scene_name,voxel_threshold, soft_overlap_threshold, hard_overlap_threshold) self.generate_sequence(root_dir, scene_name,voxel_threshold)
# except Exception as e:
# Log.error(f"Scene <{scene_name}> Failed, Error: {e}")
cnt += 1 cnt += 1
status_manager.set_progress("generate_strategy", "strategy_generator", "scene", total, total) 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)) status_manager.set_progress("generate_strategy", "strategy_generator", "dataset", len(dataset_name_list), len(dataset_name_list))
@@ -75,27 +67,36 @@ class StrategyGenerator(Runner):
def load_experiment(self, backup_name=None): def load_experiment(self, backup_name=None):
super().load_experiment(backup_name) super().load_experiment(backup_name)
def generate_sequence(self, root, model_dir, 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) status_manager.set_status("generate_strategy", "strategy_generator", "scene", scene_name)
frame_num = DataLoadUtil.get_scene_seq_length(root, scene_name) frame_num = DataLoadUtil.get_scene_seq_length(root, scene_name)
model_points_normals = DataLoadUtil.load_points_normals(root, scene_name) model_points_normals = DataLoadUtil.load_points_normals(root, scene_name)
model_pts = model_points_normals[:,:3] 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 = [] pts_list = []
nrm_list = []
scan_points_indices_list = [] scan_points_indices_list = []
non_zero_cnt = 0 non_zero_cnt = 0
for frame_idx in range(frame_num): for frame_idx in range(frame_num):
status_manager.set_progress("generate_strategy", "strategy_generator", "loading frame", frame_idx, 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") 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") 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)
indices = np.load(idx_path)
pts_list.append(sampled_point_cloud)
pts = np.load(pts_path)
if self.compute_with_normal:
if pts.shape[0] == 0:
nrm = np.zeros((0,3))
else:
nrm = np.load(nrm_path)
nrm_list.append(nrm)
pts_list.append(pts)
indices = np.load(idx_path)
scan_points_indices_list.append(indices) scan_points_indices_list.append(indices)
if sampled_point_cloud.shape[0] > 0: if pts.shape[0] > 0:
non_zero_cnt += 1 non_zero_cnt += 1
status_manager.set_progress("generate_strategy", "strategy_generator", "loading frame", frame_num, frame_num) status_manager.set_progress("generate_strategy", "strategy_generator", "loading frame", frame_num, frame_num)
@@ -103,7 +104,7 @@ class StrategyGenerator(Runner):
init_view_list = [] init_view_list = []
idx = 0 idx = 0
while len(init_view_list) < seq_num and idx < len(pts_list): 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) init_view_list.append(idx)
idx += 1 idx += 1
@@ -112,8 +113,13 @@ class StrategyGenerator(Runner):
for init_view in init_view_list: for init_view in init_view_list:
status_manager.set_progress("generate_strategy", "strategy_generator", "computing sequence", seq_idx, len(init_view_list)) status_manager.set_progress("generate_strategy", "strategy_generator", "computing sequence", seq_idx, len(init_view_list))
start = time.time() 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() end = time.time()
print(f"Time: {end-start}") print(f"Time: {end-start}")
data_pairs = self.generate_data_pairs(limited_useful_view) data_pairs = self.generate_data_pairs(limited_useful_view)
@@ -132,8 +138,6 @@ class StrategyGenerator(Runner):
with open(output_label_path, 'w') as f: with open(output_label_path, 'w') as f:
json.dump(seq_save_data, f) json.dump(seq_save_data, f)
seq_idx += 1 seq_idx += 1
if self.save_mesh:
DataLoadUtil.save_target_mesh_at_world_space(root, model_dir, scene_name)
status_manager.set_progress("generate_strategy", "strategy_generator", "computing sequence", len(init_view_list), len(init_view_list)) status_manager.set_progress("generate_strategy", "strategy_generator", "computing sequence", len(init_view_list), len(init_view_list))

View File

@@ -9,7 +9,7 @@ class ViewGenerator(Runner):
self.config_path = config_path self.config_path = config_path
def run(self): 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() print()
def create_experiment(self, backup_name=None): def create_experiment(self, backup_name=None):

View File

@@ -4,11 +4,28 @@ import json
import cv2 import cv2
import trimesh import trimesh
import torch import torch
import OpenEXR
import Imath
from utils.pts import PtsUtil from utils.pts import PtsUtil
class DataLoadUtil: class DataLoadUtil:
TABLE_POSITION = np.asarray([0, 0, 0.8215]) TABLE_POSITION = np.asarray([0, 0, 0.8215])
@staticmethod
def load_exr_image(file_path):
exr_file = OpenEXR.InputFile(file_path)
header = exr_file.header()
dw = header['dataWindow']
width = dw.max.x - dw.min.x + 1
height = dw.max.y - dw.min.y + 1
float_channels = ['R', 'G', 'B']
img_data = []
for channel in float_channels:
channel_data = exr_file.channel(channel)
img_data.append(np.frombuffer(channel_data, dtype=np.float16).reshape((height, width)))
img = np.stack(img_data, axis=-1)
return img
@staticmethod @staticmethod
def get_display_table_info(root, scene_name): def get_display_table_info(root, scene_name):
@@ -34,6 +51,8 @@ class DataLoadUtil:
@staticmethod @staticmethod
def get_label_num(root, scene_name): def get_label_num(root, scene_name):
label_dir = os.path.join(root, scene_name, "label") label_dir = os.path.join(root, scene_name, "label")
if not os.path.exists(label_dir):
return 0
return len(os.listdir(label_dir)) return len(os.listdir(label_dir))
@staticmethod @staticmethod
@@ -44,11 +63,6 @@ class DataLoadUtil:
path = os.path.join(label_dir, f"{seq_idx}.json") path = os.path.join(label_dir, f"{seq_idx}.json")
return path return path
@staticmethod
def get_label_path_old(root, scene_name):
path = os.path.join(root, scene_name, "label.json")
return path
@staticmethod @staticmethod
def get_scene_seq_length(root, scene_name): def get_scene_seq_length(root, scene_name):
camera_params_path = os.path.join(root, scene_name, "camera_params") camera_params_path = os.path.join(root, scene_name, "camera_params")
@@ -69,36 +83,6 @@ class DataLoadUtil:
diagonal_length = np.linalg.norm(bbox) diagonal_length = np.linalg.norm(bbox)
return diagonal_length return diagonal_length
@staticmethod
def save_mesh_at(model_dir, output_dir, object_name, scene_name, world_object_pose):
mesh = DataLoadUtil.load_mesh_at(model_dir, object_name, world_object_pose)
model_path = os.path.join(output_dir, scene_name, "world_mesh.obj")
mesh.export(model_path)
@staticmethod
def save_target_mesh_at_world_space(
root, model_dir, scene_name, display_table_as_world_space_origin=True
):
scene_info = DataLoadUtil.load_scene_info(root, scene_name)
target_name = scene_info["target_name"]
transformation = scene_info[target_name]
if display_table_as_world_space_origin:
location = transformation["location"] - DataLoadUtil.get_display_table_top(
root, scene_name
)
else:
location = transformation["location"]
rotation_euler = transformation["rotation_euler"]
pose_mat = trimesh.transformations.euler_matrix(*rotation_euler)
pose_mat[:3, 3] = location
mesh = DataLoadUtil.load_mesh_at(model_dir, target_name, pose_mat)
mesh_dir = os.path.join(root, scene_name, "mesh")
if not os.path.exists(mesh_dir):
os.makedirs(mesh_dir)
model_path = os.path.join(mesh_dir, "world_target_mesh.obj")
mesh.export(model_path)
@staticmethod @staticmethod
def load_scene_info(root, scene_name): def load_scene_info(root, scene_name):
scene_info_path = os.path.join(root, scene_name, "scene_info.json") scene_info_path = os.path.join(root, scene_name, "scene_info.json")
@@ -113,17 +97,6 @@ class DataLoadUtil:
target_pts_num_dict = json.load(f) target_pts_num_dict = json.load(f)
return target_pts_num_dict return target_pts_num_dict
@staticmethod
def load_target_object_pose(root, scene_name):
scene_info = DataLoadUtil.load_scene_info(root, scene_name)
target_name = scene_info["target_name"]
transformation = scene_info[target_name]
location = transformation["location"]
rotation_euler = transformation["rotation_euler"]
pose_mat = trimesh.transformations.euler_matrix(*rotation_euler)
pose_mat[:3, 3] = location
return pose_mat
@staticmethod @staticmethod
def load_depth(path, min_depth=0.01, max_depth=5.0, binocular=False): def load_depth(path, min_depth=0.01, max_depth=5.0, binocular=False):
@@ -161,8 +134,8 @@ class DataLoadUtil:
if binocular and not left_only: if binocular and not left_only:
def clean_mask(mask_image): def clean_mask(mask_image):
green = [0, 255, 0, 255] green = [0, 255, 0]
red = [255, 0, 0, 255] red = [255, 0, 0]
threshold = 2 threshold = 2
mask_image = np.where( mask_image = np.where(
np.abs(mask_image - green) <= threshold, green, mask_image np.abs(mask_image - green) <= threshold, green, mask_image
@@ -194,30 +167,31 @@ class DataLoadUtil:
return mask_image return mask_image
@staticmethod @staticmethod
def load_normal(path, binocular=False, left_only=False): def load_normal(path, binocular=False, left_only=False, file_type="exr"):
if binocular and not left_only: if binocular and not left_only:
normal_path_L = os.path.join( normal_path_L = os.path.join(
os.path.dirname(path), "normal", os.path.basename(path) + "_L.png" os.path.dirname(path), "normal", os.path.basename(path) + f"_L.{file_type}"
) )
normal_image_L = cv2.imread(normal_path_L, cv2.IMREAD_COLOR) normal_image_L = DataLoadUtil.load_exr_image(normal_path_L)
normal_path_R = os.path.join( normal_path_R = os.path.join(
os.path.dirname(path), "normal", os.path.basename(path) + "_R.png" os.path.dirname(path), "normal", os.path.basename(path) + f"_R.{file_type}"
) )
normal_image_R = cv2.imread(normal_path_R, cv2.IMREAD_COLOR) normal_image_R = DataLoadUtil.load_exr_image(normal_path_R)
normalized_normal_image_L = normal_image_L / 255.0 * 2.0 - 1.0 normalized_normal_image_L = normal_image_L * 2.0 - 1.0
normalized_normal_image_R = normal_image_R / 255.0 * 2.0 - 1.0 normalized_normal_image_R = normal_image_R * 2.0 - 1.0
return normalized_normal_image_L, normalized_normal_image_R return normalized_normal_image_L, normalized_normal_image_R
else: else:
if binocular and left_only: if binocular and left_only:
normal_path = os.path.join( normal_path = os.path.join(
os.path.dirname(path), "normal", os.path.basename(path) + "_L.png" os.path.dirname(path), "normal", os.path.basename(path) + f"_L.{file_type}"
) )
else: else:
normal_path = os.path.join( normal_path = os.path.join(
os.path.dirname(path), "normal", os.path.basename(path) + ".png" os.path.dirname(path), "normal", os.path.basename(path) + f".{file_type}"
) )
normal_image = cv2.imread(normal_path, cv2.IMREAD_COLOR) normal_image = DataLoadUtil.load_exr_image(normal_path)
normalized_normal_image = normal_image / 255.0 * 2.0 - 1.0 normalized_normal_image = normal_image * 2.0 - 1.0
return normalized_normal_image return normalized_normal_image
@staticmethod @staticmethod
@@ -227,20 +201,26 @@ class DataLoadUtil:
return label_data return label_data
@staticmethod @staticmethod
def load_rgb(path): def load_from_preprocessed_pts(path, file_type="npy"):
rgb_path = os.path.join(
os.path.dirname(path), "rgb", os.path.basename(path) + ".png"
)
rgb_image = cv2.imread(rgb_path, cv2.IMREAD_COLOR)
return rgb_image
@staticmethod
def load_from_preprocessed_pts(path):
npy_path = os.path.join( npy_path = os.path.join(
os.path.dirname(path), "pts", os.path.basename(path) + ".npy" os.path.dirname(path), "pts", os.path.basename(path) + "." + file_type
) )
pts = np.load(npy_path) if file_type == "txt":
pts = np.loadtxt(npy_path)
else:
pts = np.load(npy_path)
return pts return pts
@staticmethod
def load_from_preprocessed_nrm(path, file_type="npy"):
npy_path = os.path.join(
os.path.dirname(path), "nrm", os.path.basename(path) + "." + file_type
)
if file_type == "txt":
nrm = np.loadtxt(npy_path)
else:
nrm = np.load(npy_path)
return nrm
@staticmethod @staticmethod
def cam_pose_transformation(cam_pose_before): def cam_pose_transformation(cam_pose_before):
@@ -260,11 +240,12 @@ class DataLoadUtil:
label_data = json.load(f) label_data = json.load(f)
cam_to_world = np.asarray(label_data["extrinsic"]) cam_to_world = np.asarray(label_data["extrinsic"])
cam_to_world = DataLoadUtil.cam_pose_transformation(cam_to_world) cam_to_world = DataLoadUtil.cam_pose_transformation(cam_to_world)
world_to_display_table = np.eye(4)
world_to_display_table[:3, 3] = -DataLoadUtil.get_display_table_top(
root_dir, scene_name
)
if display_table_as_world_space_origin: if display_table_as_world_space_origin:
world_to_display_table = np.eye(4)
world_to_display_table[:3, 3] = -DataLoadUtil.get_display_table_top(
root_dir, scene_name
)
cam_to_world = np.dot(world_to_display_table, cam_to_world) cam_to_world = np.dot(world_to_display_table, cam_to_world)
cam_intrinsic = np.asarray(label_data["intrinsic"]) cam_intrinsic = np.asarray(label_data["intrinsic"])
cam_info = { cam_info = {

View File

@@ -1,16 +1,32 @@
import numpy as np import numpy as np
import open3d as o3d import open3d as o3d
import torch import torch
from scipy.spatial import cKDTree
class PtsUtil: class PtsUtil:
@staticmethod @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):
o3d_pc = o3d.geometry.PointCloud() voxel_indices = np.floor(point_cloud / voxel_size).astype(np.int32)
o3d_pc.points = o3d.utility.Vector3dVector(point_cloud) if require_idx:
downsampled_pc = o3d_pc.voxel_down_sample(voxel_size) _, inverse, counts = np.unique(voxel_indices, axis=0, return_inverse=True, return_counts=True)
return np.asarray(downsampled_pc.points) 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 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 @staticmethod
def random_downsample_point_cloud(point_cloud, num_points, require_idx=False): def random_downsample_point_cloud(point_cloud, num_points, require_idx=False):
@@ -78,24 +94,24 @@ class PtsUtil:
return overlapping_points return overlapping_points
@staticmethod @staticmethod
def filter_points(points, points_normals, cam_pose, voxel_size=0.002, theta=45, z_range=(0.2, 0.45)): def filter_points(points, normals, cam_pose, theta_limit=45, z_range=(0.2, 0.45)):
""" filter with normal """
normals_normalized = normals / np.linalg.norm(normals, axis=1, keepdims=True)
cos_theta = np.dot(normals_normalized, np.array([0, 0, 1]))
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 """ """ filter with z range """
points_cam = PtsUtil.transform_point_cloud(points, np.linalg.inv(cam_pose)) 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]) idx = (points_cam[:, 2] > z_range[0]) & (points_cam[:, 2] < z_range[1])
z_filtered_points = points[idx] z_filtered_points = filtered_sampled_points[idx]
z_filtered_normals = filtered_normals[idx]
""" filter with normal """ return z_filtered_points[:, :3], z_filtered_normals
sampled_points = PtsUtil.voxel_downsample_point_cloud(z_filtered_points, voxel_size)
kdtree = cKDTree(points_normals[:,:3]) @staticmethod
_, indices = kdtree.query(sampled_points) def point_to_hash(point, voxel_size):
nearest_points = points_normals[indices] return tuple(np.floor(point / voxel_size).astype(int))
normals = nearest_points[:, 3:]
camera_axis = -cam_pose[:3, 2]
normals_normalized = normals / np.linalg.norm(normals, axis=1, keepdims=True)
cos_theta = np.dot(normals_normalized, camera_axis)
theta_rad = np.deg2rad(theta)
idx = cos_theta > np.cos(theta_rad)
filtered_sampled_points= sampled_points[idx]
return filtered_sampled_points[:, :3]

View File

@@ -3,25 +3,45 @@ from scipy.spatial import cKDTree
from utils.pts import PtsUtil from utils.pts import PtsUtil
class ReconstructionUtil: class ReconstructionUtil:
@staticmethod @staticmethod
def compute_coverage_rate(target_point_cloud, combined_point_cloud, threshold=0.01): def compute_coverage_rate(target_point_cloud, combined_point_cloud, threshold=0.01):
kdtree = cKDTree(combined_point_cloud) kdtree = cKDTree(combined_point_cloud)
distances, _ = kdtree.query(target_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] coverage_rate = covered_points_num / target_point_cloud.shape[0]
return coverage_rate, covered_points_num return coverage_rate, covered_points_num
@staticmethod @staticmethod
def compute_overlap_rate(new_point_cloud, combined_point_cloud, threshold=0.01): 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*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]
return coverage_rate, covered_points_num
@staticmethod
def check_overlap(new_point_cloud, combined_point_cloud, overlap_area_threshold=25, voxel_size=0.01, require_new_added_pts_num=False):
kdtree = cKDTree(combined_point_cloud) kdtree = cKDTree(combined_point_cloud)
distances, _ = kdtree.query(new_point_cloud) distances, _ = kdtree.query(new_point_cloud)
overlapping_points = np.sum(distances < threshold) overlapping_points_num = np.sum(distances < voxel_size*2)
if new_point_cloud.shape[0] == 0: cm = 0.01
overlap_rate = 0 voxel_size_cm = voxel_size / cm
else: overlap_area = overlapping_points_num * voxel_size_cm * voxel_size_cm
overlap_rate = overlapping_points / new_point_cloud.shape[0] if require_new_added_pts_num:
return overlap_rate return overlap_area > overlap_area_threshold, len(new_point_cloud)-np.sum(distances < voxel_size*1.2)
return overlap_area > overlap_area_threshold
@staticmethod @staticmethod
@@ -37,14 +57,14 @@ class ReconstructionUtil:
return new_added_points return new_added_points
@staticmethod @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] selected_views = [init_view]
combined_point_cloud = point_cloud_list[init_view] combined_point_cloud = point_cloud_list[init_view]
history_indices = [scan_points_indices_list[init_view]] history_indices = [scan_points_indices_list[init_view]]
max_rec_pts = np.vstack(point_cloud_list) max_rec_pts = np.vstack(point_cloud_list)
downsampled_max_rec_pts = PtsUtil.voxel_downsample_point_cloud(max_rec_pts, threshold) downsampled_max_rec_pts = PtsUtil.voxel_downsample_point_cloud(max_rec_pts, threshold)
combined_point_cloud = PtsUtil.voxel_downsample_point_cloud(combined_point_cloud, threshold)
max_rec_pts_num = downsampled_max_rec_pts.shape[0] max_rec_pts_num = downsampled_max_rec_pts.shape[0]
max_real_rec_pts_coverage, _ = ReconstructionUtil.compute_coverage_rate(target_point_cloud, downsampled_max_rec_pts, threshold) max_real_rec_pts_coverage, _ = ReconstructionUtil.compute_coverage_rate(target_point_cloud, downsampled_max_rec_pts, threshold)
@@ -57,6 +77,7 @@ class ReconstructionUtil:
cnt_processed_view = 0 cnt_processed_view = 0
remaining_views.remove(init_view) remaining_views.remove(init_view)
curr_rec_pts_num = combined_point_cloud.shape[0] curr_rec_pts_num = combined_point_cloud.shape[0]
drop_output_ratio = 0.4
import time import time
while remaining_views: while remaining_views:
@@ -66,27 +87,23 @@ class ReconstructionUtil:
best_covered_num = 0 best_covered_num = 0
for view_index in remaining_views: for view_index in remaining_views:
if np.random.rand() < drop_output_ratio:
continue
if point_cloud_list[view_index].shape[0] == 0: if point_cloud_list[view_index].shape[0] == 0:
continue continue
if selected_views: if selected_views:
new_scan_points_indices = scan_points_indices_list[view_index] 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): 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: else:
overlap_threshold = soft_overlap_threshold curr_overlap_area_threshold = overlap_area_threshold * 0.5
start = time.time()
overlap_rate = ReconstructionUtil.compute_overlap_rate(point_cloud_list[view_index],combined_point_cloud, threshold) if not ReconstructionUtil.check_overlap(point_cloud_list[view_index], combined_point_cloud, overlap_area_threshold = curr_overlap_area_threshold, voxel_size=threshold):
end = time.time()
# print(f"overlap_rate Time: {end-start}")
if overlap_rate < overlap_threshold:
continue continue
start = time.time()
new_combined_point_cloud = np.vstack([combined_point_cloud, point_cloud_list[view_index]]) 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_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) 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 coverage_increase = new_coverage - current_coverage
if coverage_increase > best_coverage_increase: if coverage_increase > best_coverage_increase:
best_coverage_increase = coverage_increase best_coverage_increase = coverage_increase
@@ -95,6 +112,100 @@ class ReconstructionUtil:
best_combined_point_cloud = new_downsampled_combined_point_cloud 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_view is not None:
if best_coverage_increase <=1e-3 or best_covered_num - current_covered_num <= 5: if best_coverage_increase <=1e-3 or best_covered_num - current_covered_num <= 5:
break break
@@ -106,6 +217,7 @@ class ReconstructionUtil:
current_covered_num = best_covered_num current_covered_num = best_covered_num
curr_rec_pts_num = best_rec_pts_num curr_rec_pts_num = best_rec_pts_num
combined_point_cloud = best_combined_point_cloud combined_point_cloud = best_combined_point_cloud
combined_normal = best_combined_normal
remaining_views.remove(best_view) remaining_views.remove(best_view)
history_indices.append(scan_points_indices_list[best_view]) history_indices.append(scan_points_indices_list[best_view])
current_coverage += best_coverage_increase current_coverage += best_coverage_increase
@@ -145,18 +257,6 @@ class ReconstructionUtil:
attempts += 1 attempts += 1
return points return points
@staticmethod
def compute_covered_scan_points(scan_points, point_cloud, threshold=0.01):
tree = cKDTree(point_cloud)
covered_points = []
indices = []
for i, scan_point in enumerate(scan_points):
if tree.query_ball_point(scan_point, threshold):
covered_points.append(scan_point)
indices.append(i)
return covered_points, indices
@staticmethod @staticmethod
def check_scan_points_overlap(history_indices, indices2, threshold=5): def check_scan_points_overlap(history_indices, indices2, threshold=5):
for indices1 in history_indices: for indices1 in history_indices:

View File

@@ -1,16 +1,75 @@
import os import os
import json import json
import time
import subprocess import subprocess
import tempfile import tempfile
import shutil import shutil
import numpy as np
from utils.data_load import DataLoadUtil from utils.data_load import DataLoadUtil
from utils.reconstruction import ReconstructionUtil from utils.reconstruction import ReconstructionUtil
from utils.pts import PtsUtil from utils.pts import PtsUtil
class RenderUtil: class RenderUtil:
target_mask_label = (0, 255, 0)
display_table_mask_label = (0, 0, 255)
random_downsample_N = 32768
min_z = 0.2
max_z = 0.5
@staticmethod @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): def get_world_points_and_normal(depth, mask, normal, cam_intrinsic, cam_extrinsic, random_downsample_N):
z = depth[mask]
i, j = np.nonzero(mask)
x = (j - cam_intrinsic[0, 2]) * z / cam_intrinsic[0, 0]
y = (i - cam_intrinsic[1, 2]) * z / cam_intrinsic[1, 1]
points_camera = np.stack((x, y, z), axis=-1).reshape(-1, 3)
normal_camera = normal[mask].reshape(-1, 3)
sampled_target_points, idx = PtsUtil.random_downsample_point_cloud(
points_camera, random_downsample_N, require_idx=True
)
if len(sampled_target_points) == 0:
return np.zeros((0, 3)), np.zeros((0, 3))
sampled_normal_camera = normal_camera[idx]
points_camera_aug = np.concatenate((sampled_target_points, np.ones((sampled_target_points.shape[0], 1))), axis=-1)
points_camera_world = np.dot(cam_extrinsic, points_camera_aug.T).T[:, :3]
return points_camera_world, sampled_normal_camera
@staticmethod
def get_world_points(depth, mask, cam_intrinsic, cam_extrinsic, random_downsample_N):
z = depth[mask]
i, j = np.nonzero(mask)
x = (j - cam_intrinsic[0, 2]) * z / cam_intrinsic[0, 0]
y = (i - cam_intrinsic[1, 2]) * z / cam_intrinsic[1, 1]
points_camera = np.stack((x, y, z), axis=-1).reshape(-1, 3)
sampled_target_points = PtsUtil.random_downsample_point_cloud(
points_camera, random_downsample_N
)
points_camera_aug = np.concatenate((sampled_target_points, np.ones((sampled_target_points.shape[0], 1))), axis=-1)
points_camera_world = np.dot(cam_extrinsic, points_camera_aug.T).T[:, :3]
return points_camera_world
@staticmethod
def get_scan_points_indices(scan_points, mask, display_table_mask_label, cam_intrinsic, cam_extrinsic):
scan_points_homogeneous = np.hstack((scan_points, np.ones((scan_points.shape[0], 1))))
points_camera = np.dot(np.linalg.inv(cam_extrinsic), scan_points_homogeneous.T).T[:, :3]
points_image_homogeneous = np.dot(cam_intrinsic, points_camera.T).T
points_image_homogeneous /= points_image_homogeneous[:, 2:]
pixel_x = points_image_homogeneous[:, 0].astype(int)
pixel_y = points_image_homogeneous[:, 1].astype(int)
h, w = mask.shape[:2]
valid_indices = (pixel_x >= 0) & (pixel_x < w) & (pixel_y >= 0) & (pixel_y < h)
mask_colors = mask[pixel_y[valid_indices], pixel_x[valid_indices]]
selected_points_indices = np.where((mask_colors == display_table_mask_label).all(axis=-1))[0]
selected_points_indices = np.where(valid_indices)[0][selected_points_indices]
return selected_points_indices
@staticmethod
def render_pts(cam_pose, scene_path, script_path, scan_points, voxel_threshold=0.005, filter_degree=75, nO_to_nL_pose=None, require_full_scene=False):
nO_to_world_pose = DataLoadUtil.get_real_cam_O_from_cam_L(cam_pose, nO_to_nL_pose, scene_path=scene_path) nO_to_world_pose = DataLoadUtil.get_real_cam_O_from_cam_L(cam_pose, nO_to_nL_pose, scene_path=scene_path)
@@ -25,28 +84,58 @@ class RenderUtil:
params_data_path = os.path.join(temp_dir, "params.json") params_data_path = os.path.join(temp_dir, "params.json")
with open(params_data_path, 'w') as f: with open(params_data_path, 'w') as f:
json.dump(params, f) json.dump(params, f)
start_time = time.time()
result = subprocess.run([ result = subprocess.run([
'blender', '-b', '-P', script_path, '--', temp_dir '/home/hofee/blender-4.0.2-linux-x64/blender', '-b', '-P', script_path, '--', temp_dir
], capture_output=True, text=True) ], capture_output=True, text=True)
if result.returncode != 0: end_time = time.time()
print("Blender script failed:")
print(result.stderr) print(f"-- Time taken for blender: {end_time - start_time} seconds")
return None
path = os.path.join(temp_dir, "tmp") path = os.path.join(temp_dir, "tmp")
point_cloud = DataLoadUtil.get_target_point_cloud_world_from_path(path, binocular=True) cam_info = DataLoadUtil.load_cam_info(path, binocular=True)
cam_params = DataLoadUtil.load_cam_info(path, binocular=True) depth_L, depth_R = DataLoadUtil.load_depth(
path, cam_info["near_plane"],
''' TODO: old code: filter_points api is changed, need to update the code ''' cam_info["far_plane"],
filtered_point_cloud = PtsUtil.filter_points(point_cloud, model_points_normals, cam_pose=cam_params["cam_to_world"], voxel_size=voxel_threshold, theta=filter_degree) binocular=True
full_scene_point_cloud = None )
if require_full_scene: start_time = time.time()
depth_L, depth_R = DataLoadUtil.load_depth(path, cam_params['near_plane'], cam_params['far_plane'], binocular=True) mask_L, mask_R = DataLoadUtil.load_seg(path, binocular=True)
point_cloud_L = DataLoadUtil.get_point_cloud(depth_L, cam_params['cam_intrinsic'], cam_params['cam_to_world'])['points_world'] normal_L = DataLoadUtil.load_normal(path, binocular=True, left_only=True)
point_cloud_R = DataLoadUtil.get_point_cloud(depth_R, cam_params['cam_intrinsic'], cam_params['cam_to_world_R'])['points_world'] ''' target points '''
mask_img_L = mask_L
point_cloud_L = PtsUtil.random_downsample_point_cloud(point_cloud_L, 65536) mask_img_R = mask_R
point_cloud_R = PtsUtil.random_downsample_point_cloud(point_cloud_R, 65536)
full_scene_point_cloud = PtsUtil.get_overlapping_points(point_cloud_L, point_cloud_R) target_mask_img_L = (mask_L == RenderUtil.target_mask_label).all(axis=-1)
target_mask_img_R = (mask_R == RenderUtil.target_mask_label).all(axis=-1)
return filtered_point_cloud, full_scene_point_cloud sampled_target_points_L, sampled_target_normal_L = RenderUtil.get_world_points_and_normal(depth_L,target_mask_img_L,normal_L, cam_info["cam_intrinsic"], cam_info["cam_to_world"], RenderUtil.random_downsample_N)
sampled_target_points_R = RenderUtil.get_world_points(depth_R, target_mask_img_R, cam_info["cam_intrinsic"], cam_info["cam_to_world_R"], RenderUtil.random_downsample_N )
has_points = sampled_target_points_L.shape[0] > 0 and sampled_target_points_R.shape[0] > 0
if has_points:
target_points, overlap_idx = PtsUtil.get_overlapping_points(
sampled_target_points_L, sampled_target_points_R, voxel_threshold, require_idx=True
)
sampled_target_normal_L = sampled_target_normal_L[overlap_idx]
if has_points:
has_points = target_points.shape[0] > 0
if has_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=(RenderUtil.min_z, RenderUtil.max_z)
)
scan_points_indices_L = RenderUtil.get_scan_points_indices(scan_points, mask_img_L, RenderUtil.display_table_mask_label, cam_info["cam_intrinsic"], cam_info["cam_to_world"])
scan_points_indices_R = RenderUtil.get_scan_points_indices(scan_points, mask_img_R, RenderUtil.display_table_mask_label, cam_info["cam_intrinsic"], cam_info["cam_to_world_R"])
scan_points_indices = np.intersect1d(scan_points_indices_L, scan_points_indices_R)
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, scan_points_indices

192
utils/vis.py Normal file
View File

@@ -0,0 +1,192 @@
import numpy as np
import matplotlib.pyplot as plt
import sys
import os
import trimesh
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from utils.data_load import DataLoadUtil
from utils.pts import PtsUtil
class visualizeUtil:
@staticmethod
def save_all_cam_pos_and_cam_axis(root, scene, output_dir):
length = DataLoadUtil.get_scene_seq_length(root, scene)
all_cam_pos = []
all_cam_axis = []
for i in range(length):
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, "all_cam_pos.txt"), all_cam_pos)
np.savetxt(os.path.join(output_dir, "all_cam_axis.txt"), all_cam_axis)
@staticmethod
def save_all_combined_pts(root, scene, output_dir):
length = DataLoadUtil.get_scene_seq_length(root, scene)
all_combined_pts = []
for i in range(length):
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, "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(
root, model_dir, scene_name, display_table_as_world_space_origin=True
):
scene_info = DataLoadUtil.load_scene_info(root, scene_name)
target_name = scene_info["target_name"]
transformation = scene_info[target_name]
if display_table_as_world_space_origin:
location = transformation["location"] - DataLoadUtil.get_display_table_top(
root, scene_name
)
else:
location = transformation["location"]
rotation_euler = transformation["rotation_euler"]
pose_mat = trimesh.transformations.euler_matrix(*rotation_euler)
pose_mat[:3, 3] = location
mesh = DataLoadUtil.load_mesh_at(model_dir, target_name, pose_mat)
mesh_dir = os.path.join(root, scene_name, "mesh")
if not os.path.exists(mesh_dir):
os.makedirs(mesh_dir)
model_path = os.path.join(mesh_dir, "world_target_mesh.obj")
mesh.export(model_path)
@staticmethod
def save_points_and_normals(root, scene, frame_idx, output_dir, binocular=False):
target_mask_label = (0, 255, 0, 255)
path = DataLoadUtil.get_path(root, scene, frame_idx)
cam_info = DataLoadUtil.load_cam_info(path, binocular=binocular, display_table_as_world_space_origin=False)
depth = DataLoadUtil.load_depth(
path, cam_info["near_plane"],
cam_info["far_plane"],
binocular=binocular,
)
if isinstance(depth, tuple):
depth = depth[0]
mask = DataLoadUtil.load_seg(path, binocular=binocular, left_only=True)
normal = DataLoadUtil.load_normal(path, binocular=binocular, left_only=True)
''' target points '''
if mask is None:
target_mask_img = np.ones_like(depth, dtype=bool)
else:
target_mask_img = (mask == target_mask_label).all(axis=-1)
cam_intrinsic = cam_info["cam_intrinsic"]
z = depth[target_mask_img]
i, j = np.nonzero(target_mask_img)
x = (j - cam_intrinsic[0, 2]) * z / cam_intrinsic[0, 0]
y = (i - cam_intrinsic[1, 2]) * z / cam_intrinsic[1, 1]
random_downsample_N = 1000
points_camera = np.stack((x, y, z), axis=-1).reshape(-1, 3)
normal_camera = normal[target_mask_img].reshape(-1, 3)
sampled_target_points, idx = PtsUtil.random_downsample_point_cloud(
points_camera, random_downsample_N, require_idx=True
)
if len(sampled_target_points) == 0:
print("No target points")
sampled_normal_camera = normal_camera[idx]
sampled_visualized_normal = []
sampled_normal_camera[:, 2] = -sampled_normal_camera[:, 2]
sampled_normal_camera[:, 1] = -sampled_normal_camera[:, 1]
num_samples = 10
for i in range(len(sampled_target_points)):
sampled_visualized_normal.append([sampled_target_points[i] + 0.02*t * sampled_normal_camera[i] for t in range(num_samples)])
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(root, scene, frame_idx, output_dir, binocular=False):
path = DataLoadUtil.get_path(root, scene, frame_idx)
pts_world = DataLoadUtil.load_from_preprocessed_pts(path, "npy")
nrm_camera = DataLoadUtil.load_from_preprocessed_nrm(path, "npy")
cam_info = DataLoadUtil.load_cam_info(path, binocular=binocular)
cam_to_world = cam_info["cam_to_world"]
nrm_world = nrm_camera @ cam_to_world[:3, :3].T
visualized_nrm = []
num_samples = 10
for i in range(len(pts_world)):
for t in range(num_samples):
visualized_nrm.append(pts_world[i] - 0.02 * t * nrm_world[i])
visualized_nrm = np.array(visualized_nrm)
np.savetxt(os.path.join(output_dir, "nrm.txt"), visualized_nrm)
np.savetxt(os.path.join(output_dir, "pts.txt"), pts_world)
# ------ Debug ------
if __name__ == "__main__":
root = r"C:\Document\Local Project\nbv_rec\nbv_reconstruction\temp"
model_dir = r"H:\\AI\\Datasets\\scaled_object_box_meshes"
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_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)
visualizeUtil.save_pts_nrm(root, scene, "116", output_dir, binocular=True)