├── .gitignore
├── LICENSE
├── README.md
├── assert
├── ppl.png
└── vis.png
├── compile_op.sh
├── config.py
├── data_prepare_semantickitti.py
├── data_prepare_semanticposs.py
├── dataset
├── __init__.py
├── poss_testset.py
├── poss_trainset.py
├── semkitti_testset.py
├── semkitti_trainset.py
├── two_poss_trainset.py
└── two_semkitti_trainset.py
├── evaluate_DIS_SemanticKITTI.py
├── evaluate_DIS_SemanticPOSS.py
├── help_utils.py
├── my_env.yaml
├── network
├── BAAF.py
├── BAF_LAC.py
├── RandLANet.py
├── __init__.py
├── loss_func.py
└── pytorch_utils.py
├── test_SemanticKITTI.py
├── test_SemanticPOSS.py
├── tool
├── caculate_time.py
├── draw.py
├── draw_poss.py
├── draw_vis_compare.py
├── eval_KITTI_gap.py
└── main_poss.py
├── train_SemanticKITTI.py
├── train_SemanticPOSS.py
├── train_both_SemanticKITTI.py
├── train_both_SemanticPOSS.py
├── utils
├── __init__.py
├── config.py
├── cpp_wrappers
│ ├── compile_wrappers.sh
│ ├── cpp_subsampling
│ │ ├── grid_subsampling
│ │ │ ├── grid_subsampling.cpp
│ │ │ └── grid_subsampling.h
│ │ ├── setup.py
│ │ └── wrapper.cpp
│ └── cpp_utils
│ │ ├── cloud
│ │ ├── cloud.cpp
│ │ └── cloud.h
│ │ └── nanoflann
│ │ └── nanoflann.hpp
├── data_process.py
├── help_utils.py
├── metric.py
├── nearest_neighbors
│ ├── KDTreeTableAdaptor.h
│ ├── knn.cpp
│ ├── knn.pyx
│ ├── knn_.cxx
│ ├── knn_.h
│ ├── nanoflann.hpp
│ ├── setup.py
│ └── test.py
├── np_ioueval.py
├── semantic-kitti.yaml
├── semantic-poss.yaml
└── semkitti_vis
│ ├── __init__.py
│ ├── laserscan.py
│ └── laserscanvis.py
└── visualize_SemanticKITTI.py
/.gitignore:
--------------------------------------------------------------------------------
1 | # Byte-compiled / optimized / DLL files
2 | __pycache__/
3 | *.py[cod]
4 | *$py.class
5 |
6 | # C extensions
7 | *.so
8 |
9 | # Distribution / packaging
10 | .Python
11 | build/
12 | develop-eggs/
13 | dist/
14 | downloads/
15 | eggs/
16 | .eggs/
17 | lib/
18 | lib64/
19 | parts/
20 | sdist/
21 | var/
22 | wheels/
23 | pip-wheel-metadata/
24 | share/python-wheels/
25 | *.egg-info/
26 | .installed.cfg
27 | *.egg
28 | MANIFEST
29 |
30 | # PyInstaller
31 | # Usually these files are written by a python script from a template
32 | # before PyInstaller builds the exe, so as to inject date/other infos into it.
33 | *.manifest
34 | *.spec
35 |
36 | # Installer logs
37 | pip-log.txt
38 | pip-delete-this-directory.txt
39 |
40 | # Unit test / coverage reports
41 | htmlcov/
42 | .tox/
43 | .nox/
44 | .coverage
45 | .coverage.*
46 | .cache
47 | nosetests.xml
48 | coverage.xml
49 | *.cover
50 | *.py,cover
51 | .hypothesis/
52 | .pytest_cache/
53 |
54 | # Translations
55 | *.mo
56 | *.pot
57 |
58 | # Django stuff:
59 | *.log
60 | local_settings.py
61 | db.sqlite3
62 | db.sqlite3-journal
63 |
64 | # Flask stuff:
65 | instance/
66 | .webassets-cache
67 |
68 | # Scrapy stuff:
69 | .scrapy
70 |
71 | # Sphinx documentation
72 | docs/_build/
73 |
74 | # PyBuilder
75 | target/
76 |
77 | # Jupyter Notebook
78 | .ipynb_checkpoints
79 |
80 | # IPython
81 | profile_default/
82 | ipython_config.py
83 |
84 | # pyenv
85 | .python-version
86 |
87 | # pipenv
88 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
89 | # However, in case of collaboration, if having platform-specific dependencies or dependencies
90 | # having no cross-platform support, pipenv may install dependencies that don't work, or not
91 | # install all needed dependencies.
92 | #Pipfile.lock
93 |
94 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow
95 | __pypackages__/
96 |
97 | # Celery stuff
98 | celerybeat-schedule
99 | celerybeat.pid
100 |
101 | # SageMath parsed files
102 | *.sage.py
103 |
104 | # Environments
105 | .env
106 | .venv
107 | env/
108 | venv/
109 | ENV/
110 | env.bak/
111 | venv.bak/
112 |
113 | # Spyder project settings
114 | .spyderproject
115 | .spyproject
116 |
117 | # Rope project settings
118 | .ropeproject
119 |
120 | # mkdocs documentation
121 | /site
122 |
123 | # mypy
124 | .mypy_cache/
125 | .dmypy.json
126 | dmypy.json
127 |
128 | # Pyre type checker
129 | .pyre/
130 |
--------------------------------------------------------------------------------
/LICENSE:
--------------------------------------------------------------------------------
1 | MIT License
2 |
3 | Copyright (c) 2020 Tsunghan Wu
4 |
5 | Permission is hereby granted, free of charge, to any person obtaining a copy
6 | of this software and associated documentation files (the "Software"), to deal
7 | in the Software without restriction, including without limitation the rights
8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9 | copies of the Software, and to permit persons to whom the Software is
10 | furnished to do so, subject to the following conditions:
11 |
12 | The above copyright notice and this permission notice shall be included in all
13 | copies or substantial portions of the Software.
14 |
15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21 | SOFTWARE.
22 |
--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
1 |
2 |
PCB-RandNet: Rethinking Random Sampling for LiDAR Semantic Segmentation in Autonomous Driving Scene
3 |
4 | Huixian Cheng,
5 | Xianfeng Han,
6 | Hang Jiang, Dehong He, Guoqiang Xiao
7 |
8 | College of Computer and Information Science, Southwest University
9 | [arXiv]
10 |
11 |
12 | [comment]: <> ()
13 |
14 | [comment]: <> ( [arXiv])
15 |
16 | [comment]: <> (
)
17 |
18 |
19 |

20 |
21 |
22 | ## Updates
23 | - **2024-02-12[😋]** Our paper was very lucky to be accepted by ICRA 2024.
24 |
25 | ## Abstract
26 | LiDAR point cloud sparsity and distance-dependent long-tailed distributions
27 | make Random Sampling less suitable for this scenario.
28 | To alleviate this problem, we propose Polar Cylinder Balanced Random Sampling (PCB-RS)
29 | and Sampling Consistency Loss (SCL) to optimize the point cloud
30 | distribution after down-sampling to improve the segmentation performance under
31 | different ranges (especially long range).
32 |
33 |
34 |

35 |
36 |
37 | ## Environment Setup
38 | Install python packages
39 | ```
40 | conda env create -f my_env.yaml
41 | conda activate randla
42 | ```
43 | Note:
44 | - Not all packages are necessary, you can refer to this environment to install.
45 | - Since we are using operator `torch.take_along_dim` that are only available in pytorch version ≥ 1.9.0 in [here](https://github.com/huixiancheng/PCB-RandNet/blob/3d05d457761f5f096263b9658940bc7d01ce04eb/train_both_SemanticKITTI.py#L170), please make sure that the installed version meets the requirements.
46 |
47 | Compile C++ Wrappers
48 | ```
49 | bash compile_op.sh
50 | ```
51 |
52 | ## Prepare Data
53 | Download SemanticKITTI from [official web](http://www.semantic-kitti.org/dataset.html). Download SemanticPOSS from [official web](http://www.poss.pku.edu.cn./download.html).
54 | Then preprocess the data:
55 |
56 | ```
57 | python data_prepare_semantickitti.py
58 | python data_prepare_semanticposs.py
59 | ```
60 | Note:
61 | - Please change the data path with your own path.
62 | - Change corresponding path `self.dataset_path` in all `.py` in dataset (such as [here](https://github.com/huixiancheng/PCB-RandNet/blob/3d05d457761f5f096263b9658940bc7d01ce04eb/dataset/semkitti_trainset.py#L32)).
63 |
64 | ## Training
65 |
66 | 1. Training baseline with RS or PCB-RS
67 |
68 | ```
69 | python train_SemanticKITTI.py
70 |
71 | python train_SemanticPOSS.py
72 |
73 | Options:
74 | --backbone select the backbone to be used: choices=['randla', 'baflac', 'baaf']
75 | --checkpoint_path path to pretrained models(if any), otherwise train from start
76 | --log_dir Name of the log dir
77 | --max_epoch max epoch for the model to run
78 | --batch_size training batch size, modify to full utilize the GPU/s
79 | --val_batch_size batch size for validation
80 | --num_workers number of workers for I/O
81 | --sampling select the sampling way: RS or PCB-RS. choices=['random', 'polar']
82 | --seed set random seed
83 | --step set length of dataset: 0 mean use all data || 4 mean use 1/4 dataset (Only use for SemanticKITTI)
84 | --grid resolution of polar cylinder
85 | ```
86 |
87 | 2. Training model with PCB-RS and SCL
88 | ```
89 | python train_both_SemanticKITTI.py
90 |
91 | python train_both_SemanticPOSS.py
92 |
93 | Options: Similar to before.
94 | ```
95 |
96 | ## Test
97 |
98 | ```
99 | python test_SemanticKITTI.py
100 |
101 | python test_SemanticPOSS.py
102 |
103 | Options:
104 | --infer_type all: infer all points in specified sequence, sub: subsamples in specified sequence
105 | --sampling select the sampling way: RS or PCB-RS. choices=['random', 'polar']
106 | --backbone select the backbone to be used: choices=['randla', 'baflac', 'baaf']
107 | --checkpoint_path required. path to the model to test
108 | --test_id sequence id to test
109 | --result_dir result dir of predictions
110 | --step set length of dataset: 0 mean use all data || 4 mean use 1/4 dataset (Only use for SemanticKITTI)
111 | --grid resolution of polar cylinder
112 | ```
113 | **Note:**
114 | - **For SemanticKITTI dataset, if your want to infer all data of 08 sequence, please make sure you have more than 32G of free RAM. If your hardware does not meet this requirement, please refer to this [issue](https://github.com/tsunghan-wu/RandLA-Net-pytorch/issues/10).**
115 |
116 | - **Due to the use of `torch_data.IterableDataset`, the num_workers needs to be set to 0 and cannot use multi-threaded acceleration, so the test speed will be very slow. Welcome to propose a feasible accelerated PR for this. (This [reference](https://medium.com/speechmatics/how-to-build-a-streaming-dataloader-with-pytorch-a66dd891d9dd) may be useful)**
117 |
118 | ## Other Utils
119 | 1. Evaluation with Distances
120 | ```
121 | python evaluate_DIS_SemanticKITTI.py || python evaluate_DIS_SemanticPOSS.py
122 | ```
123 | 2. Visualization
124 | ```
125 | python visualize_SemanticKITTI.py
126 | ```
127 | This code is not used, so it has not been tested and may have bugs. If you want to use this code to visualize SemanticKITTI and SemanticPOSS, please refer to [this repo](https://github.com/PRBonn/semantic-kitti-api) and [this repo](https://github.com/huixiancheng/CENet).
128 |
129 | 3. Others in tool.
130 | ```
131 | caculate_time.py Use to get Time Consumption Comparison between RS, PCB-RS, and FPS.
132 | draw.py / draw_poss.py Use to draw performance plot at different distances.
133 | draw_vis_compare.py Use to generate qualitative visualization and quantitative statistical analysis images between RS and PCB-RS similar to the paper.
134 | eval_KITTI_gap.py Use to calculate the difference in performance of the model under different sampling methods.
135 | main_poss.py Use to count distance distribution
136 | ```
137 |
138 | ## Others to Note and Clarify
139 |
140 | - All our experiments were trained on a single NVIDIA RTX 3090 GPU using mixed precision.
141 | - We set `torch.backends.cudnn.enabled = False` and use `with torch.cuda.amp.autocast():` to reduce training time and save GPU memory. However, it is not certain whether it will have an impact on performance.
142 | - In `train_SemanticKITTI.py`, We use `num_classes = 20` instead of `num_classes = 19`. The reason is that I think the [ops](https://github.com/huixiancheng/PCB-RandNet/blob/main/network/loss_func.py) used to calculate the loss in the original code is very complex, and its main point is to ignore the loss when `label == 0`. Therefore, I just changed the num_classes to 20 and set the weight of `label == 0` to 0 to achieve it.
143 | - In fact, the criterion corresponding to `compute_loss` should be `self.criterion = nn.CrossEntropyLoss(weight=class_weights, reduction='none')` instead of `self.criterion = nn.CrossEntropyLoss(weight=class_weights)`. This may affect the scale of the computed loss, but should not affect performance.
144 | - There are two ways to calculate SCL, one is the [L1 loss](https://github.com/huixiancheng/PCB-RandNet/blob/05f2d4c796ac68d45745ddea2b7d0119c43cc0a5/train_both_SemanticKITTI.py#L172) used by default, and the other is method [here](https://github.com/huixiancheng/PCB-RandNet/blob/05f2d4c796ac68d45745ddea2b7d0119c43cc0a5/train_both_SemanticKITTI.py#L127-L135) from CPGNET.
145 | - In terms of loss size, the default L1 loss is better used with `reduction='mean'` and `consistency_loss_l1` is better used with `reduction='none'`. Some of our experimental results show that the performance of the models trained in both settings is comparable.
146 | - Since the two losses used L_ce and L_scl may be correlated. Therefore the uncertainty weighting method may not converge using random initialization. In our experiments, we set `ce_sigma==0.27` and `l1_sigma==0.62` because they are close to our initial experimental values with `seed==1024`. Some of our experiments show that the model converges effectively with `ce_sigma in range [0.2, 0.3]` and `l1_sigma in range [0.6, 0.8]`.
147 | - For BAF-LAC backbone, our implemented model is as consistent as possible with the [original Repo](https://github.com/Xiangxu-0103/BAF-LAC).
148 | - For BAAF backbone, our implementation is slightly different from the [original Repo](https://github.com/ShiQiu0419/BAAF-Net), which will result in some performance differences. While the original Repo used FPS sampling, I changed to RS sampling to keep the framework structure consistent. The original Repo used an auxiliary loss, in order to keep simple, we did not use this part.
149 | - Since I' m not good at coding & optimization, the code of [PCB-RS](https://github.com/huixiancheng/PCB-RandNet/blob/05f2d4c796ac68d45745ddea2b7d0119c43cc0a5/dataset/semkitti_trainset.py#L130-L230) is very rough, if you are interesting in optimizing & accelerating this part, welcome to PR !
150 |
151 | ## Pretrained Models and Logs:
152 | | **KITTI Result** | **POSS Result** | **Ablation Study** |
153 | | ---------------- | --------------- | ------------------ |
154 | | [Google Drive](https://drive.google.com/drive/folders/1cWHgWbXgPN-33y1k1CIIH2vxyt_NdW3P?usp=sharing) | [Google Drive](https://drive.google.com/drive/folders/13-RdlHdVM7-kVaH3hS8eAaTL7F139itF?usp=sharing) | [Google Drive](https://drive.google.com/drive/folders/1rVvV7lu7luEM2gUzTyXaaqCeXi7rqffy?usp=sharing)
155 |
156 | ## Acknowledgement
157 | - This repo is heavily based on [RandLA-Net-pytorch](https://github.com/tsunghan-wu/RandLA-Net-pytorch), many thanks to the author for the excellent Pytorch reproduction. :thumbsup:
158 | - If you are looking for a better performance reproduction of RandLA-Net, please refer to this [Repo](https://github.com/charlesChen02/RandLA-pytorch).
159 | - Our code refers in varying degrees to the Repos listed below: [RandLA-Net](https://github.com/QingyongHu/RandLA-Net),
160 | [BAAF-Net](https://github.com/ShiQiu0419/BAAF-Net),
161 | [BAF-LAC](https://github.com/Xiangxu-0103/BAF-LAC), [PolarNet](https://github.com/edwardzhou130/PolarSeg),
162 | [Cylinder3D](https://github.com/xinge008/Cylinder3D), [R-Drop](https://github.com/dropreg/R-Drop),
163 | [JS3C-Net](https://github.com/yanx27/JS3C-Net), [Pointnet_Pointnet2_pytorch](https://github.com/yanx27/Pointnet_Pointnet2_pytorch),
164 | [AutomaticWeightedLoss](https://github.com/Mikoto10032/AutomaticWeightedLoss), [MINet](https://github.com/sj-li/MINet). Thanks a lot for their open source code. :heart:
165 | - Thanks to PhD. Qiu & PhD. Li for the code help on [BAAF-Net](https://github.com/ShiQiu0419/BAAF-Net) and [CPGNet](https://arxiv.org/abs/2204.09914) in the email. :heart:
166 | - Thanks for the answers provided in this stackoverflow [question](https://stackoverflow.com/questions/42232540/how-to-find-indices-of-a-reordered-numpy-array).
167 | - The paper was inspired by a helpful discussion with RandLA-Net authors at China3DV in Xiamen in June 2021, thanks the insightful comments and feedback from [PhD. Hu](https://github.com/QingyongHu). :heart:
168 |
169 |
170 | ## Citation
171 | If you find this work helpful, please consider citing our paper:
172 | ~~~
173 | @article{cheng2022pcb,
174 | title={PCB-RandNet: Rethinking Random Sampling for LIDAR Semantic Segmentation in Autonomous Driving Scene},
175 | author={Cheng, Huixian and Han, XianFeng and Jiang, Hang and He, Dehong and Xiao, Guoqiang},
176 | journal={arXiv preprint arXiv:2209.13797},
177 | year={2022}
178 | }
179 | ~~~
180 |
--------------------------------------------------------------------------------
/assert/ppl.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/huixiancheng/PCB-RandNet/304c84b67380a23ab66516581b348c7db2f9f0df/assert/ppl.png
--------------------------------------------------------------------------------
/assert/vis.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/huixiancheng/PCB-RandNet/304c84b67380a23ab66516581b348c7db2f9f0df/assert/vis.png
--------------------------------------------------------------------------------
/compile_op.sh:
--------------------------------------------------------------------------------
1 | cd utils/nearest_neighbors
2 | python setup.py install --home="."
3 | cd ../../
4 |
5 | cd utils/cpp_wrappers
6 | sh compile_wrappers.sh
7 | cd ../../../
--------------------------------------------------------------------------------
/config.py:
--------------------------------------------------------------------------------
1 | class ConfigSemanticKITTI:
2 | k_n = 16 # KNN
3 | num_layers = 4 # Number of layers
4 | num_points = 4096 * 11 # Number of input points
5 | num_classes = 20 # Number of valid classes
6 | sub_grid_size = 0.06 # preprocess_parameter
7 |
8 | sub_sampling_ratio = [4, 4, 4, 4] # sampling ratio of random sampling at each layer
9 | d_out = [16, 64, 128, 256] # feature dimension
10 | num_sub_points = [num_points // 4, num_points // 16, num_points // 64, num_points // 256]
11 |
12 | class ConfigSemanticKITTI_BAF:
13 | k_n = 16 # KNN
14 | num_layers = 4 # Number of layers
15 | num_points = 4096 * 11 # Number of input points
16 | num_classes = 20 # Number of valid classes
17 | sub_grid_size = 0.06 # preprocess_parameter
18 |
19 | sub_sampling_ratio = [4, 4, 4, 4] # sampling ratio of random sampling at each layer
20 | d_out = [16, 32, 64, 128] # feature dimension
21 | num_sub_points = [num_points // 4, num_points // 16, num_points // 64, num_points // 256]
22 |
23 | class ConfigSemanticPOSS:
24 | k_n = 16 # KNN
25 | num_layers = 4 # Number of layers
26 | num_points = 4096 * 11 # Number of input points
27 | num_classes = 11 # Number of valid classes
28 | sub_grid_size = 0.06 # preprocess_parameter
29 |
30 | sub_sampling_ratio = [4, 4, 4, 4] # sampling ratio of random sampling at each layer
31 | d_out = [16, 64, 128, 256] # feature dimension
32 | num_sub_points = [num_points // 4, num_points // 16, num_points // 64, num_points // 256]
33 |
34 |
35 | class ConfigSemanticPOSS_BAF:
36 | k_n = 16 # KNN
37 | num_layers = 4 # Number of layers
38 | num_points = 4096 * 11 # Number of input points
39 | num_classes = 11 # Number of valid classes
40 | sub_grid_size = 0.06 # preprocess_parameter
41 |
42 | sub_sampling_ratio = [4, 4, 4, 4] # sampling ratio of random sampling at each layer
43 | d_out = [16, 32, 64, 128] # feature dimension
44 | num_sub_points = [num_points // 4, num_points // 16, num_points // 64, num_points // 256]
45 |
--------------------------------------------------------------------------------
/data_prepare_semantickitti.py:
--------------------------------------------------------------------------------
1 | import os
2 | import yaml
3 | import pickle
4 | import argparse
5 | import numpy as np
6 | from os.path import join, exists
7 | from sklearn.neighbors import KDTree
8 | from utils.data_process import DataProcessing as DP
9 | from tqdm import tqdm
10 | parser = argparse.ArgumentParser()
11 | parser.add_argument('--src_path', default='/home/chx/Work/semantic-kitti/sequences', help='source dataset path [default: None]')
12 | parser.add_argument('--dst_path', default='/home/chx/Work/semantic-kitti-randla/sequences', help='destination dataset path [default: None]')
13 | parser.add_argument('--grid_size', type=float, default=0.06, help='Subsample Grid Size [default: 0.06]')
14 | parser.add_argument('--yaml_config', default='utils/semantic-kitti.yaml', help='semantic-kitti.yaml path')
15 | FLAGS = parser.parse_args()
16 |
17 |
18 | data_config = FLAGS.yaml_config
19 | DATA = yaml.safe_load(open(data_config, 'r'))
20 | remap_dict = DATA["learning_map"]
21 | max_key = max(remap_dict.keys())
22 | remap_lut = np.zeros((max_key + 100), dtype=np.int32)
23 | remap_lut[list(remap_dict.keys())] = list(remap_dict.values())
24 |
25 | grid_size = FLAGS.grid_size
26 | dataset_path = FLAGS.src_path
27 | output_path = FLAGS.dst_path
28 | seq_list = np.sort(os.listdir(dataset_path))
29 |
30 |
31 | for seq_id in tqdm(seq_list, total=len(seq_list)):
32 | print('sequence ' + seq_id + ' start')
33 | seq_path = join(dataset_path, seq_id)
34 | seq_path_out = join(output_path, seq_id)
35 | pc_path = join(seq_path, 'velodyne')
36 | pc_path_out = join(seq_path_out, 'velodyne')
37 | KDTree_path_out = join(seq_path_out, 'KDTree')
38 | os.makedirs(seq_path_out) if not exists(seq_path_out) else None
39 | os.makedirs(pc_path_out) if not exists(pc_path_out) else None
40 | os.makedirs(KDTree_path_out) if not exists(KDTree_path_out) else None
41 |
42 | if int(seq_id) < 11:
43 | # continue
44 | label_path = join(seq_path, 'labels')
45 | label_path_out = join(seq_path_out, 'labels')
46 | os.makedirs(label_path_out) if not exists(label_path_out) else None
47 | scan_list = np.sort(os.listdir(pc_path))
48 |
49 | # scan_list = scan_list[0:len(scan_list):4]
50 |
51 | for scan_id in tqdm(scan_list, total=len(scan_list)):
52 | # print(scan_id)
53 | points = DP.load_pc_kitti(join(pc_path, scan_id))
54 | labels = DP.load_label_kitti(join(label_path, str(scan_id[:-4]) + '.label'), remap_lut)
55 | sub_points, sub_labels = DP.grid_sub_sampling(points, labels=labels, grid_size=grid_size)
56 | search_tree = KDTree(sub_points)
57 | KDTree_save = join(KDTree_path_out, str(scan_id[:-4]) + '.pkl')
58 | np.save(join(pc_path_out, scan_id)[:-4], sub_points)
59 | np.save(join(label_path_out, scan_id)[:-4], sub_labels)
60 | with open(KDTree_save, 'wb') as f:
61 | pickle.dump(search_tree, f)
62 | if seq_id == '08':
63 | proj_path = join(seq_path_out, 'proj')
64 | os.makedirs(proj_path) if not exists(proj_path) else None
65 | proj_inds = np.squeeze(search_tree.query(points, return_distance=False))
66 | proj_inds = proj_inds.astype(np.int32)
67 | proj_save = join(proj_path, str(scan_id[:-4]) + '_proj.pkl')
68 | with open(proj_save, 'wb') as f:
69 | pickle.dump([proj_inds], f)
70 | else:
71 | proj_path = join(seq_path_out, 'proj')
72 | os.makedirs(proj_path) if not exists(proj_path) else None
73 | scan_list = np.sort(os.listdir(pc_path))
74 | for scan_id in tqdm(scan_list, total=len(scan_list)):
75 | # print(scan_id)
76 | points = DP.load_pc_kitti(join(pc_path, scan_id))
77 | sub_points = DP.grid_sub_sampling(points, grid_size=0.06)
78 | search_tree = KDTree(sub_points)
79 | proj_inds = np.squeeze(search_tree.query(points, return_distance=False))
80 | proj_inds = proj_inds.astype(np.int32)
81 | KDTree_save = join(KDTree_path_out, str(scan_id[:-4]) + '.pkl')
82 | proj_save = join(proj_path, str(scan_id[:-4]) + '_proj.pkl')
83 | np.save(join(pc_path_out, scan_id)[:-4], sub_points)
84 | with open(KDTree_save, 'wb') as f:
85 | pickle.dump(search_tree, f)
86 | with open(proj_save, 'wb') as f:
87 | pickle.dump([proj_inds], f)
88 |
--------------------------------------------------------------------------------
/data_prepare_semanticposs.py:
--------------------------------------------------------------------------------
1 | import os
2 | import yaml
3 | import pickle
4 | import argparse
5 | import numpy as np
6 | from os.path import join, exists
7 | from sklearn.neighbors import KDTree
8 | from utils.data_process import DataProcessing as DP
9 | from tqdm import tqdm
10 | parser = argparse.ArgumentParser()
11 | parser.add_argument('--src_path', default='/home/chx/Work/SemanticPOSS_dataset/sequences', help='source dataset path [default: None]')
12 | parser.add_argument('--dst_path', default='/home/chx/Work/SemanticPOSS_dataset_grid_0.01/sequences', help='destination dataset path [default: None]')
13 | parser.add_argument('--grid_size', type=float, default=0.01, help='Subsample Grid Size [default: 0.06]')
14 | parser.add_argument('--yaml_config', default='utils/semantic-poss.yaml', help='semantic-kitti.yaml path')
15 | FLAGS = parser.parse_args()
16 |
17 |
18 | data_config = FLAGS.yaml_config
19 | DATA = yaml.safe_load(open(data_config, 'r'))
20 | remap_dict = DATA["learning_map"]
21 | max_key = max(remap_dict.keys())
22 | remap_lut = np.zeros((max_key + 100), dtype=np.int32)
23 | remap_lut[list(remap_dict.keys())] = list(remap_dict.values())
24 | grid_size = FLAGS.grid_size
25 | dataset_path = FLAGS.src_path
26 | output_path = FLAGS.dst_path
27 | seq_list = np.sort(os.listdir(dataset_path))
28 |
29 | for seq_id in tqdm(seq_list, total=len(seq_list)):
30 | # seq_id = '03'
31 | print('sequence ' + seq_id + ' start')
32 | seq_path = join(dataset_path, seq_id)
33 | seq_path_out = join(output_path, seq_id)
34 | pc_path = join(seq_path, 'velodyne')
35 | pc_path_out = join(seq_path_out, 'velodyne')
36 | KDTree_path_out = join(seq_path_out, 'KDTree')
37 | os.makedirs(seq_path_out) if not exists(seq_path_out) else None
38 | os.makedirs(pc_path_out) if not exists(pc_path_out) else None
39 | os.makedirs(KDTree_path_out) if not exists(KDTree_path_out) else None
40 |
41 | label_path = join(seq_path, 'labels')
42 | label_path_out = join(seq_path_out, 'labels')
43 | os.makedirs(label_path_out) if not exists(label_path_out) else None
44 | scan_list = np.sort(os.listdir(pc_path))
45 | for scan_id in tqdm(scan_list, total=len(scan_list)):
46 | # print(scan_id)
47 | points = DP.load_pc_kitti(join(pc_path, scan_id))
48 | labels = DP.load_label_kitti(join(label_path, str(scan_id[:-4]) + '.label'), remap_lut)
49 | sub_points, sub_labels = DP.grid_sub_sampling(points, labels=labels, grid_size=grid_size)
50 | search_tree = KDTree(sub_points)
51 | KDTree_save = join(KDTree_path_out, str(scan_id[:-4]) + '.pkl')
52 | np.save(join(pc_path_out, scan_id)[:-4], sub_points)
53 | np.save(join(label_path_out, scan_id)[:-4], sub_labels)
54 | with open(KDTree_save, 'wb') as f:
55 | pickle.dump(search_tree, f)
56 | if seq_id == '03':
57 | proj_path = join(seq_path_out, 'proj')
58 | os.makedirs(proj_path) if not exists(proj_path) else None
59 | proj_inds = np.squeeze(search_tree.query(points, return_distance=False))
60 | proj_inds = proj_inds.astype(np.int32)
61 | proj_save = join(proj_path, str(scan_id[:-4]) + '_proj.pkl')
62 | with open(proj_save, 'wb') as f:
63 | pickle.dump([proj_inds], f)
64 |
65 |
--------------------------------------------------------------------------------
/dataset/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/huixiancheng/PCB-RandNet/304c84b67380a23ab66516581b348c7db2f9f0df/dataset/__init__.py
--------------------------------------------------------------------------------
/evaluate_DIS_SemanticKITTI.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # This file is covered by the LICENSE file in the root of this project.
3 |
4 | import os
5 | import sys
6 | import yaml
7 | import argparse
8 | import numpy as np
9 | from tqdm import tqdm
10 | from utils.np_ioueval import iouEval
11 |
12 |
13 | DISTANCES = [(1e-8, 100.0),
14 | (1e-8, 10.0),
15 | (10.0, 20.0),
16 | (20.0, 30.0),
17 | (30.0, 40.0),
18 | (40.0, 50.0),
19 | (50.0, 100.0)]
20 |
21 | def get_args():
22 | parser = argparse.ArgumentParser("./evaluate_semantics.py")
23 | parser.add_argument(
24 | '--eval_type', '-e',
25 | default='all',
26 | type=str, choices=['all', 'sub'],
27 | help='Eval ALL or just eval Subsample')
28 | parser.add_argument(
29 | '--dataset', '-d',
30 | type=str,
31 | default='/home/chx/Work/semantic-kitti/sequences',
32 | help='Dataset dir. No Default',
33 | )
34 | parser.add_argument(
35 | '--predictions', '-p',
36 | type=str,
37 | default='/home/chx/Work/PT-RandLA-Net/96-96-32-1024-polar',
38 | help='Prediction dir. Same organization as dataset, but predictions in'
39 | 'each sequences "prediction" directory. No Default. If no option is set'
40 | ' we look for the labels in the same directory as dataset'
41 | )
42 | parser.add_argument(
43 | '--sequences', '-s',
44 | nargs="+",
45 | default=["08"],
46 | help='evaluated sequences',
47 | )
48 | parser.add_argument(
49 | '--datacfg', '-dc',
50 | type=str,
51 | required=False,
52 | default="utils/semantic-kitti.yaml",
53 | help='Dataset config file. Defaults to %(default)s',
54 | )
55 | parser.add_argument(
56 | '--limit', '-l',
57 | type=int,
58 | required=False,
59 | default=None,
60 | help='Limit to the first "--limit" points of each scan. Useful for'
61 | ' evaluating single scan from aggregated pointcloud.'
62 | ' Defaults to %(default)s',
63 | )
64 | FLAGS = parser.parse_args()
65 |
66 | # fill in real predictions dir
67 | if FLAGS.predictions is None:
68 | FLAGS.predictions = FLAGS.dataset
69 | return FLAGS
70 |
71 |
72 | def load_label(data_root, sequences, sub_dir_name, ext):
73 | label_names = []
74 | for sequence in sequences:
75 | sequence = '{0:02d}'.format(int(sequence))
76 | label_paths = os.path.join(data_root, str(sequence), sub_dir_name)
77 | # populate the label names
78 | seq_label_names = [os.path.join(dp, f) for dp, dn, fn in os.walk(
79 | os.path.expanduser(label_paths)) for f in fn if f".{ext}" in f]
80 | seq_label_names.sort()
81 | label_names.extend(seq_label_names)
82 | return label_names
83 |
84 |
85 | if __name__ == '__main__':
86 |
87 | FLAGS = get_args()
88 | # print summary of what we will do
89 | print("*" * 80)
90 | print("INTERFACE:")
91 | print("Eval What:", FLAGS.eval_type)
92 | print("Data: ", FLAGS.dataset)
93 | print("Predictions: ", FLAGS.predictions)
94 | print("Sequences: ", FLAGS.sequences)
95 | print("Config: ", FLAGS.datacfg)
96 | print("Limit: ", FLAGS.limit)
97 | print("*" * 80)
98 |
99 | print("Opening data config file %s" % FLAGS.datacfg)
100 | DATA = yaml.safe_load(open(FLAGS.datacfg, 'r'))
101 |
102 | # get number of interest classes, and the label mappings
103 | class_strings = DATA["labels"]
104 | class_ignore = DATA["learning_ignore"]
105 | class_inv_remap = DATA["learning_map_inv"]
106 | nr_classes = len(class_inv_remap)
107 | data_config = FLAGS.datacfg
108 | DATA = yaml.safe_load(open(data_config, 'r'))
109 | remap_dict = DATA["learning_map"]
110 | max_key = max(remap_dict.keys())
111 | remap_lut = np.zeros((max_key + 100), dtype=np.int32)
112 | remap_lut[list(remap_dict.keys())] = list(remap_dict.values())
113 |
114 | # create evaluator
115 | ignore = []
116 | for cl, ign in class_ignore.items():
117 | if ign:
118 | x_cl = int(cl)
119 | ignore.append(x_cl)
120 | print("Ignoring xentropy class ", x_cl, " in IoU evaluation")
121 | # create evaluator
122 |
123 | # create evaluator
124 | evaluators = []
125 | for i in range(len(DISTANCES)):
126 | evaluators.append(iouEval(nr_classes, ignore))
127 | evaluators[i].reset()
128 |
129 |
130 | # get label paths
131 | if FLAGS.eval_type == "sub":
132 | label_names = load_label(FLAGS.dataset, FLAGS.sequences, "labels", "npy")
133 | else:
134 | label_names = load_label(FLAGS.dataset, FLAGS.sequences, "labels", "label")
135 |
136 | py_names = load_label(FLAGS.dataset, FLAGS.sequences, "velodyne", "bin")
137 | # py_names = py_names[0:len(py_names):4]
138 |
139 | # get predictions paths
140 | pred_names = load_label(FLAGS.predictions, FLAGS.sequences, "predictions", "label")
141 | # label_names = label_names[0:len(label_names):4]
142 |
143 | print(len(label_names), len(pred_names))
144 | assert(len(label_names) == len(pred_names))
145 |
146 |
147 | print("Evaluating sequences")
148 | N = len(label_names)
149 | # open each file, get the tensor, and make the iou comparison
150 | for i in tqdm(range(N)):
151 | label_file = label_names[i]
152 | pred_file = pred_names[i]
153 | points_file = py_names[i]
154 | # open label
155 | if FLAGS.eval_type == "sub":
156 | label = np.load(label_file)
157 | label = label.reshape((-1)) # reshape to vector
158 | else:
159 | label = np.fromfile(label_file, dtype=np.int32)
160 | label = label.reshape((-1))
161 | sem_label = label & 0xFFFF # semantic label in lower half
162 | inst_label = label >> 16 # instance id in upper half
163 | assert ((sem_label + (inst_label << 16) == label).all())
164 | label = remap_lut[sem_label]
165 |
166 | if FLAGS.limit is not None:
167 | label = label[:FLAGS.limit] # limit to desired length
168 | # open prediction
169 |
170 | pred = np.fromfile(pred_file, dtype=np.int32)
171 | pred = pred.reshape((-1)) # reshape to vector
172 | pred = remap_lut[pred]
173 | if FLAGS.limit is not None:
174 | pred = pred[:FLAGS.limit] # limit to desired length
175 | # add single scan to evaluation
176 |
177 | xyzr = np.fromfile(points_file, dtype=np.float32)
178 | xyzr = xyzr.reshape((-1, 4))
179 | points = xyzr[:, 0:3]
180 | if FLAGS.limit is not None:
181 | points = points[:FLAGS.limit] # limit to desired length
182 | depth = np.linalg.norm(points, 2, axis=1)
183 |
184 | # evaluate for all distances
185 | for idx in range(len(DISTANCES)):
186 | # select by range
187 | lrange = DISTANCES[idx][0]
188 | hrange = DISTANCES[idx][1]
189 | mask = np.logical_and(depth > lrange, depth <= hrange)
190 |
191 | # mask by distance
192 | # mask_depth = depth[mask]
193 | # print("mask range, ", mask_depth.max(), mask_depth.min())
194 | mask_label = label[mask]
195 | mask_pred = pred[mask]
196 |
197 | # add single scan to evaluation
198 | evaluators[idx].addBatch(mask_pred, mask_label)
199 |
200 | # print for all ranges
201 | print("*" * 80)
202 | for idx in range(len(DISTANCES)):
203 | # when I am done, print the evaluation
204 | m_accuracy = evaluators[idx].getacc()
205 | m_jaccard, class_jaccard = evaluators[idx].getIoU()
206 |
207 | # print for spreadsheet
208 | sys.stdout.write('range {lrange}m to {hrange}m,'.format(lrange=DISTANCES[idx][0],
209 | hrange=DISTANCES[idx][1]))
210 | for i, jacc in enumerate(class_jaccard):
211 | if i not in ignore:
212 | sys.stdout.write('{jacc:.3f}'.format(jacc=jacc.item()))
213 | sys.stdout.write(",")
214 | sys.stdout.write('{jacc:.3f}'.format(jacc=m_jaccard.item()))
215 | sys.stdout.write(",")
216 | sys.stdout.write('{acc:.3f}'.format(acc=m_accuracy.item()))
217 | sys.stdout.write('\n')
218 | sys.stdout.flush()
219 |
220 | # if FLAGS.codalab:
221 | # results = {}
222 | # for idx in range(len(DISTANCES)):
223 | # # make string for distance
224 | # d_str = str(DISTANCES[idx][-1])+"m_"
225 | #
226 | # # get values for this distance range
227 | # m_accuracy = evaluators[idx].getacc()
228 | # m_jaccard, class_jaccard = evaluators[idx].getIoU()
229 | #
230 | # # put in dictionary
231 | # results[d_str+"accuracy_mean"] = float(m_accuracy)
232 | # results[d_str+"iou_mean"] = float(m_jaccard)
233 | # for i, jacc in enumerate(class_jaccard):
234 | # if i not in ignore:
235 | # results[d_str+"iou_"+class_strings[class_inv_remap[i]]] = float(jacc)
236 | # # save to file
237 | # with open('segmentation_scores_distance.txt', 'w') as yaml_file:
238 | # yaml.dump(results, yaml_file, default_flow_style=False)
--------------------------------------------------------------------------------
/evaluate_DIS_SemanticPOSS.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # This file is covered by the LICENSE file in the root of this project.
3 |
4 | import os
5 | import sys
6 | import yaml
7 | import argparse
8 | import numpy as np
9 | from tqdm import tqdm
10 | from utils.np_ioueval import iouEval
11 |
12 |
13 | DISTANCES = [(1e-8, 300.0),
14 | (1e-8, 10.0),
15 | (10.0, 20.0),
16 | (20.0, 30.0),
17 | (30.0, 40.0),
18 | (40.0, 50.0),
19 | (50.0, 60.0),
20 | (60.0, 70.0),
21 | (70.0, 80.0),
22 | (80.0, 300.0),]
23 |
24 | def get_args():
25 | parser = argparse.ArgumentParser("./evaluate_semantics.py")
26 | parser.add_argument(
27 | '--eval_type', '-e',
28 | default='all',
29 | type=str, choices=['all', 'sub'],
30 | help='Eval ALL or just eval Subsample')
31 | parser.add_argument(
32 | '--dataset', '-d',
33 | type=str,
34 | default='/home/chx/Work/SemanticPOSS_dataset/sequences',
35 | help='Dataset dir. No Default',
36 | )
37 | parser.add_argument(
38 | '--predictions', '-p',
39 | type=str,
40 | default='/home/chx/Work/PT-RandLA-Net/poss_result/test-baaf-random',
41 | help='Prediction dir. Same organization as dataset, but predictions in'
42 | 'each sequences "prediction" directory. No Default. If no option is set'
43 | ' we look for the labels in the same directory as dataset'
44 | )
45 | parser.add_argument(
46 | '--sequences', '-s',
47 | nargs="+",
48 | default=["03"],
49 | help='evaluated sequences',
50 | )
51 |
52 | parser.add_argument(
53 | '--datacfg', '-dc',
54 | type=str,
55 | required=False,
56 | default="utils/semantic-poss.yaml",
57 | help='Dataset config file. Defaults to %(default)s',
58 | )
59 | parser.add_argument(
60 | '--limit', '-l',
61 | type=int,
62 | required=False,
63 | default=None,
64 | help='Limit to the first "--limit" points of each scan. Useful for'
65 | ' evaluating single scan from aggregated pointcloud.'
66 | ' Defaults to %(default)s',
67 | )
68 | FLAGS = parser.parse_args()
69 |
70 | # fill in real predictions dir
71 | if FLAGS.predictions is None:
72 | FLAGS.predictions = FLAGS.dataset
73 | return FLAGS
74 |
75 |
76 | def load_label(data_root, sequences, sub_dir_name, ext):
77 | label_names = []
78 | for sequence in sequences:
79 | sequence = '{0:02d}'.format(int(sequence))
80 | label_paths = os.path.join(data_root, str(sequence), sub_dir_name)
81 | # populate the label names
82 | seq_label_names = [os.path.join(dp, f) for dp, dn, fn in os.walk(
83 | os.path.expanduser(label_paths)) for f in fn if f".{ext}" in f]
84 | seq_label_names.sort()
85 | label_names.extend(seq_label_names)
86 | return label_names
87 |
88 |
89 | if __name__ == '__main__':
90 |
91 | FLAGS = get_args()
92 | # print summary of what we will do
93 | print("*" * 80)
94 | print("INTERFACE:")
95 | print("Eval What:", FLAGS.eval_type)
96 | print("Data: ", FLAGS.dataset)
97 | print("Predictions: ", FLAGS.predictions)
98 | print("Sequences: ", FLAGS.sequences)
99 | print("Config: ", FLAGS.datacfg)
100 | print("Limit: ", FLAGS.limit)
101 | print("*" * 80)
102 |
103 | print("Opening data config file %s" % FLAGS.datacfg)
104 | DATA = yaml.safe_load(open(FLAGS.datacfg, 'r'))
105 |
106 | # get number of interest classes, and the label mappings
107 | class_strings = DATA["labels"]
108 | class_ignore = DATA["learning_ignore"]
109 | class_inv_remap = DATA["learning_map_inv"]
110 | nr_classes = len(class_inv_remap)
111 | data_config = FLAGS.datacfg
112 | DATA = yaml.safe_load(open(data_config, 'r'))
113 | remap_dict = DATA["learning_map"]
114 | max_key = max(remap_dict.keys())
115 | remap_lut = np.zeros((max_key + 100), dtype=np.int32)
116 | remap_lut[list(remap_dict.keys())] = list(remap_dict.values())
117 |
118 | # create evaluator
119 | ignore = []
120 | for cl, ign in class_ignore.items():
121 | if ign:
122 | x_cl = int(cl)
123 | ignore.append(x_cl)
124 | print("Ignoring xentropy class ", x_cl, " in IoU evaluation")
125 | # create evaluator
126 |
127 | # create evaluator
128 | evaluators = []
129 | for i in range(len(DISTANCES)):
130 | evaluators.append(iouEval(nr_classes, ignore))
131 | evaluators[i].reset()
132 |
133 |
134 | # get label paths
135 | if FLAGS.eval_type == "sub":
136 | label_names = load_label(FLAGS.dataset, FLAGS.sequences, "labels", "npy")
137 | else:
138 | label_names = load_label(FLAGS.dataset, FLAGS.sequences, "labels", "label")
139 |
140 | py_names = load_label(FLAGS.dataset, FLAGS.sequences, "velodyne", "bin")
141 | # py_names = py_names[0:len(py_names):4]
142 |
143 | # get predictions paths
144 | pred_names = load_label(FLAGS.predictions, FLAGS.sequences, "predictions", "label")
145 |
146 | # label_names = label_names[0:len(label_names):4]
147 |
148 | print(len(label_names), len(pred_names))
149 | assert(len(label_names) == len(pred_names))
150 |
151 |
152 | print("Evaluating sequences")
153 | N = len(label_names)
154 | # open each file, get the tensor, and make the iou comparison
155 | for i in tqdm(range(N)):
156 | label_file = label_names[i]
157 | pred_file = pred_names[i]
158 | points_file = py_names[i]
159 |
160 |
161 | # open label
162 | if FLAGS.eval_type == "sub":
163 | label = np.load(label_file)
164 | label = label.reshape((-1)) # reshape to vector
165 | else:
166 | label = np.fromfile(label_file, dtype=np.int32)
167 | label = label.reshape((-1))
168 | sem_label = label & 0xFFFF # semantic label in lower half
169 | inst_label = label >> 16 # instance id in upper half
170 | assert ((sem_label + (inst_label << 16) == label).all())
171 | label = remap_lut[sem_label]
172 |
173 | if FLAGS.limit is not None:
174 | label = label[:FLAGS.limit] # limit to desired length
175 | # open prediction
176 |
177 | pred = np.fromfile(pred_file, dtype=np.int32)
178 | pred = pred.reshape((-1)) # reshape to vector
179 | pred = remap_lut[pred]
180 | if FLAGS.limit is not None:
181 | pred = pred[:FLAGS.limit] # limit to desired length
182 | # add single scan to evaluation
183 |
184 | xyzr = np.fromfile(points_file, dtype=np.float32)
185 | xyzr = xyzr.reshape((-1, 4))
186 | points = xyzr[:, 0:3]
187 | if FLAGS.limit is not None:
188 | points = points[:FLAGS.limit] # limit to desired length
189 | depth = np.linalg.norm(points, 2, axis=1)
190 | # print(depth.shape, pred.shape, label.shape)
191 | # evaluate for all distances
192 | for idx in range(len(DISTANCES)):
193 | # select by range
194 | lrange = DISTANCES[idx][0]
195 | hrange = DISTANCES[idx][1]
196 | mask = np.logical_and(depth > lrange, depth <= hrange)
197 |
198 | # mask by distance
199 | # mask_depth = depth[mask]
200 | # print("mask range, ", mask_depth.max(), mask_depth.min())
201 | mask_label = label[mask]
202 | mask_pred = pred[mask]
203 |
204 | # add single scan to evaluation
205 | evaluators[idx].addBatch(mask_pred, mask_label)
206 |
207 | # print for all ranges
208 | print("*" * 80)
209 | for idx in range(len(DISTANCES)):
210 | # when I am done, print the evaluation
211 | m_accuracy = evaluators[idx].getacc()
212 | m_jaccard, class_jaccard = evaluators[idx].getIoU()
213 |
214 | # print for spreadsheet
215 | sys.stdout.write('range {lrange}m to {hrange}m,'.format(lrange=DISTANCES[idx][0],
216 | hrange=DISTANCES[idx][1]))
217 | for i, jacc in enumerate(class_jaccard):
218 | if i not in ignore:
219 | sys.stdout.write('{jacc:.3f}'.format(jacc=jacc.item()))
220 | sys.stdout.write(",")
221 | sys.stdout.write('{jacc:.3f}'.format(jacc=m_jaccard.item()))
222 | sys.stdout.write(",")
223 | sys.stdout.write('{acc:.3f}'.format(acc=m_accuracy.item()))
224 | sys.stdout.write('\n')
225 | sys.stdout.flush()
226 |
227 | # if FLAGS.codalab:
228 | # results = {}
229 | # for idx in range(len(DISTANCES)):
230 | # # make string for distance
231 | # d_str = str(DISTANCES[idx][-1])+"m_"
232 | #
233 | # # get values for this distance range
234 | # m_accuracy = evaluators[idx].getacc()
235 | # m_jaccard, class_jaccard = evaluators[idx].getIoU()
236 | #
237 | # # put in dictionary
238 | # results[d_str+"accuracy_mean"] = float(m_accuracy)
239 | # results[d_str+"iou_mean"] = float(m_jaccard)
240 | # for i, jacc in enumerate(class_jaccard):
241 | # if i not in ignore:
242 | # results[d_str+"iou_"+class_strings[class_inv_remap[i]]] = float(jacc)
243 | # # save to file
244 | # with open('segmentation_scores_distance.txt', 'w') as yaml_file:
245 | # yaml.dump(results, yaml_file, default_flow_style=False)
--------------------------------------------------------------------------------
/help_utils.py:
--------------------------------------------------------------------------------
1 | import os
2 | import logging
3 | import torch
4 | import numpy as np
5 |
6 | def seed_torch(seed=1024):
7 | import random
8 | import numpy as np
9 | random.seed(seed)
10 | os.environ['PYTHONHASHSEED'] = str(seed)
11 | np.random.seed(seed)
12 | torch.manual_seed(seed)
13 | torch.cuda.manual_seed(seed)
14 | torch.cuda.manual_seed_all(seed) # if you are using multi-GPU.
15 | torch.backends.cudnn.deterministic = False
16 | print("We use the seed: {}".format(seed))
17 |
18 | def my_worker_init_fn(worker_id):
19 | np.random.seed(np.random.get_state()[1][0] + worker_id)
20 |
21 | def get_logger(filename, verbosity=0, name=None):
22 | level_dict = {0: logging.DEBUG, 1: logging.INFO, 2: logging.WARNING}
23 | formatter = logging.Formatter(
24 | "[%(asctime)s][%(levelname)s] %(message)s", '%Y%m%d %H:%M:%S'
25 | )
26 | logger = logging.getLogger(name)
27 | logger.setLevel(level_dict[verbosity])
28 |
29 | fh = logging.FileHandler(filename, "w")
30 | fh.setFormatter(formatter)
31 | logger.addHandler(fh)
32 |
33 | sh = logging.StreamHandler()
34 | sh.setFormatter(formatter)
35 | logger.addHandler(sh)
36 | return logger
37 |
38 | class AverageMeter(object):
39 | """Computes and stores the average and current value"""
40 | def __init__(self):
41 | self.reset()
42 | def reset(self):
43 | self.val = 0
44 | self.avg = 0
45 | self.sum = 0
46 | self.count = 0
47 | def update(self, val, n=1):
48 | self.val = val
49 | self.sum += val * n
50 | self.count += n
51 | self.avg = self.sum / self.count
52 |
53 | def copyFiles(target_path):
54 | import shutil
55 | target_path = target_path + 'code' + "/"
56 | if not os.path.isdir(target_path):
57 | os.makedirs(target_path)
58 | shutil.copyfile("train_SemanticKITTI.py", target_path + "train_SemanticKITTI.py")
59 | shutil.copyfile("train_both_SemanticKITTI.py", target_path + "train_both_SemanticKITTI.py")
60 | shutil.copyfile("train_SemanticPOSS.py", target_path + "train_SemanticPOSS.py")
61 | shutil.copyfile("train_both_SemanticPOSS.py", target_path + "train_both_SemanticPOSS.py")
62 | shutil.copyfile("test_SemanticKITTI.py", target_path + "test_SemanticKITTI.py")
63 | shutil.copyfile("test_SemanticPOSS.py", target_path + "test_SemanticPOSS.py")
64 |
65 | shutil.copyfile("config.py", target_path + "/config.py")
66 | # shutil.copyfile("network/ResNet.py", target_path + "/ResNet.py")
67 |
68 | temp_path = target_path + '/' + 'dataset'
69 | if not os.path.isdir(temp_path):
70 | os.makedirs(temp_path)
71 | shutil.rmtree(temp_path)
72 | shutil.copytree('dataset', temp_path)
73 |
74 | temp_path = target_path + '/' + 'network'
75 | if not os.path.isdir(temp_path):
76 | os.makedirs(temp_path)
77 | shutil.rmtree(temp_path)
78 | shutil.copytree('network', temp_path)
--------------------------------------------------------------------------------
/my_env.yaml:
--------------------------------------------------------------------------------
1 | name: randla
2 | channels:
3 | - pytorch
4 | - defaults
5 | dependencies:
6 | - _libgcc_mutex=0.1=main
7 | - _openmp_mutex=5.1=1_gnu
8 | - blas=1.0=mkl
9 | - bzip2=1.0.8=h7b6447c_0
10 | - ca-certificates=2022.4.26=h06a4308_0
11 | - certifi=2021.5.30=py36h06a4308_0
12 | - cpuonly=1.0=0
13 | - cudatoolkit=11.3.1=h2bc3f7f_2
14 | - dataclasses=0.8=pyh4f3eec9_6
15 | - ffmpeg=4.3=hf484d3e_0
16 | - freetype=2.11.0=h70c0345_0
17 | - gmp=6.2.1=h2531618_2
18 | - gnutls=3.6.15=he1e5248_0
19 | - intel-openmp=2022.0.1=h06a4308_3633
20 | - jpeg=9b=h024ee3a_2
21 | - lame=3.100=h7b6447c_0
22 | - lcms2=2.12=h3be6417_0
23 | - ld_impl_linux-64=2.38=h1181459_1
24 | - libffi=3.3=he6710b0_2
25 | - libgcc-ng=11.2.0=h1234567_0
26 | - libgomp=11.2.0=h1234567_0
27 | - libiconv=1.16=h7f8727e_2
28 | - libidn2=2.3.2=h7f8727e_0
29 | - libpng=1.6.37=hbc83047_0
30 | - libstdcxx-ng=11.2.0=h1234567_0
31 | - libtasn1=4.16.0=h27cfd23_0
32 | - libtiff=4.2.0=h85742a9_0
33 | - libunistring=0.9.10=h27cfd23_0
34 | - libuv=1.40.0=h7b6447c_0
35 | - libwebp-base=1.2.2=h7f8727e_0
36 | - lz4-c=1.9.3=h295c915_1
37 | - mkl=2020.2=256
38 | - mkl-service=2.3.0=py36he8ac12f_0
39 | - mkl_fft=1.3.0=py36h54f3939_0
40 | - mkl_random=1.1.1=py36h0573a6f_0
41 | - ncurses=6.3=h7f8727e_2
42 | - nettle=3.7.3=hbbd107a_1
43 | - ninja=1.10.2=h06a4308_5
44 | - ninja-base=1.10.2=hd09550d_5
45 | - numpy-base=1.19.2=py36hfa32c7d_0
46 | - olefile=0.46=py36_0
47 | - openh264=2.1.1=h4ff587b_0
48 | - openjpeg=2.4.0=h3ad879b_0
49 | - openssl=1.1.1o=h7f8727e_0
50 | - pip=21.2.2=py36h06a4308_0
51 | - python=3.6.13=h12debd9_1
52 | - readline=8.1.2=h7f8727e_1
53 | - setuptools=58.0.4=py36h06a4308_0
54 | - six=1.16.0=pyhd3eb1b0_1
55 | - sqlite=3.38.3=hc218d9a_0
56 | - tk=8.6.11=h1ccaba5_1
57 | - typing_extensions=4.1.1=pyh06a4308_0
58 | - wheel=0.37.1=pyhd3eb1b0_0
59 | - xz=5.2.5=h7f8727e_1
60 | - zlib=1.2.12=h7f8727e_2
61 | - zstd=1.4.9=haebb681_0
62 | - pip:
63 | - absl-py==1.0.0
64 | - cachetools==4.2.4
65 | - charset-normalizer==2.0.12
66 | - cycler==0.11.0
67 | - cython==0.29.30
68 | - decorator==4.4.2
69 | - dropblock==0.3.0
70 | - freetype-py==2.2.0
71 | - future==0.18.2
72 | - glfw==2.5.3
73 | - google-auth==2.6.6
74 | - google-auth-oauthlib==0.4.6
75 | - googledrivedownloader==0.4
76 | - grpcio==1.46.3
77 | - hsluv==5.0.2
78 | - idna==3.3
79 | - imgui==1.4.1
80 | - importlib-metadata==4.8.3
81 | - importlib-resources==5.4.0
82 | - isodate==0.6.1
83 | - jinja2==3.0.3
84 | - joblib==1.1.0
85 | - kiwisolver==1.3.1
86 | - llvmlite==0.36.0
87 | - markdown==3.3.7
88 | - markupsafe==2.0.1
89 | - matplotlib==3.3.4
90 | - networkx==2.5.1
91 | - numba==0.53.1
92 | - numpy==1.18.0
93 | - oauthlib==3.2.0
94 | - pandas==1.1.5
95 | - pillow==8.4.0
96 | - plyfile==0.7.4
97 | - protobuf==3.19.4
98 | - pyasn1==0.4.8
99 | - pyasn1-modules==0.2.8
100 | - pyopengl==3.1.6
101 | - pyparsing==3.0.9
102 | - pyqt5==5.15.6
103 | - pyqt5-qt5==5.15.2
104 | - pyqt5-sip==12.9.1
105 | - python-dateutil==2.8.2
106 | - pytz==2022.1
107 | - pyyaml==6.0
108 | - rdflib==5.0.0
109 | - requests==2.27.1
110 | - requests-oauthlib==1.3.1
111 | - rsa==4.8
112 | - scikit-learn==0.22.2
113 | - scipy==1.5.4
114 | - tensorboard==2.9.0
115 | - tensorboard-data-server==0.6.1
116 | - tensorboard-plugin-wit==1.8.1
117 | - threadpoolctl==3.1.0
118 | - torch==1.9.0+cu111
119 | - torch-geometric==2.0.3
120 | - torch-scatter==2.0.7
121 | - torch-sparse==0.6.9
122 | - torchaudio==0.9.0
123 | - torchvision==0.10.0+cu111
124 | - tqdm==4.64.0
125 | - urllib3==1.26.9
126 | - vispy==0.10.0
127 | - werkzeug==2.0.3
128 | - yacs==0.1.8
129 | - zipp==3.6.0
130 | #prefix: /home/chx/anaconda3/envs/randla
131 |
--------------------------------------------------------------------------------
/network/RandLANet.py:
--------------------------------------------------------------------------------
1 | import torch
2 | import torch.nn as nn
3 | import torch.nn.functional as F
4 | import network.pytorch_utils as pt_utils
5 |
6 | class Network(nn.Module):
7 |
8 | def __init__(self, config, learn=False):
9 | super().__init__()
10 | self.config = config
11 | self.learn = learn
12 | self.fc0 = pt_utils.Conv1d(3, 8, kernel_size=1, bn=True)
13 |
14 | self.dilated_res_blocks = nn.ModuleList()
15 | d_in = 8
16 | for i in range(self.config.num_layers):
17 | d_out = self.config.d_out[i]
18 | self.dilated_res_blocks.append(Dilated_res_block(d_in, d_out))
19 | d_in = 2 * d_out
20 |
21 | d_out = d_in
22 | self.decoder_0 = pt_utils.Conv2d(d_in, d_out, kernel_size=(1, 1), bn=True)
23 |
24 | self.decoder_blocks = nn.ModuleList()
25 | for j in range(self.config.num_layers):
26 | if j < 3:
27 | d_in = d_out + 2 * self.config.d_out[-j-2]
28 | d_out = 2 * self.config.d_out[-j-2]
29 | else:
30 | d_in = 4 * self.config.d_out[-4]
31 | d_out = 2 * self.config.d_out[-4]
32 | self.decoder_blocks.append(pt_utils.Conv2d(d_in, d_out, kernel_size=(1, 1), bn=True))
33 | self.fc1 = pt_utils.Conv2d(d_out, 64, kernel_size=(1, 1), bn=True)
34 | self.fc2 = pt_utils.Conv2d(64, 32, kernel_size=(1, 1), bn=True)
35 | self.dropout = nn.Dropout(0.5)
36 | self.fc3 = pt_utils.Conv2d(32, self.config.num_classes, kernel_size=(1, 1), bn=False, activation=None)
37 |
38 | if self.learn:
39 | self.ce_sigma = nn.Parameter(torch.Tensor([0.27]), requires_grad=True)
40 | self.l1_sigma = nn.Parameter(torch.Tensor([0.62]), requires_grad=True)
41 | # self.ce_sigma = nn.Parameter(torch.Tensor(1).uniform_(0.2, 1.0), requires_grad=True)
42 | # self.l1_sigma = nn.Parameter(torch.Tensor(1).uniform_(0.2, 1.0), requires_grad=True)
43 |
44 | def forward(self, end_points):
45 | features = end_points['features'] # Batch*channel*npoints
46 | features = self.fc0(features)
47 |
48 | features = features.unsqueeze(dim=3) # Batch*channel*npoints*1
49 |
50 | # ###########################Encoder############################
51 | f_encoder_list = []
52 | for i in range(self.config.num_layers):
53 | f_encoder_i = self.dilated_res_blocks[i](features, end_points['xyz'][i], end_points['neigh_idx'][i])
54 | f_sampled_i = self.random_sample(f_encoder_i, end_points['sub_idx'][i])
55 | features = f_sampled_i
56 | if i == 0:
57 | f_encoder_list.append(f_encoder_i)
58 | f_encoder_list.append(f_sampled_i)
59 | # ###########################Encoder############################
60 | features = self.decoder_0(f_encoder_list[-1])
61 | # for i in f_encoder_list:
62 | # print(i.shape)
63 | # ###########################Decoder############################
64 | f_decoder_list = []
65 | for j in range(self.config.num_layers):
66 | f_interp_i = self.nearest_interpolation(features, end_points['interp_idx'][-j - 1])
67 | f_decoder_i = self.decoder_blocks[j](torch.cat([f_encoder_list[-j - 2], f_interp_i], dim=1))
68 |
69 | features = f_decoder_i
70 | f_decoder_list.append(f_decoder_i)
71 | # ###########################Decoder############################
72 |
73 | features = self.fc1(features)
74 | features = self.fc2(features)
75 | features = self.dropout(features)
76 | features = self.fc3(features)
77 | f_out = features.squeeze(3)
78 |
79 | if self.learn:
80 | return f_out, [self.ce_sigma, self.l1_sigma]
81 | else:
82 | return f_out
83 |
84 |
85 | @staticmethod
86 | def random_sample(feature, pool_idx):
87 | """
88 | :param feature: [B, N, d] input features matrix
89 | :param pool_idx: [B, N', max_num] N' < N, N' is the selected position after pooling
90 | :return: pool_features = [B, N', d] pooled features matrix
91 | """
92 | feature = feature.squeeze(dim=3) # batch*channel*npoints
93 | num_neigh = pool_idx.shape[-1]
94 | d = feature.shape[1]
95 | batch_size = pool_idx.shape[0]
96 | pool_idx = pool_idx.reshape(batch_size, -1) # batch*(npoints,nsamples)
97 | pool_features = torch.gather(feature, 2, pool_idx.unsqueeze(1).repeat(1, feature.shape[1], 1))
98 | pool_features = pool_features.reshape(batch_size, d, -1, num_neigh)
99 | pool_features = pool_features.max(dim=3, keepdim=True)[0] # batch*channel*npoints*1
100 | return pool_features
101 |
102 | @staticmethod
103 | def nearest_interpolation(feature, interp_idx):
104 | """
105 | :param feature: [B, N, d] input features matrix
106 | :param interp_idx: [B, up_num_points, 1] nearest neighbour index
107 | :return: [B, up_num_points, d] interpolated features matrix
108 | """
109 | feature = feature.squeeze(dim=3) # batch*channel*npoints
110 | batch_size = interp_idx.shape[0]
111 | up_num_points = interp_idx.shape[1]
112 | interp_idx = interp_idx.reshape(batch_size, up_num_points)
113 | interpolated_features = torch.gather(feature, 2, interp_idx.unsqueeze(1).repeat(1, feature.shape[1], 1))
114 | interpolated_features = interpolated_features.unsqueeze(3) # batch*channel*npoints*1
115 | return interpolated_features
116 |
117 |
118 | class Dilated_res_block(nn.Module):
119 | def __init__(self, d_in, d_out):
120 | super().__init__()
121 |
122 | self.mlp1 = pt_utils.Conv2d(d_in, d_out//2, kernel_size=(1, 1), bn=True)
123 | self.lfa = Building_block(d_out)
124 | self.mlp2 = pt_utils.Conv2d(d_out, d_out*2, kernel_size=(1, 1), bn=True, activation=None)
125 | self.shortcut = pt_utils.Conv2d(d_in, d_out*2, kernel_size=(1, 1), bn=True, activation=None)
126 |
127 | def forward(self, feature, xyz, neigh_idx):
128 | f_pc = self.mlp1(feature) # Batch*channel*npoints*1
129 | f_pc = self.lfa(xyz, f_pc, neigh_idx) # Batch*d_out*npoints*1
130 | f_pc = self.mlp2(f_pc)
131 | shortcut = self.shortcut(feature)
132 | return F.leaky_relu(f_pc+shortcut, negative_slope=0.2)
133 |
134 |
135 | class Building_block(nn.Module):
136 | def __init__(self, d_out): # d_in = d_out//2
137 | super().__init__()
138 | self.mlp1 = pt_utils.Conv2d(10, d_out//2, kernel_size=(1, 1), bn=True)
139 | self.att_pooling_1 = Att_pooling(d_out, d_out//2)
140 |
141 | self.mlp2 = pt_utils.Conv2d(d_out//2, d_out//2, kernel_size=(1, 1), bn=True)
142 | self.att_pooling_2 = Att_pooling(d_out, d_out)
143 |
144 | def forward(self, xyz, feature, neigh_idx): # feature: Batch*channel*npoints*1
145 | f_xyz = self.relative_pos_encoding(xyz, neigh_idx) # batch*npoint*nsamples*10
146 | f_xyz = f_xyz.permute((0, 3, 1, 2)) # batch*10*npoint*nsamples
147 | f_xyz = self.mlp1(f_xyz)
148 | # batch*npoint*nsamples*channel
149 | f_neighbours = self.gather_neighbour(feature.squeeze(-1).permute((0, 2, 1)), neigh_idx)
150 | f_neighbours = f_neighbours.permute((0, 3, 1, 2)) # batch*channel*npoint*nsamples
151 | f_concat = torch.cat([f_neighbours, f_xyz], dim=1)
152 | f_pc_agg = self.att_pooling_1(f_concat) # Batch*channel*npoints*1
153 |
154 | f_xyz = self.mlp2(f_xyz)
155 | # batch*npoint*nsamples*channel
156 | f_neighbours = self.gather_neighbour(f_pc_agg.squeeze(-1).permute((0, 2, 1)), neigh_idx)
157 | f_neighbours = f_neighbours.permute((0, 3, 1, 2)) # batch*channel*npoint*nsamples
158 | f_concat = torch.cat([f_neighbours, f_xyz], dim=1)
159 | f_pc_agg = self.att_pooling_2(f_concat)
160 | return f_pc_agg
161 |
162 | def relative_pos_encoding(self, xyz, neigh_idx):
163 | neighbor_xyz = self.gather_neighbour(xyz, neigh_idx) # batch*npoint*nsamples*3
164 |
165 | xyz_tile = xyz.unsqueeze(2).repeat(1, 1, neigh_idx.shape[-1], 1) # batch*npoint*nsamples*3
166 | relative_xyz = xyz_tile - neighbor_xyz # batch*npoint*nsamples*3
167 | # batch*npoint*nsamples*1
168 | relative_dis = torch.sqrt(torch.sum(torch.pow(relative_xyz, 2), dim=-1, keepdim=True))
169 | # batch*npoint*nsamples*10
170 | relative_feature = torch.cat([relative_dis, relative_xyz, xyz_tile, neighbor_xyz], dim=-1)
171 | return relative_feature
172 |
173 | @staticmethod
174 | def gather_neighbour(pc, neighbor_idx): # pc: batch*npoint*channel
175 | # gather the coordinates or features of neighboring points
176 | batch_size = pc.shape[0]
177 | num_points = pc.shape[1]
178 | d = pc.shape[2]
179 | index_input = neighbor_idx.reshape(batch_size, -1)
180 | features = torch.gather(pc, 1, index_input.unsqueeze(-1).repeat(1, 1, pc.shape[2]))
181 | features = features.reshape(batch_size, num_points, neighbor_idx.shape[-1], d) # batch*npoint*nsamples*channel
182 | return features
183 |
184 |
185 | class Att_pooling(nn.Module):
186 | def __init__(self, d_in, d_out):
187 | super().__init__()
188 | self.fc = nn.Conv2d(d_in, d_in, (1, 1), bias=False)
189 | self.mlp = pt_utils.Conv2d(d_in, d_out, kernel_size=(1, 1), bn=True)
190 |
191 | def forward(self, feature_set):
192 |
193 | att_activation = self.fc(feature_set)
194 | att_scores = F.softmax(att_activation, dim=3)
195 | f_agg = feature_set * att_scores
196 | f_agg = torch.sum(f_agg, dim=3, keepdim=True)
197 | f_agg = self.mlp(f_agg)
198 | return f_agg
199 |
200 | if __name__ == "__main__":
201 | from dataset.semkitti_trainset import SemanticKITTI
202 | from torch.utils.data import DataLoader
203 | from config import ConfigSemanticKITTI as cfg
204 | import os
205 | os.environ["CUDA_VISIBLE_DEVICES"] = '1'
206 | train_dataset = SemanticKITTI('training', sampling_way='random')
207 |
208 | train_loader = DataLoader(
209 | train_dataset, batch_size=1,shuffle=True,
210 | num_workers=0, collate_fn=train_dataset.collate_fn,)
211 |
212 | net = Network(cfg)
213 | net.cuda()
214 | for batch_idx, batch_data in enumerate(train_loader):
215 | for key in batch_data:
216 | if type(batch_data[key]) is list:
217 | for i in range(cfg.num_layers):
218 | batch_data[key][i] = batch_data[key][i].cuda(non_blocking=True)
219 | else:
220 | batch_data[key] = batch_data[key].cuda(non_blocking=True)
221 |
222 | with torch.no_grad():
223 | end_points = net(batch_data)
224 |
225 | exit()
226 |
--------------------------------------------------------------------------------
/network/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/huixiancheng/PCB-RandNet/304c84b67380a23ab66516581b348c7db2f9f0df/network/__init__.py
--------------------------------------------------------------------------------
/network/loss_func.py:
--------------------------------------------------------------------------------
1 | import torch
2 |
3 |
4 | def compute_loss(semantic_out, end_points, dataset, criterion):
5 |
6 | logits = semantic_out
7 | labels = end_points['labels']
8 |
9 | end_points = {}
10 | logits = logits.transpose(1, 2).reshape(-1, dataset.num_classes)
11 | labels = labels.reshape(-1)
12 |
13 | # Boolean mask of points that should be ignored
14 | ignored_bool = (labels == 0)
15 |
16 | for ign_label in dataset.ignored_labels:
17 | ignored_bool = ignored_bool | (labels == ign_label)
18 |
19 | # Collect logits and labels that are not ignored
20 | valid_idx = ignored_bool == 0
21 | valid_logits = logits[valid_idx, :]
22 | valid_labels_init = labels[valid_idx]
23 | # Reduce label values in the range of logit shape
24 | reducing_list = torch.arange(0, dataset.num_classes).long().to(logits.device)
25 | inserted_value = torch.zeros((1,)).long().to(logits.device)
26 | for ign_label in dataset.ignored_labels:
27 | reducing_list = torch.cat([reducing_list[:ign_label], inserted_value, reducing_list[ign_label:]], 0)
28 | valid_labels = torch.gather(reducing_list, 0, valid_labels_init)
29 | loss = criterion(valid_logits, valid_labels).mean()
30 | end_points['valid_logits'], end_points['valid_labels'] = valid_logits, valid_labels
31 | return loss, end_points
32 |
--------------------------------------------------------------------------------
/network/pytorch_utils.py:
--------------------------------------------------------------------------------
1 | import torch.nn as nn
2 | from typing import List, Tuple
3 |
4 |
5 | class SharedMLP(nn.Sequential):
6 |
7 | def __init__(
8 | self,
9 | args: List[int],
10 | *,
11 | bn: bool = False,
12 | activation=nn.ReLU(inplace=True),
13 | preact: bool = False,
14 | first: bool = False,
15 | name: str = "",
16 | instance_norm: bool = False
17 | ):
18 | super().__init__()
19 |
20 | for i in range(len(args) - 1):
21 | self.add_module(
22 | name + 'layer{}'.format(i),
23 | Conv2d(
24 | args[i],
25 | args[i + 1],
26 | bn=(not first or not preact or (i != 0)) and bn,
27 | activation=activation
28 | if (not first or not preact or (i != 0)) else None,
29 | preact=preact,
30 | instance_norm=instance_norm
31 | )
32 | )
33 |
34 |
35 | class _ConvBase(nn.Sequential):
36 |
37 | def __init__(
38 | self,
39 | in_size,
40 | out_size,
41 | kernel_size,
42 | stride,
43 | padding,
44 | activation,
45 | bn,
46 | init,
47 | conv=None,
48 | batch_norm=None,
49 | bias=True,
50 | preact=False,
51 | name="",
52 | instance_norm=False,
53 | instance_norm_func=None
54 | ):
55 | super().__init__()
56 |
57 | bias = bias and (not bn)
58 | conv_unit = conv(
59 | in_size,
60 | out_size,
61 | kernel_size=kernel_size,
62 | stride=stride,
63 | padding=padding,
64 | bias=bias
65 | )
66 | init(conv_unit.weight)
67 | if bias:
68 | nn.init.constant_(conv_unit.bias, 0)
69 |
70 | if bn:
71 | if not preact:
72 | bn_unit = batch_norm(out_size)
73 | else:
74 | bn_unit = batch_norm(in_size)
75 | if instance_norm:
76 | if not preact:
77 | in_unit = instance_norm_func(out_size, affine=False, track_running_stats=False)
78 | else:
79 | in_unit = instance_norm_func(in_size, affine=False, track_running_stats=False)
80 |
81 | if preact:
82 | if bn:
83 | self.add_module(name + 'bn', bn_unit)
84 |
85 | if activation is not None:
86 | self.add_module(name + 'activation', activation)
87 |
88 | if not bn and instance_norm:
89 | self.add_module(name + 'in', in_unit)
90 |
91 | self.add_module(name + 'conv', conv_unit)
92 |
93 | if not preact:
94 | if bn:
95 | self.add_module(name + 'bn', bn_unit)
96 |
97 | if activation is not None:
98 | self.add_module(name + 'activation', activation)
99 |
100 | if not bn and instance_norm:
101 | self.add_module(name + 'in', in_unit)
102 |
103 |
104 | class _BNBase(nn.Sequential):
105 |
106 | def __init__(self, in_size, batch_norm=None, name=""):
107 | super().__init__()
108 | self.add_module(name + "bn", batch_norm(in_size, eps=1e-6, momentum=0.99))
109 |
110 | nn.init.constant_(self[0].weight, 1.0)
111 | nn.init.constant_(self[0].bias, 0)
112 |
113 |
114 | class BatchNorm1d(_BNBase):
115 |
116 | def __init__(self, in_size: int, *, name: str = ""):
117 | super().__init__(in_size, batch_norm=nn.BatchNorm1d, name=name)
118 |
119 |
120 | class BatchNorm2d(_BNBase):
121 |
122 | def __init__(self, in_size: int, name: str = ""):
123 | super().__init__(in_size, batch_norm=nn.BatchNorm2d, name=name)
124 |
125 |
126 | class Conv1d(_ConvBase):
127 |
128 | def __init__(
129 | self,
130 | in_size: int,
131 | out_size: int,
132 | *,
133 | kernel_size: int = 1,
134 | stride: int = 1,
135 | padding: int = 0,
136 | activation=nn.LeakyReLU(negative_slope=0.2, inplace=True),
137 | bn: bool = False,
138 | init=nn.init.kaiming_normal_,
139 | bias: bool = True,
140 | preact: bool = False,
141 | name: str = "",
142 | instance_norm=False
143 | ):
144 | super().__init__(
145 | in_size,
146 | out_size,
147 | kernel_size,
148 | stride,
149 | padding,
150 | activation,
151 | bn,
152 | init,
153 | conv=nn.Conv1d,
154 | batch_norm=BatchNorm1d,
155 | bias=bias,
156 | preact=preact,
157 | name=name,
158 | instance_norm=instance_norm,
159 | instance_norm_func=nn.InstanceNorm1d
160 | )
161 |
162 |
163 | class Conv2d(_ConvBase):
164 |
165 | def __init__(
166 | self,
167 | in_size: int,
168 | out_size: int,
169 | *,
170 | kernel_size: Tuple[int, int] = (1, 1),
171 | stride: Tuple[int, int] = (1, 1),
172 | padding: Tuple[int, int] = (0, 0),
173 | activation=nn.LeakyReLU(negative_slope=0.2, inplace=True),
174 | bn: bool = False,
175 | init=nn.init.kaiming_normal_,
176 | bias: bool = True,
177 | preact: bool = False,
178 | name: str = "",
179 | instance_norm=False
180 | ):
181 | super().__init__(
182 | in_size,
183 | out_size,
184 | kernel_size,
185 | stride,
186 | padding,
187 | activation,
188 | bn,
189 | init,
190 | conv=nn.Conv2d,
191 | batch_norm=BatchNorm2d,
192 | bias=bias,
193 | preact=preact,
194 | name=name,
195 | instance_norm=instance_norm,
196 | instance_norm_func=nn.InstanceNorm2d
197 | )
198 |
199 |
200 | class FC(nn.Sequential):
201 |
202 | def __init__(
203 | self,
204 | in_size: int,
205 | out_size: int,
206 | *,
207 | activation=nn.ReLU(inplace=True),
208 | bn: bool = False,
209 | init=None,
210 | preact: bool = False,
211 | name: str = ""
212 | ):
213 | super().__init__()
214 |
215 | fc = nn.Linear(in_size, out_size, bias=not bn)
216 | if init is not None:
217 | init(fc.weight)
218 | if not bn:
219 | nn.init.constant(fc.bias, 0)
220 |
221 | if preact:
222 | if bn:
223 | self.add_module(name + 'bn', BatchNorm1d(in_size))
224 |
225 | if activation is not None:
226 | self.add_module(name + 'activation', activation)
227 |
228 | self.add_module(name + 'fc', fc)
229 |
230 | if not preact:
231 | if bn:
232 | self.add_module(name + 'bn', BatchNorm1d(out_size))
233 |
234 | if activation is not None:
235 | self.add_module(name + 'activation', activation)
236 |
237 |
238 | def set_bn_momentum_default(bn_momentum):
239 |
240 | def fn(m):
241 | if isinstance(m, (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d)):
242 | m.momentum = bn_momentum
243 |
244 | return fn
245 |
246 |
247 | class BNMomentumScheduler(object):
248 |
249 | def __init__(
250 | self, model, bn_lambda, last_epoch=-1,
251 | setter=set_bn_momentum_default
252 | ):
253 | if not isinstance(model, nn.Module):
254 | raise RuntimeError(
255 | "Class '{}' is not a PyTorch nn Module".format(
256 | type(model).__name__
257 | )
258 | )
259 |
260 | self.model = model
261 | self.setter = setter
262 | self.lmbd = bn_lambda
263 |
264 | self.step(last_epoch + 1)
265 | self.last_epoch = last_epoch
266 |
267 | def step(self, epoch=None):
268 | if epoch is None:
269 | epoch = self.last_epoch + 1
270 |
271 | self.last_epoch = epoch
272 | self.model.apply(self.setter(self.lmbd(epoch)))
273 |
--------------------------------------------------------------------------------
/test_SemanticKITTI.py:
--------------------------------------------------------------------------------
1 | # Common
2 | import os
3 | import time
4 | import yaml
5 | import logging
6 | import warnings
7 | import argparse
8 | import numpy as np
9 | from tqdm import tqdm
10 | # torch
11 | import torch
12 | import torch.nn as nn
13 | from torch.utils.data import DataLoader
14 | # my module
15 | from dataset.semkitti_testset import SemanticKITTI
16 |
17 | import pickle
18 | from help_utils import get_logger
19 |
20 | np.random.seed(0)
21 |
22 | # os.environ["CUDA_VISIBLE_DEVICES"] = '1'
23 | torch.backends.cudnn.enabled = False
24 |
25 | warnings.filterwarnings("ignore")
26 | parser = argparse.ArgumentParser()
27 | parser.add_argument('--infer_type', default='all', type=str, choices=['all', 'sub'], help='Infer ALL or just infer Subsample')
28 | parser.add_argument('--sampling', type=str, default='random', choices=['random', 'polar'], help='Polar sample or not')
29 | parser.add_argument('--backbone', type=str, default='randla', choices=['randla', 'baflac', 'baaf'])
30 | parser.add_argument('--checkpoint_path', default=None, help='Model checkpoint path [default: None]')
31 | parser.add_argument('--test_id', default='08', type=str, help='Predicted sequence id [default: 08]')
32 | parser.add_argument('--result_dir', default='test_pred', help='Dump dir to save prediction [default: result/]')
33 | parser.add_argument('--yaml_config', default='utils/semantic-kitti.yaml', help='semantic-kitti.yaml path')
34 | parser.add_argument('--batch_size', type=int, default=12, help='Batch Size during training [default: 30]')
35 | parser.add_argument('--step', type=int, default=0, help='sub dataset size')
36 | parser.add_argument('--grid', nargs='+', type=int, default=[64, 64, 16], help='grid size of BEV representation')
37 | FLAGS = parser.parse_args()
38 |
39 | if FLAGS.backbone == 'baflac':
40 | from config import ConfigSemanticKITTI_BAF as cfg
41 | else:
42 | from config import ConfigSemanticKITTI as cfg
43 |
44 | class Tester:
45 | def __init__(self):
46 | # Init Logging
47 | os.makedirs(FLAGS.result_dir, exist_ok=True)
48 | log_fname = os.path.join(FLAGS.result_dir, 'log_test.txt')
49 | self.logger = get_logger(log_fname, name="Tester")
50 | argsDict = FLAGS.__dict__
51 | for eachArg, value in argsDict.items():
52 | self.logger.info(eachArg + ' : ' + str(value))
53 |
54 | self.logger.info(vars(cfg))
55 | # load yaml file
56 | self.remap_lut = self.load_yaml(FLAGS.yaml_config)
57 |
58 | test_dataset = SemanticKITTI('test', test_id=FLAGS.test_id, batch_size=FLAGS.batch_size,
59 | sampling_way=FLAGS.sampling, step=FLAGS.step, grid=FLAGS.grid)
60 |
61 | device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
62 | if FLAGS.backbone == 'baflac':
63 | from network.BAF_LAC import BAF_LAC
64 | self.logger.info("Use Baseline: BAF-LAC")
65 | self.net = BAF_LAC(cfg)
66 | self.net.to(device)
67 | collate_fn = test_dataset.collate_fn_baf_lac
68 |
69 | elif FLAGS.backbone == 'randla':
70 | from network.RandLANet import Network
71 | self.logger.info("Use Baseline: Rand-LA")
72 | self.net = Network(cfg)
73 | self.net.to(device)
74 | collate_fn = test_dataset.collate_fn
75 | elif FLAGS.backbone == 'baaf':
76 | from network.BAAF import Network
77 | self.logger.info("Use Baseline: BAAF")
78 | self.net = Network(cfg)
79 | self.net.to(device)
80 | collate_fn = test_dataset.collate_fn
81 | else:
82 | raise TypeError("1~5!! can can need !!!")
83 |
84 | # get_dataset & dataloader
85 | self.test_loader = DataLoader(test_dataset, batch_size=None, collate_fn=collate_fn,
86 | pin_memory=True, num_workers=0)
87 |
88 | # Load module
89 | CHECKPOINT_PATH = FLAGS.checkpoint_path
90 | if CHECKPOINT_PATH is not None and os.path.isfile(CHECKPOINT_PATH):
91 | self.logger.info("Loading from " + CHECKPOINT_PATH)
92 | checkpoint = torch.load(CHECKPOINT_PATH)
93 | try:
94 | self.logger.info("Loading strict=True")
95 | self.net.load_state_dict(checkpoint['model_state_dict'], strict=True)
96 | except:
97 | self.logger.info("Loading strict=False")
98 | self.net.load_state_dict(checkpoint['model_state_dict'], strict=False)
99 | self.test_dataset = test_dataset
100 | # Initialize testing probability
101 | self.test_dataset.init_prob()
102 | self.test_probs = self.init_prob()
103 | self.test_smooth = 0.98
104 |
105 | def load_yaml(self, path):
106 | DATA = yaml.safe_load(open(path, 'r'))
107 | # get number of interest classes, and the label mappings
108 | remapdict = DATA["learning_map_inv"]
109 | # make lookup table for mapping
110 | maxkey = max(remapdict.keys())
111 | # +100 hack making lut bigger just in case there are unknown labels
112 | remap_lut = np.zeros((maxkey + 100), dtype=np.int32)
113 | remap_lut[list(remapdict.keys())] = list(remapdict.values())
114 | return remap_lut
115 |
116 | def init_prob(self):
117 | probs = []
118 | for item in self.test_dataset.possibility:
119 | prob = np.zeros(shape=[len(item), self.test_dataset.num_classes], dtype=np.float32)
120 | probs.append(prob)
121 | return probs
122 |
123 | def test(self):
124 | self.logger.info("Start Testing")
125 | self.rolling_predict()
126 | # Merge Probability
127 | self.merge_and_store()
128 |
129 | def rolling_predict(self):
130 | self.net.eval() # set model to eval mode (for bn and dp)
131 |
132 | iter_loader = iter(self.test_loader)
133 |
134 | start = time.time()
135 | step_id = 0
136 | with torch.no_grad():
137 | min_possibility = self.test_dataset.min_possibility
138 | while np.min(min_possibility) <= 0.5:
139 | batch_data, input_inds, cloud_inds, min_possibility = next(iter_loader)
140 | for key in batch_data:
141 | if type(batch_data[key]) is list:
142 | for i in range(cfg.num_layers):
143 | batch_data[key][i] = batch_data[key][i].cuda(non_blocking=True)
144 | else:
145 | batch_data[key] = batch_data[key].cuda(non_blocking=True)
146 | # Forward pass
147 | with torch.cuda.amp.autocast():
148 | semantic_out = self.net(batch_data)
149 | semantic_out = semantic_out.transpose(1, 2)
150 | # update prediction (multi-thread)
151 | self.update_predict(semantic_out, batch_data, input_inds, cloud_inds)
152 |
153 | if step_id % 10 == 0:
154 | end = time.time()
155 | self.logger.info('Step {:3d}, Min possibility = {:.6e}, Time = {:.3f}'.format(step_id, np.min(min_possibility), end-start))
156 | start = time.time()
157 | step_id += 1
158 |
159 |
160 | def update_predict(self, semantic_out, batch_data, input_inds, cloud_inds):
161 | # Store logits into list
162 | B = semantic_out.size(0)
163 | semantic_out = semantic_out.cpu().numpy()
164 | for j in range(B):
165 | probs = semantic_out[j]
166 | inds = input_inds[j]
167 | c_i = cloud_inds[j][0]
168 | self.test_probs[c_i][inds] = \
169 | self.test_smooth * self.test_probs[c_i][inds] + (1 - self.test_smooth) * probs
170 |
171 | def merge_and_store(self):
172 | # initialize result directory
173 | root_dir = os.path.join(FLAGS.result_dir, self.test_dataset.test_scan_number, 'predictions')
174 | os.makedirs(root_dir, exist_ok=True)
175 | self.logger.info(f'mkdir {root_dir}')
176 | N = len(self.test_probs)
177 | for j in tqdm(range(N)):
178 | if FLAGS.infer_type == 'all':
179 | proj_path = os.path.join(self.test_dataset.dataset_path, self.test_dataset.test_scan_number, 'proj')
180 | proj_file = os.path.join(proj_path, self.test_dataset.data_list[j][1] + '_proj.pkl')
181 | if os.path.isfile(proj_file):
182 | with open(proj_file, 'rb') as f:
183 | proj_inds = pickle.load(f)
184 | probs = self.test_probs[j][proj_inds[0], :]
185 | pred = np.argmax(probs, 1).astype(np.uint32)
186 | elif FLAGS.infer_type == 'sub':
187 | pred = np.argmax(self.test_probs[j], 1).astype(np.uint32)
188 | else:
189 | raise TypeError("Choose what you want to infer")
190 |
191 | # pred += 1
192 |
193 | pred = self.remap(pred)
194 | name = self.test_dataset.data_list[j][1] + '.label'
195 | output_path = os.path.join(root_dir, name)
196 | pred.tofile(output_path)
197 |
198 |
199 | def remap(self, label):
200 | upper_half = label >> 16 # get upper half for instances
201 | lower_half = label & 0xFFFF # get lower half for semantics
202 | lower_half = self.remap_lut[lower_half] # do the remapping of semantics
203 | label = (upper_half << 16) + lower_half # reconstruct full label
204 | label = label.astype(np.uint32)
205 | return label
206 |
207 |
208 | def main():
209 | tester = Tester()
210 | tester.test()
211 |
212 |
213 | if __name__ == '__main__':
214 | main()
215 |
--------------------------------------------------------------------------------
/test_SemanticPOSS.py:
--------------------------------------------------------------------------------
1 | # Common
2 | import os
3 | import time
4 | import yaml
5 | import logging
6 | import warnings
7 | import argparse
8 | import numpy as np
9 | from tqdm import tqdm
10 | # torch
11 | import torch
12 | import torch.nn as nn
13 | from torch.utils.data import DataLoader
14 | # my module
15 | from config import ConfigSemanticPOSS as cfg
16 | from dataset.poss_testset import SemanticPOSS
17 |
18 | import pickle
19 | from help_utils import get_logger
20 |
21 | np.random.seed(0)
22 | # os.environ["CUDA_VISIBLE_DEVICES"] = '1'
23 | torch.backends.cudnn.enabled = False
24 |
25 | warnings.filterwarnings("ignore")
26 | parser = argparse.ArgumentParser()
27 | parser.add_argument('--infer_type', default='all', type=str, choices=['all', 'sub'], help='Infer ALL or just infer Subsample')
28 | parser.add_argument('--sampling', type=str, default='random', choices=['random', 'polar'], help='Polar sample or not')
29 | parser.add_argument('--backbone', type=str, default='randla', choices=['randla', 'baflac', 'baaf'])
30 |
31 | parser.add_argument('--checkpoint_path', default=None, help='Model checkpoint path [default: None]')
32 | parser.add_argument('--test_id', default='03', type=str, help='Predicted sequence id [default: 08]')
33 | parser.add_argument('--result_dir', default='test_POSS_pred', help='Dump dir to save prediction [default: result/]')
34 | parser.add_argument('--yaml_config', default='utils/semantic-poss.yaml', help='semantic-kitti.yaml path')
35 | parser.add_argument('--batch_size', type=int, default=12, help='Batch Size during training [default: 30]')
36 | parser.add_argument('--grid', nargs='+', type=int, default=[64, 64, 16], help='grid size of BEV representation')
37 | FLAGS = parser.parse_args()
38 |
39 | if FLAGS.backbone == 'baflac':
40 | from config import ConfigSemanticPOSS_BAF as cfg
41 | else:
42 | from config import ConfigSemanticPOSS as cfg
43 |
44 | class Tester:
45 | def __init__(self):
46 | # Init Logging
47 | os.makedirs(FLAGS.result_dir, exist_ok=True)
48 | log_fname = os.path.join(FLAGS.result_dir, 'log_test.txt')
49 | self.logger = get_logger(log_fname, name="Tester")
50 | argsDict = FLAGS.__dict__
51 | for eachArg, value in argsDict.items():
52 | self.logger.info(eachArg + ' : ' + str(value))
53 |
54 | self.logger.info(vars(cfg))
55 | # load yaml file
56 | self.remap_lut = self.load_yaml(FLAGS.yaml_config)
57 |
58 | test_dataset = SemanticPOSS('test', test_id=FLAGS.test_id, batch_size=FLAGS.batch_size,
59 | sampling_way=FLAGS.sampling, grid=FLAGS.grid)
60 |
61 | device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
62 | if FLAGS.backbone == 'baflac':
63 | from network.BAF_LAC import BAF_LAC
64 | self.logger.info("Use Baseline: BAF-LAC")
65 | self.net = BAF_LAC(cfg)
66 | self.net.to(device)
67 | collate_fn = test_dataset.collate_fn_baf_lac
68 |
69 | elif FLAGS.backbone == 'randla':
70 | from network.RandLANet import Network
71 | self.logger.info("Use Baseline: Rand-LA")
72 | self.net = Network(cfg)
73 | self.net.to(device)
74 | collate_fn = test_dataset.collate_fn
75 | elif FLAGS.backbone == 'baaf':
76 | from network.BAAF import Network
77 | self.logger.info("Use Baseline: BAAF")
78 | self.net = Network(cfg)
79 | self.net.to(device)
80 | collate_fn = test_dataset.collate_fn
81 | else:
82 | raise TypeError("1~5!! can can need !!!")
83 |
84 | # get_dataset & dataloader
85 | self.test_loader = DataLoader(test_dataset, batch_size=None, collate_fn=collate_fn,
86 | pin_memory=True, num_workers=0)
87 |
88 | # Load module
89 | CHECKPOINT_PATH = FLAGS.checkpoint_path
90 | if CHECKPOINT_PATH is not None and os.path.isfile(CHECKPOINT_PATH):
91 | self.logger.info("Loading from " + CHECKPOINT_PATH)
92 | checkpoint = torch.load(CHECKPOINT_PATH)
93 | self.net.load_state_dict(checkpoint['model_state_dict'], strict=False)
94 |
95 | self.test_dataset = test_dataset
96 | # Initialize testing probability
97 | self.test_dataset.init_prob()
98 | self.test_probs = self.init_prob()
99 | self.test_smooth = 0.98
100 |
101 | def load_yaml(self, path):
102 | DATA = yaml.safe_load(open(path, 'r'))
103 | # get number of interest classes, and the label mappings
104 | remapdict = DATA["learning_map_inv"]
105 | # make lookup table for mapping
106 | maxkey = max(remapdict.keys())
107 | # +100 hack making lut bigger just in case there are unknown labels
108 | remap_lut = np.zeros((maxkey + 100), dtype=np.int32)
109 | remap_lut[list(remapdict.keys())] = list(remapdict.values())
110 | return remap_lut
111 |
112 | def init_prob(self):
113 | probs = []
114 | for item in self.test_dataset.possibility:
115 | prob = np.zeros(shape=[len(item), self.test_dataset.num_classes], dtype=np.float32)
116 | probs.append(prob)
117 | return probs
118 |
119 | def test(self):
120 | self.logger.info("Start Testing")
121 | self.rolling_predict()
122 | # Merge Probability
123 | self.merge_and_store()
124 |
125 | def rolling_predict(self):
126 | self.net.eval() # set model to eval mode (for bn and dp)
127 |
128 | iter_loader = iter(self.test_loader)
129 |
130 | start = time.time()
131 | step_id = 0
132 | with torch.no_grad():
133 | min_possibility = self.test_dataset.min_possibility
134 | while np.min(min_possibility) <= 0.5:
135 | batch_data, input_inds, cloud_inds, min_possibility = next(iter_loader)
136 | for key in batch_data:
137 | if type(batch_data[key]) is list:
138 | for i in range(cfg.num_layers):
139 | batch_data[key][i] = batch_data[key][i].cuda(non_blocking=True)
140 | else:
141 | batch_data[key] = batch_data[key].cuda(non_blocking=True)
142 | # Forward pass
143 | with torch.cuda.amp.autocast():
144 | semantic_out = self.net(batch_data)
145 | semantic_out = semantic_out.transpose(1, 2)
146 | # update prediction (multi-thread)
147 | self.update_predict(semantic_out, batch_data, input_inds, cloud_inds)
148 |
149 | if step_id % 10 == 0:
150 | end = time.time()
151 | self.logger.info('Step {:3d}, Min possibility = {:.6e}, Time = {:.3f}'.format(step_id, np.min(min_possibility), end-start))
152 | start = time.time()
153 | step_id += 1
154 |
155 |
156 | def update_predict(self, semantic_out, batch_data, input_inds, cloud_inds):
157 | # Store logits into list
158 | B = semantic_out.size(0)
159 | semantic_out = semantic_out.cpu().numpy()
160 | for j in range(B):
161 | probs = semantic_out[j]
162 | inds = input_inds[j]
163 | c_i = cloud_inds[j][0]
164 | self.test_probs[c_i][inds] = \
165 | self.test_smooth * self.test_probs[c_i][inds] + (1 - self.test_smooth) * probs
166 |
167 | def merge_and_store(self):
168 | # initialize result directory
169 | root_dir = os.path.join(FLAGS.result_dir, self.test_dataset.test_scan_number, 'predictions')
170 | os.makedirs(root_dir, exist_ok=True)
171 | self.logger.info(f'mkdir {root_dir}')
172 | N = len(self.test_probs)
173 | for j in tqdm(range(N)):
174 | if FLAGS.infer_type == 'all':
175 | proj_path = os.path.join(self.test_dataset.dataset_path, self.test_dataset.test_scan_number, 'proj')
176 | proj_file = os.path.join(proj_path, self.test_dataset.data_list[j][1] + '_proj.pkl')
177 | if os.path.isfile(proj_file):
178 | with open(proj_file, 'rb') as f:
179 | proj_inds = pickle.load(f)
180 | probs = self.test_probs[j][proj_inds[0], :]
181 | pred = np.argmax(probs, 1).astype(np.uint32)
182 | elif FLAGS.infer_type == 'sub':
183 | pred = np.argmax(self.test_probs[j], 1).astype(np.uint32)
184 | else:
185 | raise TypeError("Choose what you want to infer")
186 |
187 | pred += 1
188 |
189 | pred = self.remap(pred)
190 | name = self.test_dataset.data_list[j][1] + '.label'
191 | output_path = os.path.join(root_dir, name)
192 | pred.tofile(output_path)
193 |
194 |
195 | def remap(self, label):
196 | upper_half = label >> 16 # get upper half for instances
197 | lower_half = label & 0xFFFF # get lower half for semantics
198 | lower_half = self.remap_lut[lower_half] # do the remapping of semantics
199 | label = (upper_half << 16) + lower_half # reconstruct full label
200 | label = label.astype(np.uint32)
201 | return label
202 |
203 |
204 | def main():
205 | tester = Tester()
206 | tester.test()
207 |
208 |
209 | if __name__ == '__main__':
210 | main()
211 |
--------------------------------------------------------------------------------
/tool/caculate_time.py:
--------------------------------------------------------------------------------
1 | import time
2 | from matplotlib import pyplot as plt
3 | from utils.data_process import DataProcessing as DP
4 | from config import ConfigSemanticKITTI as cfg
5 | from os.path import join
6 | import numpy as np
7 | import pickle
8 | from tqdm import tqdm
9 |
10 | name = 'SemanticKITTI'
11 | dataset_path = '/home/chx/Work/semantic-kitti-randla/sequences'
12 |
13 | # name = 'SemanticPOSS'
14 | # dataset_path = '/home/chx/Work/SemanticPOSS_dataset_grid_0.01/sequences'
15 |
16 | num_classes = cfg.num_classes
17 | ignored_labels = np.sort([0])
18 |
19 | mode = 'training'
20 | sampling = ['random', 'polar', 'fps']
21 | # sampling_way = 'random'
22 |
23 |
24 | seq_list = ['00', '01', '02', '03', '04', '05', '06', '07', '09', '10']
25 | data_list = DP.get_file_list(dataset_path, seq_list)
26 | data_list = sorted(data_list)
27 | print("Using {} scans from sequences {}".format(len(data_list), seq_list))
28 |
29 |
30 | import random
31 | seed = 2
32 | random.seed(seed)
33 | random.shuffle(data_list)
34 |
35 |
36 | def get_data(file_path):
37 | seq_id = file_path[0]
38 | frame_id = file_path[1]
39 | kd_tree_path = join(dataset_path, seq_id, 'KDTree', frame_id + '.pkl')
40 | # read pkl with search tree
41 | with open(kd_tree_path, 'rb') as f:
42 | search_tree = pickle.load(f)
43 | points = np.array(search_tree.data, copy=False)
44 | # load labels
45 | label_path = join(dataset_path, seq_id, 'labels', frame_id + '.npy')
46 | labels = np.squeeze(np.load(label_path))
47 | return points, search_tree, labels
48 | def cart2polar(input_xyz):
49 | rho = np.sqrt(input_xyz[:, 0] ** 2 + input_xyz[:, 1] ** 2)
50 | phi = np.arctan2(input_xyz[:, 1], input_xyz[:, 0])
51 | return np.stack((rho, phi, input_xyz[:, 2]), axis=1)
52 | def polar_samplr(mid_xyz, grid):
53 | xyz_pol = cart2polar(mid_xyz)
54 |
55 | max_bound_r = min(np.percentile(xyz_pol[:, 0], 100, axis=0), 50)
56 | min_bound_r = max(np.percentile(xyz_pol[:, 0], 0, axis=0), 3)
57 | max_bound_p = np.max(xyz_pol[:, 1], axis=0)
58 | min_bound_p = np.min(xyz_pol[:, 1], axis=0)
59 | max_bound_z = min(np.max(xyz_pol[:, 2], axis=0), 1.5)
60 | min_bound_z = max(np.min(xyz_pol[:, 2], axis=0), -3)
61 | max_bound = np.concatenate(([max_bound_r], [max_bound_p], [max_bound_z]))
62 | min_bound = np.concatenate(([min_bound_r], [min_bound_p], [min_bound_z]))
63 | cur_grid_size = np.asarray(grid)
64 | crop_range = max_bound - min_bound
65 | intervals = crop_range / (cur_grid_size - 1)
66 |
67 | if (intervals == 0).any(): print("Zero interval!")
68 |
69 | grid_ind = (np.floor((np.clip(xyz_pol, min_bound, max_bound) - min_bound) / intervals)).astype(np.int)
70 | keys, revers, counts = np.unique(grid_ind, return_inverse=True, return_counts=True, axis=0)
71 |
72 | idx = np.argsort(revers)
73 | mid_xyz = mid_xyz[idx]
74 | slic = counts.cumsum()
75 | slic = np.insert(slic, 0, 0)
76 | left_xyz, right_xyz = np.zeros([0, 3]), np.zeros([0, 3])
77 |
78 | small = counts[counts < 4]
79 | new_nums = len(mid_xyz) // 4 - sum(small)
80 | new_grid = len(counts) - len(small)
81 | sample_list = []
82 | for i in range(new_grid):
83 | curr = new_nums // new_grid
84 | sample_list.append(curr)
85 | new_nums -= curr
86 | new_grid -= 1
87 | sample_list = np.array(sample_list)
88 | sample_list = DP.shuffle_idx(sample_list)
89 | idx = 0
90 | for i in range(len(counts)):
91 | select_xyz = mid_xyz[slic[i]:slic[i + 1]]
92 | np.random.shuffle(select_xyz)
93 | nubs = counts[i]
94 | if nubs >= 4:
95 | downs_n = sample_list[idx]
96 | idx += 1
97 | # downs_n = new_num_points_each_grid
98 | left_xyz = np.concatenate((left_xyz, select_xyz[0:downs_n]), axis=0)
99 | right_xyz = np.concatenate((right_xyz, select_xyz[downs_n:]), axis=0)
100 | else:
101 | left_xyz = np.concatenate((left_xyz, select_xyz), axis=0)
102 |
103 | supp = len(mid_xyz) // 4 - len(left_xyz)
104 | if supp == 0:
105 | pass
106 | elif supp > 0:
107 | np.random.shuffle(right_xyz)
108 | left_xyz = np.concatenate((left_xyz, right_xyz[0:supp]))
109 | right_xyz = right_xyz[supp:]
110 | else:
111 | np.random.shuffle(right_xyz)
112 | supp = len(mid_xyz) // 4
113 | left_xyz, supp_xyz = left_xyz[:supp], left_xyz[supp:]
114 | right_xyz = np.concatenate((supp_xyz, right_xyz))
115 |
116 | np.random.shuffle(left_xyz)
117 | return left_xyz
118 | # np.random.shuffle(right_xyz)
119 | # select_points = np.concatenate((left_xyz, right_xyz), axis=0)
120 | # return select_points
121 | def farthest_point_sample(point, npoint):
122 | """
123 | Input:
124 | xyz: pointcloud data, [N, D]
125 | npoint: number of samples
126 | Return:
127 | centroids: sampled pointcloud index, [npoint, D]
128 | """
129 | N, D = point.shape
130 | xyz = point[:, :3]
131 | centroids = np.zeros((npoint,))
132 | distance = np.ones((N,)) * 1e10
133 | farthest = np.random.randint(0, N)
134 | for i in range(npoint):
135 | centroids[i] = farthest
136 | centroid = xyz[farthest, :]
137 | dist = np.sum((xyz - centroid) ** 2, -1)
138 | mask = dist < distance
139 | distance[mask] = dist[mask]
140 | farthest = np.argmax(distance, -1)
141 | point = point[centroids.astype(np.int32)]
142 | return point
143 |
144 | # for sampling_way in sampling:
145 |
146 |
147 | a = [11264, 2816, 704, 176]
148 |
149 | start = 2
150 | for start in range(1, 5):
151 | for sampling_way in sampling:
152 | test = []
153 | for list in range(30):
154 | points, search_tree, labels = get_data(data_list[list])
155 | np.random.seed(seed)
156 | pick_idx = np.random.choice(len(points), 1)
157 | # print(data_list[list], len(points), pick_idx)
158 | center_point = points[pick_idx, :].reshape(1, -1)
159 | select_idx = search_tree.query(center_point, k=cfg.num_points)[1][0]
160 |
161 | if sampling_way == "random":
162 | st = time.time()
163 | select_idx = DP.shuffle_idx(select_idx)
164 | select_points = points[select_idx]
165 | for sub in a[:start]:
166 | select_points = select_points[0:sub]
167 | ed = time.time() - st
168 | test.append(ed)
169 | if list==0: print('check', select_points.shape)
170 | elif sampling_way == "fps":
171 | fps_xyz = points[select_idx]
172 | st = time.time()
173 | for sub in a[:start]:
174 | fps_xyz = farthest_point_sample(fps_xyz, sub)
175 | ed = time.time() - st
176 | test.append(ed)
177 | if list==0: print('check', fps_xyz.shape)
178 |
179 | elif sampling_way == "polar":
180 | select_points = points[select_idx]
181 | grid = [64, 64, 16]
182 | st = time.time()
183 | select_points = polar_samplr(select_points, grid)
184 | for sub in a[:start]:
185 | select_points = select_points[0:sub]
186 | ed = time.time() - st
187 | test.append(ed)
188 | if list==0: print('check', select_points.shape)
189 |
190 | print(sampling_way)
191 | print(a[:start])
192 | print("Mean time:{}\t std:{}".format(np.mean(test), np.std(test)))
193 | start += 1
194 |
195 |
196 |
--------------------------------------------------------------------------------
/tool/draw.py:
--------------------------------------------------------------------------------
1 | import matplotlib.pyplot as plt
2 | import matplotlib.ticker as ticker
3 | from matplotlib.pyplot import MultipleLocator
4 |
5 |
6 | plt.rcParams['font.family'] = 'Times New Roman'
7 | plt.rcParams.update({'font.size': 15})
8 | # plt.rc('font',family='Times New Roman', size=12)
9 |
10 | # plt.title('title name') # Title
11 | fig = plt.figure(figsize=(8, 6), dpi=300)
12 | plt.xlabel('Distance to sensor [m]') # x value
13 | plt.ylabel('Mean IoU [%]') # y value
14 | plt.grid(True)
15 |
16 | dist = [1, 2, 3, 4, 5, 6]
17 |
18 | # Polar_Balance_Random = [56.66, 51.62, 44.07, 32.34, 24.54, 10.16]
19 | Random = [58.73, 52.51, 44.83, 32.10, 23.61, 9.37]
20 |
21 | # plt.plot(dist, Polar_Balance_Random, color='r', linestyle='-', marker='o', label='Polar Random Sampling')
22 | plt.plot(dist, Random, color='b', linestyle='-', marker='s', label='Random Sampling')
23 |
24 | x = [1, 2, 3, 4, 5, 6]
25 | labels = ['(0, 10]', '(10, 20]', '(20, 30]', '(30, 40]', '(40, 50]', '(50, max)']
26 | plt.xticks(x, labels)
27 |
28 | # plt.xlim(0, 7)
29 | plt.ylim(top=60)
30 |
31 | # plt.gca().xaxis.set_major_locator(ticker.MultipleLocator(10))
32 |
33 | plt.legend()
34 |
35 | # **************************************************************************************************************************
36 | fig.set_size_inches(8, 6, forward=True)
37 | fig.tight_layout()
38 |
39 | plt.savefig("distance.png") #, bbox_inches='tight'
40 | plt.show()
41 | plt.close()
42 |
--------------------------------------------------------------------------------
/tool/draw_poss.py:
--------------------------------------------------------------------------------
1 | import matplotlib.pyplot as plt
2 | import matplotlib.ticker as ticker
3 | from matplotlib.pyplot import MultipleLocator
4 |
5 |
6 | plt.rcParams['font.family'] = 'Times New Roman'
7 | plt.rcParams.update({'font.size': 12})
8 | # plt.rc('font',family='Times New Roman', size=12)
9 |
10 | # plt.title('title name') # Title
11 | fig = plt.figure(figsize=(8, 6), dpi=300)
12 | plt.xlabel('Distance to sensor [m]') # x value
13 | plt.ylabel('Mean IoU [%]') # y value
14 | plt.grid(True)
15 |
16 | dist = [1, 2, 3, 4, 5, 6, 7, 8, 9]
17 |
18 | Polar_Balance_Random = [62.92, 59.77, 56.52, 48.89, 38.54, 34.00, 27.12, 21.41, 12.35]
19 | Random = [59.58, 60.78, 54.43, 43.71, 32.88, 27.05, 22.14, 18.39, 11.97]
20 |
21 | plt.plot(dist, Polar_Balance_Random, color='r', linestyle='-', marker='o', label='Polar Cylinder Balanced Random Sampling')
22 | plt.plot(dist, Random, color='b', linestyle='-', marker='s', label='Random Sampling')
23 |
24 | x = [1, 2, 3, 4, 5, 6, 7, 8, 9]
25 | labels = ['(0, 10]', '(10, 20]', '(20, 30]', '(30, 40]', '(40, 50]', '(50, 60]', '(60, 70]', '(70, 80]', '(80, max)']
26 | plt.xticks(x, labels)
27 |
28 | # plt.xlim(0, 7)
29 | plt.ylim(top=64)
30 |
31 | # plt.gca().xaxis.set_major_locator(ticker.MultipleLocator(10))
32 |
33 | plt.legend()
34 |
35 | # **************************************************************************************************************************
36 | fig.set_size_inches(8, 6, forward=True)
37 | fig.tight_layout()
38 |
39 | plt.savefig("distance_poss.png") #, bbox_inches='tight'
40 | plt.show()
41 | plt.close()
42 |
--------------------------------------------------------------------------------
/tool/eval_KITTI_gap.py:
--------------------------------------------------------------------------------
1 | # Common
2 | import os
3 | import logging
4 | import warnings
5 | import argparse
6 | import numpy as np
7 | from tqdm import tqdm
8 | # torch
9 | import torch
10 | import torch.nn as nn
11 | import torch.optim as optim
12 | from torch.utils.data import DataLoader
13 |
14 | from dataset.two_semkitti_trainset import SemanticKITTI
15 | from utils.metric import compute_acc, IoUCalculator, iouEval
16 |
17 | import torch.nn.functional as F
18 | from help_utils import seed_torch, my_worker_init_fn, get_logger, copyFiles, AverageMeter
19 |
20 | os.environ["CUDA_VISIBLE_DEVICES"] = '0'
21 |
22 | warnings.filterwarnings("ignore")
23 | parser = argparse.ArgumentParser()
24 | parser.add_argument('--backbone', type=str, default='randla', choices=['randla', 'baflac', 'baaf'])
25 | parser.add_argument('--checkpoint_path', default=None, help='Model checkpoint path [default: None]')
26 | parser.add_argument('--log_dir', type=str, default='GAP+TEST', help='Dump dir to save model checkpoint [default: log]')
27 | parser.add_argument('--val_batch_size', type=int, default=1, help='Batch Size during training [default: 30]')
28 | parser.add_argument('--num_workers', type=int, default=0, help='Number of workers [default: 5]')
29 | parser.add_argument('--seed', type=int, default=1024, help='Polar sample or not')
30 | parser.add_argument('--step', type=int, default=0, help='sub dataset size')
31 | FLAGS = parser.parse_args()
32 |
33 | seed_torch(FLAGS.seed)
34 | torch.backends.cudnn.enabled = False
35 |
36 | if FLAGS.backbone == 'baflac':
37 | from config import ConfigSemanticKITTI_BAF as cfg
38 | else:
39 | from config import ConfigSemanticKITTI as cfg
40 |
41 | class Trainer:
42 | def __init__(self):
43 | # Init Logging
44 | save_path = './save_semantic_poss/' + FLAGS.log_dir + '/'
45 | if not (os.path.exists(save_path)):
46 | os.makedirs(save_path)
47 | copyFiles(save_path)
48 | self.log_dir = save_path
49 | log_fname = os.path.join(self.log_dir, 'log_train.txt')
50 | self.logger = get_logger(log_fname, name='Train')
51 |
52 | argsDict = FLAGS.__dict__
53 | for eachArg, value in argsDict.items():
54 | self.logger.info(eachArg + ' : ' + str(value))
55 |
56 | # tensorboard writer
57 | val_dataset = SemanticKITTI('validation', step=FLAGS.step)
58 |
59 | # Network & Optimizer
60 | self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
61 | if FLAGS.backbone == 'baflac':
62 | from network.BAF_LAC import BAF_LAC
63 | self.logger.info("Use Baseline: BAF-LAC")
64 | self.net = BAF_LAC(cfg)
65 | self.net.to(self.device)
66 | collate_fn = val_dataset.collate_fn_baf_lac
67 |
68 | elif FLAGS.backbone == 'randla':
69 | from network.RandLANet import Network
70 | self.logger.info("Use Baseline: Rand-LA")
71 | self.net = Network(cfg)
72 | self.net.to(self.device)
73 | collate_fn = val_dataset.collate_fn
74 |
75 | elif FLAGS.backbone == 'baaf':
76 | from network.BAAF import Network
77 | self.logger.info("Use Baseline: Rand-LA")
78 | self.net = Network(cfg)
79 | self.net.to(self.device)
80 | collate_fn = val_dataset.collate_fn
81 |
82 | else:
83 | raise TypeError("1~5~!! can can need !!!")
84 |
85 | self.val_loader = DataLoader(
86 | val_dataset, batch_size=FLAGS.val_batch_size,
87 | shuffle=False, num_workers=FLAGS.num_workers,
88 | worker_init_fn=my_worker_init_fn, collate_fn=collate_fn, pin_memory=False)
89 |
90 | # self.logger.info((str(self.net)))
91 | pytorch_total_params = sum(p.numel() for p in self.net.parameters() if p.requires_grad)
92 | self.logger.info("Number of parameters: {} ".format(pytorch_total_params / 1000000) + "M")
93 |
94 | # Load module
95 | self.highest_val_iou = 0
96 | self.start_epoch = 0
97 | CHECKPOINT_PATH = FLAGS.checkpoint_path
98 | if CHECKPOINT_PATH is not None and os.path.isfile(CHECKPOINT_PATH):
99 | self.logger.info("Load Pretrain from " + CHECKPOINT_PATH)
100 | checkpoint = torch.load(CHECKPOINT_PATH)
101 | try:
102 | self.logger.info("Loading strict=True")
103 | self.net.load_state_dict(checkpoint['model_state_dict'], strict=True)
104 | except:
105 | self.logger.info("Loading strict=False")
106 | self.net.load_state_dict(checkpoint['model_state_dict'], strict=False)
107 | self.start_epoch = checkpoint['epoch']
108 |
109 | # Loss Function
110 | self.val_dataset = val_dataset
111 |
112 | self.evaluator = iouEval(20, self.device, 0)
113 | self.evaluator2 = iouEval(20, self.device, 0)
114 |
115 | def train(self):
116 | self.net.eval() # set model to eval mode (for bn and dp)
117 | # iou_calc_polar = IoUCalculator(cfg)
118 | # iou_calc_random = IoUCalculator(cfg)
119 | self.evaluator.reset()
120 | self.evaluator2.reset()
121 | tqdm_loader = tqdm(self.val_loader, total=len(self.val_loader))
122 | with torch.no_grad():
123 | cnn = []
124 | cnn2 = []
125 | for batch_idx, (polar_data, random_data, idx) in enumerate(tqdm_loader):
126 | for key in polar_data:
127 | if type(polar_data[key]) is list:
128 | for i in range(cfg.num_layers):
129 | polar_data[key][i] = polar_data[key][i].cuda(non_blocking=True)
130 | else:
131 | polar_data[key] = polar_data[key].cuda(non_blocking=True)
132 | for key in random_data:
133 | if type(random_data[key]) is list:
134 | for i in range(cfg.num_layers):
135 | random_data[key][i] = random_data[key][i].cuda(non_blocking=True)
136 | else:
137 | random_data[key] = random_data[key].cuda(non_blocking=True)
138 | import time
139 | st = time.time()
140 | semantic_out = self.net(polar_data)
141 | torch.cuda.synchronize()
142 | ed = time.time() - st
143 | cnn.append(ed)
144 | print(ed)
145 | argmax = F.softmax(semantic_out, dim=1).argmax(dim=1)
146 | self.evaluator.addBatch(argmax, polar_data['labels'])
147 |
148 | # end_points = self.get_pred(semantic_out, polar_data, self.val_dataset)
149 | # acc, end_points = compute_acc(end_points)
150 | # iou_calc_polar.add_data(end_points)
151 | st = time.time()
152 | semantic_out = self.net(random_data)
153 | torch.cuda.synchronize()
154 | ed = time.time() - st
155 | cnn2.append(ed)
156 | print(ed)
157 | argmax = F.softmax(semantic_out, dim=1).argmax(dim=1)
158 | self.evaluator2.addBatch(argmax, random_data['labels'])
159 | # end_points = self.get_pred(semantic_out, random_data, self.val_dataset)
160 | # acc, end_points = compute_acc(end_points)
161 | # iou_calc_random.add_data(end_points)
162 | if batch_idx==50:
163 | print("Mean CNN inference time:{}\t std:{}".format(np.mean(cnn), np.std(cnn)))
164 | print("Mean KNN inference time:{}\t std:{}".format(np.mean(cnn2), np.std(cnn2)))
165 | exit()
166 | # mean_iou, iou_list = iou_calc_polar.compute_iou()
167 | mean_iou, iou_list = self.evaluator.getIoU()
168 | self.logger.info('mean IoU:{:.2f}'.format(mean_iou * 100))
169 | s = 'IoU:'
170 | for iou_tmp in iou_list:
171 | s += '{:5.2f} '.format(100 * iou_tmp)
172 | self.logger.info(s)
173 |
174 | # mean_iou, iou_list = iou_calc_random.compute_iou()
175 | mean_iou, iou_list = self.evaluator2.getIoU()
176 | self.logger.info('mean IoU:{:.2f}'.format(mean_iou * 100))
177 | s = 'IoU:'
178 | for iou_tmp in iou_list:
179 | s += '{:5.2f} '.format(100 * iou_tmp)
180 | self.logger.info(s)
181 |
182 | @staticmethod
183 | def get_pred(semantic_out, end_points, dataset):
184 | logits = semantic_out
185 | labels = end_points['labels']
186 | end_points = {}
187 | logits = logits.transpose(1, 2).reshape(-1, dataset.num_classes)
188 | labels = labels.reshape(-1)
189 |
190 | # Boolean mask of points that should be ignored
191 | ignored_bool = (labels == 0)
192 |
193 | for ign_label in dataset.ignored_labels:
194 | ignored_bool = ignored_bool | (labels == ign_label)
195 |
196 | # Collect logits and labels that are not ignored
197 | valid_idx = ignored_bool == 0
198 | valid_logits = logits[valid_idx, :]
199 | valid_labels_init = labels[valid_idx]
200 | # Reduce label values in the range of logit shape
201 | reducing_list = torch.arange(0, dataset.num_classes).long().to(logits.device)
202 | inserted_value = torch.zeros((1,)).long().to(logits.device)
203 | for ign_label in dataset.ignored_labels:
204 | reducing_list = torch.cat([reducing_list[:ign_label], inserted_value, reducing_list[ign_label:]], 0)
205 | valid_labels = torch.gather(reducing_list, 0, valid_labels_init)
206 | end_points['valid_logits'], end_points['valid_labels'] = valid_logits, valid_labels
207 | return end_points
208 | def main():
209 | trainer = Trainer()
210 | trainer.train()
211 |
212 |
213 | if __name__ == '__main__':
214 | main()
215 |
--------------------------------------------------------------------------------
/tool/main_poss.py:
--------------------------------------------------------------------------------
1 | import yaml
2 | import os
3 | import torch
4 | import time
5 | import numpy as np
6 | import torch.nn as nn
7 | import shutil
8 | from tqdm import tqdm
9 |
10 | EXTENSIONS_SCAN = ['.npy']
11 | def is_scan(filename):
12 | return any(filename.endswith(ext) for ext in EXTENSIONS_SCAN)
13 |
14 | root = '/home/chx/Work/SemanticPOSS_dataset_grid_0.01/sequences' # dataset
15 | DATA = yaml.safe_load(open("../utils/semantic-poss.yaml", 'r'))
16 |
17 | lidar_list = []
18 |
19 | sequences = DATA["split"]["train"]
20 | for seq in sequences:
21 | # to string
22 | seq = '{0:02d}'.format(int(seq))
23 | print("parsing seq {}".format(seq))
24 |
25 | # get paths for each
26 | scan_path = os.path.join(root, seq, "velodyne")
27 | # get files
28 | print(scan_path)
29 | scan_files = [os.path.join(dp, f) for dp, dn, fn in os.walk(
30 | os.path.expanduser(scan_path)) for f in fn if is_scan(f)]
31 | # extend list
32 |
33 |
34 | lidar_list.extend(scan_files)
35 |
36 | # sort for correspondance
37 | lidar_list.sort()
38 |
39 | # bound = np.arange(0, 20)
40 | bound = np.arange(0, 200)
41 | DISTANCES = np.append(bound, [200, 10**3])
42 | numbers = [0]*(len(DISTANCES)-1)
43 | print(DISTANCES, len(DISTANCES))
44 | for i in tqdm(range(len(lidar_list))):
45 | filename = lidar_list[i]
46 | # print(filename)
47 | scan = np.load(filename)
48 | scan = scan.reshape((-1, 3))
49 | points = scan[:, 0:3] # get xyz
50 | depth = np.linalg.norm(points, 2, axis=1)
51 |
52 | for idx in range(len(DISTANCES)-1):
53 | # select by range
54 | lrange = DISTANCES[idx]
55 | hrange = DISTANCES[idx+1]
56 |
57 | # print(lrange, hrange)
58 | mask = np.logical_and(depth > lrange, depth <= hrange)
59 | nums = len(depth[mask])
60 | numbers[idx] += nums
61 | np.save("d.npy", numbers)
62 | print(numbers)
63 |
64 |
--------------------------------------------------------------------------------
/train_SemanticPOSS.py:
--------------------------------------------------------------------------------
1 | # Common
2 | import os
3 | import logging
4 | import warnings
5 | import argparse
6 | import numpy as np
7 | from tqdm import tqdm
8 | # torch
9 | import torch
10 | import torch.nn as nn
11 | import torch.optim as optim
12 | from torch.utils.data import DataLoader
13 | # my module
14 | from dataset.poss_trainset import SemanticPOSS
15 |
16 | import torch.nn.functional as F
17 | from network.loss_func import compute_loss
18 | from utils.metric import compute_acc, IoUCalculator, iouEval
19 | from help_utils import seed_torch, my_worker_init_fn, get_logger, copyFiles, AverageMeter
20 |
21 | # os.environ["CUDA_VISIBLE_DEVICES"] = '1'
22 | # warnings.filterwarnings("ignore")
23 |
24 | parser = argparse.ArgumentParser()
25 | parser.add_argument('--backbone', type=str, default='randla', choices=['randla', 'baflac', 'baaf'])
26 | parser.add_argument('--checkpoint_path', default=None, help='Model checkpoint path [default: None]')
27 | parser.add_argument('--log_dir', type=str, default='base', help='Dump dir to save model checkpoint [default: log]')
28 | parser.add_argument('--max_epoch', type=int, default=100, help='Epoch to run [default: 100]')
29 | parser.add_argument('--batch_size', type=int, default=6, help='Batch Size during training [default: 5]')
30 | parser.add_argument('--val_batch_size', type=int, default=8, help='Batch Size during training [default: 30]')
31 | parser.add_argument('--num_workers', type=int, default=6, help='Number of workers [default: 5]')
32 | parser.add_argument('--sampling', type=str, default='random', choices=['random', 'polar'], help='Polar sample or not')
33 | parser.add_argument('--seed', type=int, default=1024, help='Polar sample or not')
34 | parser.add_argument('--grid', nargs='+', type=int, default=[64, 64, 16], help='grid size of BEV representation')
35 | FLAGS = parser.parse_args()
36 |
37 | seed_torch(FLAGS.seed)
38 | torch.backends.cudnn.enabled = False
39 |
40 | if FLAGS.backbone == 'baflac':
41 | from config import ConfigSemanticPOSS_BAF as cfg
42 | else:
43 | from config import ConfigSemanticPOSS as cfg
44 |
45 | class Trainer:
46 | def __init__(self):
47 | # Init Logging
48 | save_path = './save_semantic_poss/' + FLAGS.log_dir + '/'
49 | if not (os.path.exists(save_path)):
50 | os.makedirs(save_path)
51 | copyFiles(save_path)
52 | self.log_dir = save_path
53 | log_fname = os.path.join(self.log_dir, 'log_train.txt')
54 | self.logger = get_logger(log_fname, name="Trainer")
55 |
56 | argsDict = FLAGS.__dict__
57 | for eachArg, value in argsDict.items():
58 | self.logger.info(eachArg + ' : ' + str(value))
59 |
60 | train_dataset = SemanticPOSS('training', sampling_way=FLAGS.sampling, grid=FLAGS.grid)
61 | val_dataset = SemanticPOSS('validation', sampling_way=FLAGS.sampling, grid=FLAGS.grid)
62 |
63 | # Network & Optimizer
64 | self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
65 | if FLAGS.backbone == 'baflac':
66 | from network.BAF_LAC import BAF_LAC
67 | self.logger.info("Use Baseline: BAF-LAC")
68 | self.net = BAF_LAC(cfg)
69 | self.net.to(self.device)
70 | collate_fn = train_dataset.collate_fn_baf_lac
71 |
72 | elif FLAGS.backbone == 'randla':
73 | from network.RandLANet import Network
74 | self.logger.info("Use Baseline: Rand-LA")
75 | self.net = Network(cfg)
76 | self.net.to(self.device)
77 | collate_fn = train_dataset.collate_fn
78 |
79 | elif FLAGS.backbone == 'baaf':
80 | from network.BAAF import Network
81 | self.logger.info("Use Baseline: BAAF")
82 | self.net = Network(cfg)
83 | self.net.to(self.device)
84 | collate_fn = train_dataset.collate_fn
85 |
86 | else:
87 | raise TypeError("1~5~!! can can need !!!")
88 |
89 | self.train_loader = DataLoader(
90 | train_dataset, batch_size=FLAGS.batch_size,
91 | shuffle=True, num_workers=FLAGS.num_workers, pin_memory=False,
92 | worker_init_fn=my_worker_init_fn, collate_fn=collate_fn)
93 |
94 | self.val_loader = DataLoader(
95 | val_dataset, batch_size=FLAGS.val_batch_size,
96 | shuffle=False, num_workers=FLAGS.num_workers, pin_memory=False,
97 | worker_init_fn=my_worker_init_fn, collate_fn=collate_fn)
98 |
99 |
100 | self.logger.info((str(self.net)))
101 | pytorch_total_params = sum(p.numel() for p in self.net.parameters() if p.requires_grad)
102 | self.logger.info("Number of parameters: {} ".format(pytorch_total_params / 1000000) + "M")
103 |
104 | # Load the Adam optimizer
105 | self.optimizer = optim.Adam(self.net.parameters(), lr=0.01)
106 | self.scheduler = optim.lr_scheduler.ExponentialLR(self.optimizer, 0.95)
107 |
108 | # Load module
109 | self.highest_val_iou = 0
110 | self.start_epoch = 0
111 | CHECKPOINT_PATH = FLAGS.checkpoint_path
112 | if CHECKPOINT_PATH is not None and os.path.isfile(CHECKPOINT_PATH):
113 | self.logger.info("Load Pretrain")
114 | checkpoint = torch.load(CHECKPOINT_PATH)
115 | self.net.load_state_dict(checkpoint['model_state_dict'], strict=True)
116 | self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
117 | self.scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
118 | self.start_epoch = checkpoint['epoch']
119 |
120 | # Loss Function
121 | # class_weights = torch.tensor([[24.88170566, 41.76906, 8.30893469, 32.16651947, 2.53649364, 39.0391128,
122 | # 40.95852351, 4.19310616, 28.07720036, 12.26610562, 4.45008737]]).cuda()
123 |
124 | class_weights = train_dataset.get_class_weight()
125 | self.logger.info(class_weights)
126 | class_weights = torch.from_numpy(class_weights).float().cuda()
127 | self.criterion = nn.CrossEntropyLoss(weight=class_weights, reduction='none')
128 |
129 | self.evaluator = iouEval(12, self.device, 0)
130 | self.train_dataset = train_dataset
131 | self.val_dataset = val_dataset
132 |
133 | def train_one_epoch(self):
134 | self.net.train() # set model to training mode
135 | losses = AverageMeter()
136 | tqdm_loader = tqdm(self.train_loader, total=len(self.train_loader))
137 |
138 | scaler = torch.cuda.amp.GradScaler()
139 | for batch_idx, batch_data in enumerate(tqdm_loader):
140 | for key in batch_data:
141 | if type(batch_data[key]) is list:
142 | for i in range(cfg.num_layers):
143 | batch_data[key][i] = batch_data[key][i].cuda(non_blocking=True)
144 | else:
145 | batch_data[key] = batch_data[key].cuda(non_blocking=True)
146 |
147 | self.optimizer.zero_grad()
148 |
149 | with torch.cuda.amp.autocast():
150 | semantic_out = self.net(batch_data)
151 | # loss = self.criterion(semantic_out, batch_data['labels']).mean()
152 | loss, end_points = compute_loss(semantic_out, batch_data, self.train_dataset, self.criterion)
153 |
154 | scaler.scale(loss).backward()
155 | scaler.step(self.optimizer)
156 | scaler.update()
157 |
158 | losses.update(loss.item())
159 | if batch_idx % 50 == 0:
160 | lr = self.optimizer.param_groups[0]['lr']
161 | self.logger.info('Step {:08d} || Lr={:.6f} || L_out={loss.val:.4f}/({loss.avg:.4f})'.format(batch_idx, lr, loss=losses))
162 | self.scheduler.step()
163 |
164 | def train(self):
165 | for epoch in range(self.start_epoch, FLAGS.max_epoch):
166 | self.cur_epoch = epoch
167 | self.logger.info('**** EPOCH %03d ****' % (epoch))
168 |
169 | self.train_one_epoch()
170 | self.logger.info('**** EVAL EPOCH %03d ****' % (epoch))
171 | checkpoint_file = os.path.join(self.log_dir, 'checkpoint.tar')
172 | self.save_checkpoint(checkpoint_file)
173 |
174 | mean_iou = self.validate()
175 | # Save best checkpoint
176 | if mean_iou > self.highest_val_iou:
177 | self.logger.info('**** Current: %03f Best: %03f ****' % (mean_iou, self.highest_val_iou))
178 | self.highest_val_iou = mean_iou
179 | checkpoint_file = os.path.join(self.log_dir, 'checkpoint-best.tar')
180 | self.save_checkpoint(checkpoint_file)
181 | else:
182 | self.logger.info('**** Current: %03f Best: %03f ****' % (mean_iou, self.highest_val_iou))
183 |
184 |
185 | def validate(self):
186 | # torch.cuda.empty_cache()
187 | self.net.eval() # set model to eval mode (for bn and dp)
188 |
189 | # self.evaluator.reset()
190 | iou_calc = IoUCalculator(cfg)
191 |
192 | tqdm_loader = tqdm(self.val_loader, total=len(self.val_loader))
193 | with torch.no_grad():
194 | for batch_idx, batch_data in enumerate(tqdm_loader):
195 | for key in batch_data:
196 | if type(batch_data[key]) is list:
197 | for i in range(cfg.num_layers):
198 | batch_data[key][i] = batch_data[key][i].cuda(non_blocking=True)
199 | else:
200 | batch_data[key] = batch_data[key].cuda(non_blocking=True)
201 |
202 | # Forward pass
203 | # torch.cuda.synchronize()
204 | semantic_out = self.net(batch_data)
205 | # argmax = F.softmax(semantic_out, dim=1).argmax(dim=1)
206 | # self.evaluator.addBatch(argmax, batch_data['labels'])
207 |
208 | loss, end_points = compute_loss(semantic_out, batch_data, self.train_dataset, self.criterion)
209 | acc, end_points = compute_acc(end_points)
210 | iou_calc.add_data(end_points)
211 |
212 | # mean_iou, iou_list = self.evaluator.getIoU()
213 | mean_iou, iou_list = iou_calc.compute_iou()
214 | self.logger.info('mean IoU:{:.1f}'.format(mean_iou * 100))
215 | s = 'IoU:'
216 | for iou_tmp in iou_list:
217 | s += '{:5.2f} '.format(100 * iou_tmp)
218 | self.logger.info(s)
219 |
220 | return mean_iou
221 |
222 | def save_checkpoint(self, fname):
223 | save_dict = {
224 | 'epoch': self.cur_epoch+1, # after training one epoch, the start_epoch should be epoch+1
225 | 'optimizer_state_dict': self.optimizer.state_dict(),
226 | 'scheduler_state_dict': self.scheduler.state_dict()
227 | }
228 | # with nn.DataParallel() the net is added as a submodule of DataParallel
229 | try:
230 | save_dict['model_state_dict'] = self.net.module.state_dict()
231 | except AttributeError:
232 | save_dict['model_state_dict'] = self.net.state_dict()
233 | torch.save(save_dict, fname)
234 |
235 |
236 | def main():
237 | trainer = Trainer()
238 | trainer.train()
239 |
240 |
241 | if __name__ == '__main__':
242 | main()
243 |
--------------------------------------------------------------------------------
/utils/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/huixiancheng/PCB-RandNet/304c84b67380a23ab66516581b348c7db2f9f0df/utils/__init__.py
--------------------------------------------------------------------------------
/utils/config.py:
--------------------------------------------------------------------------------
1 | class ConfigSemanticKITTI:
2 | k_n = [25, 16, 16, 16] # KNN
3 | num_layers = 4 # Number of layers
4 | num_points = 4096 * 11 # Number of input points
5 | num_classes = 19 # Number of valid classes
6 | sub_grid_size = 0.06 # preprocess_parameter
7 |
8 | batch_size = 6 # batch_size during training
9 | val_batch_size = 20 # batch_size during validation and test
10 | train_steps = 500 # Number of steps per epochs
11 | val_steps = 100 # Number of validation steps per epoch
12 |
13 | sub_sampling_ratio = [4, 4, 4, 4] # sampling ratio of random sampling at each layer
14 | d_out = [16, 64, 128, 256] # feature dimension
15 | num_sub_points = [num_points // 4, num_points // 16, num_points // 64, num_points // 256]
16 |
17 | noise_init = 3.5 # noise initial parameter
18 | max_epoch = 100 # maximum epoch during training
19 | learning_rate = 1e-2 # initial learning rate
20 | lr_decays = {i: 0.95 for i in range(0, 500)} # decay rate of learning rate
21 |
22 | train_sum_dir = 'train_log'
23 | saving = True
24 | saving_path = None
25 |
26 |
27 | class ConfigS3DIS:
28 | k_n = 16 # KNN
29 | num_layers = 5 # Number of layers
30 | num_points = 40960 # Number of input points
31 | num_classes = 13 # Number of valid classes
32 | sub_grid_size = 0.04 # preprocess_parameter
33 |
34 | batch_size = 6 # batch_size during training
35 | val_batch_size = 20 # batch_size during validation and test
36 | train_steps = 500 # Number of steps per epochs
37 | val_steps = 100 # Number of validation steps per epoch
38 |
39 | sub_sampling_ratio = [4, 4, 4, 4, 2] # sampling ratio of random sampling at each layer
40 | d_out = [16, 64, 128, 256, 512] # feature dimension
41 |
42 | noise_init = 3.5 # noise initial parameter
43 | max_epoch = 100 # maximum epoch during training
44 | learning_rate = 1e-2 # initial learning rate
45 | lr_decays = {i: 0.95 for i in range(0, 500)} # decay rate of learning rate
46 |
47 | train_sum_dir = 'train_log'
48 | saving = True
49 | saving_path = None
50 |
51 |
52 | class ConfigSemantic3D:
53 | k_n = 16 # KNN
54 | num_layers = 5 # Number of layers
55 | num_points = 65536 # Number of input points
56 | num_classes = 8 # Number of valid classes
57 | sub_grid_size = 0.06 # preprocess_parameter
58 |
59 | batch_size = 4 # batch_size during training
60 | val_batch_size = 16 # batch_size during validation and test
61 | train_steps = 500 # Number of steps per epochs
62 | val_steps = 100 # Number of validation steps per epoch
63 |
64 | sub_sampling_ratio = [4, 4, 4, 4, 2] # sampling ratio of random sampling at each layer
65 | d_out = [16, 64, 128, 256, 512] # feature dimension
66 |
67 | noise_init = 3.5 # noise initial parameter
68 | max_epoch = 100 # maximum epoch during training
69 | learning_rate = 1e-2 # initial learning rate
70 | lr_decays = {i: 0.95 for i in range(0, 500)} # decay rate of learning rate
71 |
72 | train_sum_dir = 'train_log'
73 | saving = True
74 | saving_path = None
75 |
76 | augment_scale_anisotropic = True
77 | augment_symmetries = [True, False, False]
78 | augment_rotation = 'vertical'
79 | augment_scale_min = 0.8
80 | augment_scale_max = 1.2
81 | augment_noise = 0.001
82 | augment_occlusion = 'none'
83 | augment_color = 0.8
84 |
--------------------------------------------------------------------------------
/utils/cpp_wrappers/compile_wrappers.sh:
--------------------------------------------------------------------------------
1 | #!/bin/bash
2 |
3 | # Compile cpp subsampling
4 | cd cpp_subsampling
5 | python3 setup.py build_ext --inplace
6 | cd ..
7 |
8 |
--------------------------------------------------------------------------------
/utils/cpp_wrappers/cpp_subsampling/grid_subsampling/grid_subsampling.cpp:
--------------------------------------------------------------------------------
1 |
2 | #include "grid_subsampling.h"
3 |
4 |
5 | void grid_subsampling(vector& original_points,
6 | vector& subsampled_points,
7 | vector& original_features,
8 | vector& subsampled_features,
9 | vector& original_classes,
10 | vector& subsampled_classes,
11 | float sampleDl,
12 | int verbose) {
13 |
14 | // Initiate variables
15 | // ******************
16 |
17 | // Number of points in the cloud
18 | size_t N = original_points.size();
19 |
20 | // Dimension of the features
21 | size_t fdim = original_features.size() / N;
22 | size_t ldim = original_classes.size() / N;
23 |
24 | // Limits of the cloud
25 | PointXYZ minCorner = min_point(original_points);
26 | PointXYZ maxCorner = max_point(original_points);
27 | PointXYZ originCorner = floor(minCorner * (1/sampleDl)) * sampleDl;
28 |
29 | // Dimensions of the grid
30 | size_t sampleNX = (size_t)floor((maxCorner.x - originCorner.x) / sampleDl) + 1;
31 | size_t sampleNY = (size_t)floor((maxCorner.y - originCorner.y) / sampleDl) + 1;
32 | //size_t sampleNZ = (size_t)floor((maxCorner.z - originCorner.z) / sampleDl) + 1;
33 |
34 | // Check if features and classes need to be processed
35 | bool use_feature = original_features.size() > 0;
36 | bool use_classes = original_classes.size() > 0;
37 |
38 |
39 | // Create the sampled map
40 | // **********************
41 |
42 | // Verbose parameters
43 | int i = 0;
44 | int nDisp = N / 100;
45 |
46 | // Initiate variables
47 | size_t iX, iY, iZ, mapIdx;
48 | unordered_map data;
49 |
50 | for (auto& p : original_points)
51 | {
52 | // Position of point in sample map
53 | iX = (size_t)floor((p.x - originCorner.x) / sampleDl);
54 | iY = (size_t)floor((p.y - originCorner.y) / sampleDl);
55 | iZ = (size_t)floor((p.z - originCorner.z) / sampleDl);
56 | mapIdx = iX + sampleNX*iY + sampleNX*sampleNY*iZ;
57 |
58 | // If not already created, create key
59 | if (data.count(mapIdx) < 1)
60 | data.emplace(mapIdx, SampledData(fdim, ldim));
61 |
62 | // Fill the sample map
63 | if (use_feature && use_classes)
64 | data[mapIdx].update_all(p, original_features.begin() + i * fdim, original_classes.begin() + i * ldim);
65 | else if (use_feature)
66 | data[mapIdx].update_features(p, original_features.begin() + i * fdim);
67 | else if (use_classes)
68 | data[mapIdx].update_classes(p, original_classes.begin() + i * ldim);
69 | else
70 | data[mapIdx].update_points(p);
71 |
72 | // Display
73 | i++;
74 | if (verbose > 1 && i%nDisp == 0)
75 | std::cout << "\rSampled Map : " << std::setw(3) << i / nDisp << "%";
76 |
77 | }
78 |
79 | // Divide for barycentre and transfer to a vector
80 | subsampled_points.reserve(data.size());
81 | if (use_feature)
82 | subsampled_features.reserve(data.size() * fdim);
83 | if (use_classes)
84 | subsampled_classes.reserve(data.size() * ldim);
85 | for (auto& v : data)
86 | {
87 | subsampled_points.push_back(v.second.point * (1.0 / v.second.count));
88 | if (use_feature)
89 | {
90 | float count = (float)v.second.count;
91 | transform(v.second.features.begin(),
92 | v.second.features.end(),
93 | v.second.features.begin(),
94 | [count](float f) { return f / count;});
95 | subsampled_features.insert(subsampled_features.end(),v.second.features.begin(),v.second.features.end());
96 | }
97 | if (use_classes)
98 | {
99 | for (int i = 0; i < ldim; i++)
100 | subsampled_classes.push_back(max_element(v.second.labels[i].begin(), v.second.labels[i].end(),
101 | [](const pair&a, const pair&b){return a.second < b.second;})->first);
102 | }
103 | }
104 |
105 | return;
106 | }
107 |
--------------------------------------------------------------------------------
/utils/cpp_wrappers/cpp_subsampling/grid_subsampling/grid_subsampling.h:
--------------------------------------------------------------------------------
1 |
2 |
3 | #include "../../cpp_utils/cloud/cloud.h"
4 |
5 | #include
6 | #include
7 |
8 | using namespace std;
9 |
10 | class SampledData
11 | {
12 | public:
13 |
14 | // Elements
15 | // ********
16 |
17 | int count;
18 | PointXYZ point;
19 | vector features;
20 | vector> labels;
21 |
22 |
23 | // Methods
24 | // *******
25 |
26 | // Constructor
27 | SampledData()
28 | {
29 | count = 0;
30 | point = PointXYZ();
31 | }
32 |
33 | SampledData(const size_t fdim, const size_t ldim)
34 | {
35 | count = 0;
36 | point = PointXYZ();
37 | features = vector(fdim);
38 | labels = vector>(ldim);
39 | }
40 |
41 | // Method Update
42 | void update_all(const PointXYZ p, vector::iterator f_begin, vector::iterator l_begin)
43 | {
44 | count += 1;
45 | point += p;
46 | transform (features.begin(), features.end(), f_begin, features.begin(), plus());
47 | int i = 0;
48 | for(vector::iterator it = l_begin; it != l_begin + labels.size(); ++it)
49 | {
50 | labels[i][*it] += 1;
51 | i++;
52 | }
53 | return;
54 | }
55 | void update_features(const PointXYZ p, vector::iterator f_begin)
56 | {
57 | count += 1;
58 | point += p;
59 | transform (features.begin(), features.end(), f_begin, features.begin(), plus());
60 | return;
61 | }
62 | void update_classes(const PointXYZ p, vector::iterator l_begin)
63 | {
64 | count += 1;
65 | point += p;
66 | int i = 0;
67 | for(vector::iterator it = l_begin; it != l_begin + labels.size(); ++it)
68 | {
69 | labels[i][*it] += 1;
70 | i++;
71 | }
72 | return;
73 | }
74 | void update_points(const PointXYZ p)
75 | {
76 | count += 1;
77 | point += p;
78 | return;
79 | }
80 | };
81 |
82 |
83 |
84 | void grid_subsampling(vector& original_points,
85 | vector& subsampled_points,
86 | vector& original_features,
87 | vector& subsampled_features,
88 | vector& original_classes,
89 | vector& subsampled_classes,
90 | float sampleDl,
91 | int verbose);
92 |
93 |
--------------------------------------------------------------------------------
/utils/cpp_wrappers/cpp_subsampling/setup.py:
--------------------------------------------------------------------------------
1 | from distutils.core import setup, Extension
2 | import numpy.distutils.misc_util
3 |
4 | # Adding OpenCV to project
5 | # ************************
6 |
7 | # Adding sources of the project
8 | # *****************************
9 |
10 | m_name = "grid_subsampling"
11 |
12 | SOURCES = ["../cpp_utils/cloud/cloud.cpp",
13 | "grid_subsampling/grid_subsampling.cpp",
14 | "wrapper.cpp"]
15 |
16 | module = Extension(m_name,
17 | sources=SOURCES,
18 | extra_compile_args=['-std=c++11',
19 | '-D_GLIBCXX_USE_CXX11_ABI=0'])
20 |
21 | setup(ext_modules=[module], include_dirs=numpy.distutils.misc_util.get_numpy_include_dirs())
22 |
23 |
24 |
25 |
26 |
27 |
28 |
29 |
30 |
--------------------------------------------------------------------------------
/utils/cpp_wrappers/cpp_subsampling/wrapper.cpp:
--------------------------------------------------------------------------------
1 | #include
2 | #include
3 | #include "grid_subsampling/grid_subsampling.h"
4 | #include
5 |
6 |
7 |
8 | // docstrings for our module
9 | // *************************
10 |
11 | static char module_docstring[] = "This module provides an interface for the subsampling of a pointcloud";
12 |
13 | static char compute_docstring[] = "function subsampling a pointcloud";
14 |
15 |
16 | // Declare the functions
17 | // *********************
18 |
19 | static PyObject *grid_subsampling_compute(PyObject *self, PyObject *args, PyObject *keywds);
20 |
21 |
22 | // Specify the members of the module
23 | // *********************************
24 |
25 | static PyMethodDef module_methods[] =
26 | {
27 | { "compute", (PyCFunction)grid_subsampling_compute, METH_VARARGS | METH_KEYWORDS, compute_docstring },
28 | {NULL, NULL, 0, NULL}
29 | };
30 |
31 |
32 | // Initialize the module
33 | // *********************
34 |
35 | static struct PyModuleDef moduledef =
36 | {
37 | PyModuleDef_HEAD_INIT,
38 | "grid_subsampling", // m_name
39 | module_docstring, // m_doc
40 | -1, // m_size
41 | module_methods, // m_methods
42 | NULL, // m_reload
43 | NULL, // m_traverse
44 | NULL, // m_clear
45 | NULL, // m_free
46 | };
47 |
48 | PyMODINIT_FUNC PyInit_grid_subsampling(void)
49 | {
50 | import_array();
51 | return PyModule_Create(&moduledef);
52 | }
53 |
54 |
55 | // Actual wrapper
56 | // **************
57 |
58 | static PyObject *grid_subsampling_compute(PyObject *self, PyObject *args, PyObject *keywds)
59 | {
60 |
61 | // Manage inputs
62 | // *************
63 |
64 | // Args containers
65 | PyObject *points_obj = NULL;
66 | PyObject *features_obj = NULL;
67 | PyObject *classes_obj = NULL;
68 |
69 | // Keywords containers
70 | static char *kwlist[] = {"points", "features", "classes", "sampleDl", "method", "verbose", NULL };
71 | float sampleDl = 0.1;
72 | const char *method_buffer = "barycenters";
73 | int verbose = 0;
74 |
75 | // Parse the input
76 | if (!PyArg_ParseTupleAndKeywords(args, keywds, "O|$OOfsi", kwlist, &points_obj, &features_obj, &classes_obj, &sampleDl, &method_buffer, &verbose))
77 | {
78 | PyErr_SetString(PyExc_RuntimeError, "Error parsing arguments");
79 | return NULL;
80 | }
81 |
82 | // Get the method argument
83 | string method(method_buffer);
84 |
85 | // Interpret method
86 | if (method.compare("barycenters") && method.compare("voxelcenters"))
87 | {
88 | PyErr_SetString(PyExc_RuntimeError, "Error parsing method. Valid method names are \"barycenters\" and \"voxelcenters\" ");
89 | return NULL;
90 | }
91 |
92 | // Check if using features or classes
93 | bool use_feature = true, use_classes = true;
94 | if (features_obj == NULL)
95 | use_feature = false;
96 | if (classes_obj == NULL)
97 | use_classes = false;
98 |
99 | // Interpret the input objects as numpy arrays.
100 | PyObject *points_array = PyArray_FROM_OTF(points_obj, NPY_FLOAT, NPY_IN_ARRAY);
101 | PyObject *features_array = NULL;
102 | PyObject *classes_array = NULL;
103 | if (use_feature)
104 | features_array = PyArray_FROM_OTF(features_obj, NPY_FLOAT, NPY_IN_ARRAY);
105 | if (use_classes)
106 | classes_array = PyArray_FROM_OTF(classes_obj, NPY_INT, NPY_IN_ARRAY);
107 |
108 | // Verify data was load correctly.
109 | if (points_array == NULL)
110 | {
111 | Py_XDECREF(points_array);
112 | Py_XDECREF(classes_array);
113 | Py_XDECREF(features_array);
114 | PyErr_SetString(PyExc_RuntimeError, "Error converting input points to numpy arrays of type float32");
115 | return NULL;
116 | }
117 | if (use_feature && features_array == NULL)
118 | {
119 | Py_XDECREF(points_array);
120 | Py_XDECREF(classes_array);
121 | Py_XDECREF(features_array);
122 | PyErr_SetString(PyExc_RuntimeError, "Error converting input features to numpy arrays of type float32");
123 | return NULL;
124 | }
125 | if (use_classes && classes_array == NULL)
126 | {
127 | Py_XDECREF(points_array);
128 | Py_XDECREF(classes_array);
129 | Py_XDECREF(features_array);
130 | PyErr_SetString(PyExc_RuntimeError, "Error converting input classes to numpy arrays of type int32");
131 | return NULL;
132 | }
133 |
134 | // Check that the input array respect the dims
135 | if ((int)PyArray_NDIM(points_array) != 2 || (int)PyArray_DIM(points_array, 1) != 3)
136 | {
137 | Py_XDECREF(points_array);
138 | Py_XDECREF(classes_array);
139 | Py_XDECREF(features_array);
140 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : points.shape is not (N, 3)");
141 | return NULL;
142 | }
143 | if (use_feature && ((int)PyArray_NDIM(features_array) != 2))
144 | {
145 | Py_XDECREF(points_array);
146 | Py_XDECREF(classes_array);
147 | Py_XDECREF(features_array);
148 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : features.shape is not (N, d)");
149 | return NULL;
150 | }
151 |
152 | if (use_classes && (int)PyArray_NDIM(classes_array) > 2)
153 | {
154 | Py_XDECREF(points_array);
155 | Py_XDECREF(classes_array);
156 | Py_XDECREF(features_array);
157 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : classes.shape is not (N,) or (N, d)");
158 | return NULL;
159 | }
160 |
161 | // Number of points
162 | int N = (int)PyArray_DIM(points_array, 0);
163 |
164 | // Dimension of the features
165 | int fdim = 0;
166 | if (use_feature)
167 | fdim = (int)PyArray_DIM(features_array, 1);
168 |
169 | //Dimension of labels
170 | int ldim = 1;
171 | if (use_classes && (int)PyArray_NDIM(classes_array) == 2)
172 | ldim = (int)PyArray_DIM(classes_array, 1);
173 |
174 | // Check that the input array respect the number of points
175 | if (use_feature && (int)PyArray_DIM(features_array, 0) != N)
176 | {
177 | Py_XDECREF(points_array);
178 | Py_XDECREF(classes_array);
179 | Py_XDECREF(features_array);
180 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : features.shape is not (N, d)");
181 | return NULL;
182 | }
183 | if (use_classes && (int)PyArray_DIM(classes_array, 0) != N)
184 | {
185 | Py_XDECREF(points_array);
186 | Py_XDECREF(classes_array);
187 | Py_XDECREF(features_array);
188 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : classes.shape is not (N,) or (N, d)");
189 | return NULL;
190 | }
191 |
192 |
193 | // Call the C++ function
194 | // *********************
195 |
196 | // Create pyramid
197 | if (verbose > 0)
198 | cout << "Computing cloud pyramid with support points: " << endl;
199 |
200 |
201 | // Convert PyArray to Cloud C++ class
202 | vector original_points;
203 | vector original_features;
204 | vector original_classes;
205 | original_points = vector((PointXYZ*)PyArray_DATA(points_array), (PointXYZ*)PyArray_DATA(points_array) + N);
206 | if (use_feature)
207 | original_features = vector((float*)PyArray_DATA(features_array), (float*)PyArray_DATA(features_array) + N*fdim);
208 | if (use_classes)
209 | original_classes = vector((int*)PyArray_DATA(classes_array), (int*)PyArray_DATA(classes_array) + N*ldim);
210 |
211 | // Subsample
212 | vector subsampled_points;
213 | vector subsampled_features;
214 | vector subsampled_classes;
215 | grid_subsampling(original_points,
216 | subsampled_points,
217 | original_features,
218 | subsampled_features,
219 | original_classes,
220 | subsampled_classes,
221 | sampleDl,
222 | verbose);
223 |
224 | // Check result
225 | if (subsampled_points.size() < 1)
226 | {
227 | PyErr_SetString(PyExc_RuntimeError, "Error");
228 | return NULL;
229 | }
230 |
231 | // Manage outputs
232 | // **************
233 |
234 | // Dimension of input containers
235 | npy_intp* point_dims = new npy_intp[2];
236 | point_dims[0] = subsampled_points.size();
237 | point_dims[1] = 3;
238 | npy_intp* feature_dims = new npy_intp[2];
239 | feature_dims[0] = subsampled_points.size();
240 | feature_dims[1] = fdim;
241 | npy_intp* classes_dims = new npy_intp[2];
242 | classes_dims[0] = subsampled_points.size();
243 | classes_dims[1] = ldim;
244 |
245 | // Create output array
246 | PyObject *res_points_obj = PyArray_SimpleNew(2, point_dims, NPY_FLOAT);
247 | PyObject *res_features_obj = NULL;
248 | PyObject *res_classes_obj = NULL;
249 | PyObject *ret = NULL;
250 |
251 | // Fill output array with values
252 | size_t size_in_bytes = subsampled_points.size() * 3 * sizeof(float);
253 | memcpy(PyArray_DATA(res_points_obj), subsampled_points.data(), size_in_bytes);
254 | if (use_feature)
255 | {
256 | size_in_bytes = subsampled_points.size() * fdim * sizeof(float);
257 | res_features_obj = PyArray_SimpleNew(2, feature_dims, NPY_FLOAT);
258 | memcpy(PyArray_DATA(res_features_obj), subsampled_features.data(), size_in_bytes);
259 | }
260 | if (use_classes)
261 | {
262 | size_in_bytes = subsampled_points.size() * ldim * sizeof(int);
263 | res_classes_obj = PyArray_SimpleNew(2, classes_dims, NPY_INT);
264 | memcpy(PyArray_DATA(res_classes_obj), subsampled_classes.data(), size_in_bytes);
265 | }
266 |
267 |
268 | // Merge results
269 | if (use_feature && use_classes)
270 | ret = Py_BuildValue("NNN", res_points_obj, res_features_obj, res_classes_obj);
271 | else if (use_feature)
272 | ret = Py_BuildValue("NN", res_points_obj, res_features_obj);
273 | else if (use_classes)
274 | ret = Py_BuildValue("NN", res_points_obj, res_classes_obj);
275 | else
276 | ret = Py_BuildValue("N", res_points_obj);
277 |
278 | // Clean up
279 | // ********
280 |
281 | Py_DECREF(points_array);
282 | Py_XDECREF(features_array);
283 | Py_XDECREF(classes_array);
284 |
285 | return ret;
286 | }
--------------------------------------------------------------------------------
/utils/cpp_wrappers/cpp_utils/cloud/cloud.cpp:
--------------------------------------------------------------------------------
1 | //
2 | //
3 | // 0==========================0
4 | // | Local feature test |
5 | // 0==========================0
6 | //
7 | // version 1.0 :
8 | // >
9 | //
10 | //---------------------------------------------------
11 | //
12 | // Cloud source :
13 | // Define usefull Functions/Methods
14 | //
15 | //----------------------------------------------------
16 | //
17 | // Hugues THOMAS - 10/02/2017
18 | //
19 |
20 |
21 | #include "cloud.h"
22 |
23 |
24 | // Getters
25 | // *******
26 |
27 | PointXYZ max_point(std::vector points)
28 | {
29 | // Initiate limits
30 | PointXYZ maxP(points[0]);
31 |
32 | // Loop over all points
33 | for (auto p : points)
34 | {
35 | if (p.x > maxP.x)
36 | maxP.x = p.x;
37 |
38 | if (p.y > maxP.y)
39 | maxP.y = p.y;
40 |
41 | if (p.z > maxP.z)
42 | maxP.z = p.z;
43 | }
44 |
45 | return maxP;
46 | }
47 |
48 | PointXYZ min_point(std::vector points)
49 | {
50 | // Initiate limits
51 | PointXYZ minP(points[0]);
52 |
53 | // Loop over all points
54 | for (auto p : points)
55 | {
56 | if (p.x < minP.x)
57 | minP.x = p.x;
58 |
59 | if (p.y < minP.y)
60 | minP.y = p.y;
61 |
62 | if (p.z < minP.z)
63 | minP.z = p.z;
64 | }
65 |
66 | return minP;
67 | }
--------------------------------------------------------------------------------
/utils/cpp_wrappers/cpp_utils/cloud/cloud.h:
--------------------------------------------------------------------------------
1 | //
2 | //
3 | // 0==========================0
4 | // | Local feature test |
5 | // 0==========================0
6 | //
7 | // version 1.0 :
8 | // >
9 | //
10 | //---------------------------------------------------
11 | //
12 | // Cloud header
13 | //
14 | //----------------------------------------------------
15 | //
16 | // Hugues THOMAS - 10/02/2017
17 | //
18 |
19 |
20 | # pragma once
21 |
22 | #include
23 | #include
24 | #include