├── .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 25 | #include 26 | #include 27 | #include 28 | #include 29 | #include 30 | 31 | #include 32 | 33 | 34 | 35 | 36 | // Point class 37 | // *********** 38 | 39 | 40 | class PointXYZ 41 | { 42 | public: 43 | 44 | // Elements 45 | // ******** 46 | 47 | float x, y, z; 48 | 49 | 50 | // Methods 51 | // ******* 52 | 53 | // Constructor 54 | PointXYZ() { x = 0; y = 0; z = 0; } 55 | PointXYZ(float x0, float y0, float z0) { x = x0; y = y0; z = z0; } 56 | 57 | // array type accessor 58 | float operator [] (int i) const 59 | { 60 | if (i == 0) return x; 61 | else if (i == 1) return y; 62 | else return z; 63 | } 64 | 65 | // opperations 66 | float dot(const PointXYZ P) const 67 | { 68 | return x * P.x + y * P.y + z * P.z; 69 | } 70 | 71 | float sq_norm() 72 | { 73 | return x*x + y*y + z*z; 74 | } 75 | 76 | PointXYZ cross(const PointXYZ P) const 77 | { 78 | return PointXYZ(y*P.z - z*P.y, z*P.x - x*P.z, x*P.y - y*P.x); 79 | } 80 | 81 | PointXYZ& operator+=(const PointXYZ& P) 82 | { 83 | x += P.x; 84 | y += P.y; 85 | z += P.z; 86 | return *this; 87 | } 88 | 89 | PointXYZ& operator-=(const PointXYZ& P) 90 | { 91 | x -= P.x; 92 | y -= P.y; 93 | z -= P.z; 94 | return *this; 95 | } 96 | 97 | PointXYZ& operator*=(const float& a) 98 | { 99 | x *= a; 100 | y *= a; 101 | z *= a; 102 | return *this; 103 | } 104 | }; 105 | 106 | 107 | // Point Opperations 108 | // ***************** 109 | 110 | inline PointXYZ operator + (const PointXYZ A, const PointXYZ B) 111 | { 112 | return PointXYZ(A.x + B.x, A.y + B.y, A.z + B.z); 113 | } 114 | 115 | inline PointXYZ operator - (const PointXYZ A, const PointXYZ B) 116 | { 117 | return PointXYZ(A.x - B.x, A.y - B.y, A.z - B.z); 118 | } 119 | 120 | inline PointXYZ operator * (const PointXYZ P, const float a) 121 | { 122 | return PointXYZ(P.x * a, P.y * a, P.z * a); 123 | } 124 | 125 | inline PointXYZ operator * (const float a, const PointXYZ P) 126 | { 127 | return PointXYZ(P.x * a, P.y * a, P.z * a); 128 | } 129 | 130 | inline std::ostream& operator << (std::ostream& os, const PointXYZ P) 131 | { 132 | return os << "[" << P.x << ", " << P.y << ", " << P.z << "]"; 133 | } 134 | 135 | inline bool operator == (const PointXYZ A, const PointXYZ B) 136 | { 137 | return A.x == B.x && A.y == B.y && A.z == B.z; 138 | } 139 | 140 | inline PointXYZ floor(const PointXYZ P) 141 | { 142 | return PointXYZ(std::floor(P.x), std::floor(P.y), std::floor(P.z)); 143 | } 144 | 145 | 146 | PointXYZ max_point(std::vector points); 147 | PointXYZ min_point(std::vector points); 148 | 149 | 150 | 151 | 152 | 153 | 154 | 155 | 156 | 157 | 158 | 159 | -------------------------------------------------------------------------------- /utils/data_process.py: -------------------------------------------------------------------------------- 1 | from os.path import join 2 | import numpy as np 3 | import os 4 | import json 5 | import pandas as pd 6 | from tqdm import tqdm 7 | import utils.cpp_wrappers.cpp_subsampling.grid_subsampling as cpp_subsampling 8 | import utils.nearest_neighbors.lib.python.nearest_neighbors as nearest_neighbors 9 | 10 | 11 | class DataProcessing: 12 | 13 | @staticmethod 14 | def load_pc_kitti(pc_path): 15 | scan = np.fromfile(pc_path, dtype=np.float32) 16 | scan = scan.reshape((-1, 4)) 17 | points = scan[:, 0:3] # get xyz 18 | return points 19 | 20 | @staticmethod 21 | def load_label_kitti(label_path, remap_lut): 22 | label = np.fromfile(label_path, dtype=np.uint32) 23 | label = label.reshape((-1)) 24 | sem_label = label & 0xFFFF # semantic label in lower half 25 | inst_label = label >> 16 # instance id in upper half 26 | assert ((sem_label + (inst_label << 16) == label).all()) 27 | sem_label = remap_lut[sem_label] 28 | return sem_label.astype(np.int32) 29 | 30 | @staticmethod 31 | def load_pc_semantic3d(filename): 32 | pc_pd = pd.read_csv(filename, header=None, delim_whitespace=True, dtype=np.float16) 33 | pc = pc_pd.values 34 | return pc 35 | 36 | @staticmethod 37 | def load_label_semantic3d(filename): 38 | label_pd = pd.read_csv(filename, header=None, delim_whitespace=True, dtype=np.uint8) 39 | cloud_labels = label_pd.values 40 | return cloud_labels 41 | 42 | @staticmethod 43 | def get_file_list(dataset_path, seq_list): 44 | data_list = [] 45 | for seq_id in seq_list: 46 | seq_path = join(dataset_path, seq_id) 47 | pc_path = join(seq_path, 'velodyne') 48 | new_data = [(seq_id, f[:-4]) for f in np.sort(os.listdir(pc_path))] 49 | data_list.extend(new_data) 50 | 51 | return data_list 52 | 53 | def get_active_list(list_root): 54 | train_list = [] 55 | pool_list = [] 56 | with open(join(list_root, 'label_data.json')) as f: 57 | train_list = json.load(f) 58 | with open(join(list_root, 'ulabel_data.json')) as f: 59 | pool_list = json.load(f) 60 | pool_list += train_list 61 | train_list = [] 62 | return train_list, pool_list 63 | 64 | @staticmethod 65 | def knn_search(support_pts, query_pts, k): 66 | """ 67 | :param support_pts: points you have, B*N1*3 68 | :param query_pts: points you want to know the neighbour index, B*N2*3 69 | :param k: Number of neighbours in knn search 70 | :return: neighbor_idx: neighboring points indexes, B*N2*k 71 | """ 72 | 73 | neighbor_idx = nearest_neighbors.knn_batch(support_pts, query_pts, k, omp=True) 74 | return neighbor_idx 75 | 76 | @staticmethod 77 | def data_aug(xyz, color, labels, idx, num_out): 78 | num_in = len(xyz) 79 | dup = np.random.choice(num_in, num_out - num_in) 80 | xyz_dup = xyz[dup, ...] 81 | xyz_aug = np.concatenate([xyz, xyz_dup], 0) 82 | color_dup = color[dup, ...] 83 | color_aug = np.concatenate([color, color_dup], 0) 84 | idx_dup = list(range(num_in)) + list(dup) 85 | idx_aug = idx[idx_dup] 86 | label_aug = labels[idx_dup] 87 | return xyz_aug, color_aug, idx_aug, label_aug 88 | 89 | @staticmethod 90 | def shuffle_idx(x): 91 | # random shuffle the index 92 | idx = np.arange(len(x)) 93 | np.random.shuffle(idx) 94 | return x[idx] 95 | 96 | @staticmethod 97 | def same_shuffl(a, b): 98 | # assert len(a) == len(b) 99 | # p = np.random.permutation(len(a)) 100 | # return a[p], b[p] 101 | 102 | assert len(a) == len(b) 103 | c = np.arange(len(a)) 104 | np.random.shuffle(c) 105 | return a[c], b[c] 106 | 107 | @staticmethod 108 | def shuffle_list(data_list): 109 | indices = np.arange(np.shape(data_list)[0]) 110 | np.random.shuffle(indices) 111 | data_list = data_list[indices] 112 | return data_list 113 | 114 | @staticmethod 115 | def grid_sub_sampling(points, features=None, labels=None, grid_size=0.1, verbose=0): 116 | """ 117 | CPP wrapper for a grid sub_sampling (method = barycenter for points and features 118 | :param points: (N, 3) matrix of input points 119 | :param features: optional (N, d) matrix of features (floating number) 120 | :param labels: optional (N,) matrix of integer labels 121 | :param grid_size: parameter defining the size of grid voxels 122 | :param verbose: 1 to display 123 | :return: sub_sampled points, with features and/or labels depending of the input 124 | """ 125 | 126 | if (features is None) and (labels is None): 127 | return cpp_subsampling.compute(points, sampleDl=grid_size, verbose=verbose) 128 | elif labels is None: 129 | return cpp_subsampling.compute(points, features=features, sampleDl=grid_size, verbose=verbose) 130 | elif features is None: 131 | return cpp_subsampling.compute(points, classes=labels, sampleDl=grid_size, verbose=verbose) 132 | else: 133 | return cpp_subsampling.compute(points, features=features, classes=labels, sampleDl=grid_size, 134 | verbose=verbose) 135 | 136 | @staticmethod 137 | def IoU_from_confusions(confusions): 138 | """ 139 | Computes IoU from confusion matrices. 140 | :param confusions: ([..., n_c, n_c] np.int32). Can be any dimension, the confusion matrices should be described by 141 | the last axes. n_c = number of classes 142 | :return: ([..., n_c] np.float32) IoU score 143 | """ 144 | 145 | # Compute TP, FP, FN. This assume that the second to last axis counts the truths (like the first axis of a 146 | # confusion matrix), and that the last axis counts the predictions (like the second axis of a confusion matrix) 147 | TP = np.diagonal(confusions, axis1=-2, axis2=-1) 148 | TP_plus_FN = np.sum(confusions, axis=-1) 149 | TP_plus_FP = np.sum(confusions, axis=-2) 150 | 151 | # Compute IoU 152 | IoU = TP / (TP_plus_FP + TP_plus_FN - TP + 1e-6) 153 | 154 | # Compute mIoU with only the actual classes 155 | mask = TP_plus_FN < 1e-3 156 | counts = np.sum(1 - mask, axis=-1, keepdims=True) 157 | mIoU = np.sum(IoU, axis=-1, keepdims=True) / (counts + 1e-6) 158 | 159 | # If class is absent, place mIoU in place of 0 IoU to get the actual mean later 160 | IoU += mask * mIoU 161 | return IoU 162 | 163 | @staticmethod 164 | def get_class_weights(data_root, paths, num_of_class): 165 | # pre-calculate the number of points in each category 166 | num_per_class = [0 for _ in range(num_of_class)] 167 | for file_path in tqdm(paths, total=len(paths)): 168 | label_path = join(data_root, file_path[0], 'labels', file_path[1] + '.npy') 169 | label = np.load(label_path).reshape(-1) 170 | inds, counts = np.unique(label, return_counts=True) 171 | for i, c in zip(inds, counts): 172 | if i == 0: # 0 : unlabeled 173 | continue 174 | else: 175 | num_per_class[i-1] += c 176 | ''' 177 | if dataset_name == 'S3DIS': 178 | num_per_class = np.array([3370714, 2856755, 4919229, 318158, 375640, 478001, 974733, 179 | 650464, 791496, 88727, 1284130, 229758, 2272837], dtype=np.int32) 180 | elif dataset_name == 'Semantic3D': 181 | num_per_class = np.array([5181602, 5012952, 6830086, 1311528, 10476365, 946982, 334860, 269353], 182 | dtype=np.int32) 183 | elif dataset_name == 'SemanticKITTI': 184 | num_per_class = np.array([55437630, 320797, 541736, 2578735, 3274484, 552662, 184064, 78858, 185 | 240942562, 17294618, 170599734, 6369672, 230413074, 101130274, 476491114, 186 | 9833174, 129609852, 4506626, 1168181]) 187 | ''' 188 | num_per_class = np.array(num_per_class) 189 | weight = num_per_class / float(sum(num_per_class)) 190 | ce_label_weight = 1 / (weight + 0.02) 191 | return np.expand_dims(ce_label_weight, axis=0) -------------------------------------------------------------------------------- /utils/help_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/huixiancheng/PCB-RandNet/304c84b67380a23ab66516581b348c7db2f9f0df/utils/help_utils.py -------------------------------------------------------------------------------- /utils/metric.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import threading 3 | from sklearn.metrics import confusion_matrix 4 | 5 | import torch 6 | 7 | class iouEval: 8 | def __init__(self, n_classes, device, ignore=None): 9 | self.n_classes = n_classes 10 | self.device = device 11 | # if ignore is larger than n_classes, consider no ignoreIndex 12 | self.ignore = torch.tensor(ignore).long() 13 | self.include = torch.tensor( 14 | [n for n in range(self.n_classes) if n not in self.ignore]).long() 15 | print("[IOU EVAL] IGNORE: ", self.ignore) 16 | print("[IOU EVAL] INCLUDE: ", self.include) 17 | self.reset() 18 | 19 | def num_classes(self): 20 | return self.n_classes 21 | 22 | def reset(self): 23 | self.conf_matrix = torch.zeros( 24 | (self.n_classes, self.n_classes), device=self.device).double() 25 | self.ones = None 26 | self.last_scan_size = None # for when variable scan size is used 27 | 28 | def addBatch(self, x, y): # x=preds, y=targets 29 | # if numpy, pass to pytorch 30 | # to tensor 31 | if isinstance(x, np.ndarray): 32 | x = torch.from_numpy(np.array(x)).long().to(self.device) 33 | if isinstance(y, np.ndarray): 34 | y = torch.from_numpy(np.array(y)).long().to(self.device) 35 | 36 | # sizes should be "batch_size x H x W" 37 | x_row = x.reshape(-1) # de-batchify 38 | y_row = y.reshape(-1) # de-batchify 39 | 40 | # idxs are labels and predictions 41 | idxs = torch.stack([x_row, y_row], dim=0) 42 | 43 | # ones is what I want to add to conf when I 44 | if self.ones is None or self.last_scan_size != idxs.shape[-1]: 45 | self.ones = torch.ones((idxs.shape[-1]), device=self.device).double() 46 | self.last_scan_size = idxs.shape[-1] 47 | 48 | # make confusion matrix (cols = gt, rows = pred) 49 | self.conf_matrix = self.conf_matrix.index_put_( 50 | tuple(idxs), self.ones, accumulate=True) 51 | 52 | # print(self.tp.shape) 53 | # print(self.fp.shape) 54 | # print(self.fn.shape) 55 | 56 | def getStats(self): 57 | # remove fp and fn from confusion on the ignore classes cols and rows 58 | conf = self.conf_matrix.clone().double() 59 | conf[self.ignore] = 0 60 | conf[:, self.ignore] = 0 61 | 62 | # get the clean stats 63 | tp = conf.diag() 64 | fp = conf.sum(dim=1) - tp 65 | fn = conf.sum(dim=0) - tp 66 | return tp, fp, fn 67 | 68 | def getIoU(self): 69 | tp, fp, fn = self.getStats() 70 | intersection = tp 71 | union = tp + fp + fn + 1e-15 72 | iou = intersection / union 73 | iou_mean = (intersection[self.include] / union[self.include]).mean() 74 | return iou_mean, iou # returns "iou mean", "iou per class" ALL CLASSES 75 | 76 | def getacc(self): 77 | tp, fp, fn = self.getStats() 78 | total_tp = tp.sum() 79 | total = tp[self.include].sum() + fp[self.include].sum() + 1e-15 80 | acc_mean = total_tp / total 81 | return acc_mean # returns "acc mean" 82 | 83 | class IoUCalculator: 84 | def __init__(self, cfg): 85 | self.gt_classes = [0 for _ in range(cfg.num_classes)] 86 | self.positive_classes = [0 for _ in range(cfg.num_classes)] 87 | self.true_positive_classes = [0 for _ in range(cfg.num_classes)] 88 | self.cfg = cfg 89 | self.lock = threading.Lock() 90 | 91 | def add_data(self, end_points): 92 | logits = end_points['valid_logits'] 93 | labels = end_points['valid_labels'] 94 | pred = logits.max(dim=1)[1] 95 | pred_valid = pred.detach().cpu().numpy() 96 | labels_valid = labels.detach().cpu().numpy() 97 | 98 | val_total_correct = 0 99 | val_total_seen = 0 100 | 101 | correct = np.sum(pred_valid == labels_valid) 102 | val_total_correct += correct 103 | val_total_seen += len(labels_valid) 104 | 105 | conf_matrix = confusion_matrix(labels_valid, pred_valid, np.arange(0, self.cfg.num_classes, 1)) 106 | self.lock.acquire() 107 | self.gt_classes += np.sum(conf_matrix, axis=1) 108 | self.positive_classes += np.sum(conf_matrix, axis=0) 109 | self.true_positive_classes += np.diagonal(conf_matrix) 110 | self.lock.release() 111 | 112 | def compute_iou(self): 113 | iou_list = [] 114 | for n in range(0, self.cfg.num_classes, 1): 115 | if float(self.gt_classes[n] + self.positive_classes[n] - self.true_positive_classes[n]) != 0: 116 | iou = self.true_positive_classes[n] / \ 117 | float(self.gt_classes[n] + self.positive_classes[n] - self.true_positive_classes[n]) 118 | iou_list.append(iou) 119 | else: 120 | iou_list.append(0.0) 121 | mean_iou = sum(iou_list) / float(self.cfg.num_classes) 122 | return mean_iou, iou_list 123 | 124 | 125 | def compute_acc(end_points): 126 | 127 | logits = end_points['valid_logits'] 128 | labels = end_points['valid_labels'] 129 | logits = logits.max(dim=1)[1] 130 | acc = (logits == labels).sum().float() / float(labels.shape[0]) 131 | end_points['acc'] = acc 132 | return acc, end_points 133 | -------------------------------------------------------------------------------- /utils/nearest_neighbors/KDTreeTableAdaptor.h: -------------------------------------------------------------------------------- 1 | /*********************************************************************** 2 | * Software License Agreement (BSD License) 3 | * 4 | * Copyright 2011-16 Jose Luis Blanco (joseluisblancoc@gmail.com). 5 | * All rights reserved. 6 | * 7 | * Redistribution and use in source and binary forms, with or without 8 | * modification, are permitted provided that the following conditions 9 | * are met: 10 | * 11 | * 1. Redistributions of source code must retain the above copyright 12 | * notice, this list of conditions and the following disclaimer. 13 | * 2. Redistributions in binary form must reproduce the above copyright 14 | * notice, this list of conditions and the following disclaimer in the 15 | * documentation and/or other materials provided with the distribution. 16 | * 17 | * THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR 18 | * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES 19 | * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. 20 | * IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT, 21 | * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT 22 | * NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, 23 | * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY 24 | * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 25 | * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF 26 | * THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 27 | *************************************************************************/ 28 | 29 | #pragma once 30 | 31 | #include "nanoflann.hpp" 32 | 33 | // #include 34 | 35 | // ===== This example shows how to use nanoflann with these types of containers: ======= 36 | //typedef std::vector > my_vector_of_vectors_t; 37 | //typedef std::vector my_vector_of_vectors_t; // This requires #include 38 | // ===================================================================================== 39 | 40 | 41 | /** A simple vector-of-vectors adaptor for nanoflann, without duplicating the storage. 42 | * The i'th vector represents a point in the state space. 43 | * 44 | * \tparam DIM If set to >0, it specifies a compile-time fixed dimensionality for the points in the data set, allowing more compiler optimizations. 45 | * \tparam num_t The type of the point coordinates (typically, double or float). 46 | * \tparam Distance The distance metric to use: nanoflann::metric_L1, nanoflann::metric_L2, nanoflann::metric_L2_Simple, etc. 47 | * \tparam IndexType The type for indices in the KD-tree index (typically, size_t of int) 48 | */ 49 | // template 50 | // struct KDTreeVectorAdaptor 51 | // { 52 | // typedef KDTreeVectorAdaptor self_t; 53 | // typedef typename Distance::template traits::distance_t metric_t; 54 | // typedef nanoflann::KDTreeSingleIndexAdaptor< metric_t,self_t,DIM,IndexType> index_t; 55 | 56 | // index_t* index; //! The kd-tree index for the user to call its methods as usual with any other FLANN index. 57 | // size_t dims; 58 | 59 | // /// Constructor: takes a const ref to the vector of vectors object with the data points 60 | // KDTreeVectorAdaptor(const size_t dims /* dimensionality */, const VectorType &mat, const int leaf_max_size = 10) : m_data(mat) 61 | // { 62 | // assert(mat.size() != 0); 63 | // this->dims= dims; 64 | // index = new index_t( static_cast(dims), *this /* adaptor */, nanoflann::KDTreeSingleIndexAdaptorParams(leaf_max_size ) ); 65 | // index->buildIndex(); 66 | // } 67 | 68 | // ~KDTreeVectorAdaptor() { 69 | // delete index; 70 | // } 71 | 72 | // const VectorType &m_data; 73 | 74 | // /** Query for the \a num_closest closest points to a given point (entered as query_point[0:dim-1]). 75 | // * Note that this is a short-cut method for index->findNeighbors(). 76 | // * The user can also call index->... methods as desired. 77 | // * \note nChecks_IGNORED is ignored but kept for compatibility with the original FLANN interface. 78 | // */ 79 | // inline void query(const num_t *query_point, const size_t num_closest, IndexType *out_indices, num_t *out_distances_sq, const int nChecks_IGNORED = 10) const 80 | // { 81 | // nanoflann::KNNResultSet resultSet(num_closest); 82 | // resultSet.init(out_indices, out_distances_sq); 83 | // index->findNeighbors(resultSet, query_point, nanoflann::SearchParams()); 84 | // } 85 | 86 | // /** @name Interface expected by KDTreeSingleIndexAdaptor 87 | // * @{ */ 88 | 89 | // const self_t & derived() const { 90 | // return *this; 91 | // } 92 | // self_t & derived() { 93 | // return *this; 94 | // } 95 | 96 | // // Must return the number of data points 97 | // inline size_t kdtree_get_point_count() const { 98 | // return m_data.size()/this->dims; 99 | // } 100 | 101 | // // Returns the dim'th component of the idx'th point in the class: 102 | // inline num_t kdtree_get_pt(const size_t idx, const size_t dim) const { 103 | // return m_data[idx*this->dims + dim]; 104 | // } 105 | 106 | // // Optional bounding-box computation: return false to default to a standard bbox computation loop. 107 | // // Return true if the BBOX was already computed by the class and returned in "bb" so it can be avoided to redo it again. 108 | // // Look at bb.size() to find out the expected dimensionality (e.g. 2 or 3 for point clouds) 109 | // template 110 | // bool kdtree_get_bbox(BBOX & /*bb*/) const { 111 | // return false; 112 | // } 113 | 114 | // /** @} */ 115 | 116 | // }; // end of KDTreeVectorOfVectorsAdaptor 117 | 118 | 119 | 120 | 121 | template 122 | struct KDTreeTableAdaptor 123 | { 124 | typedef KDTreeTableAdaptor self_t; 125 | typedef typename Distance::template traits::distance_t metric_t; 126 | typedef nanoflann::KDTreeSingleIndexAdaptor< metric_t,self_t,DIM,IndexType> index_t; 127 | 128 | index_t* index; //! The kd-tree index for the user to call its methods as usual with any other FLANN index. 129 | size_t dim; 130 | size_t npts; 131 | const TableType* m_data; 132 | 133 | /// Constructor: takes a const ref to the vector of vectors object with the data points 134 | KDTreeTableAdaptor(const size_t npts, const size_t dim, const TableType* mat, const int leaf_max_size = 10) : m_data(mat), dim(dim), npts(npts) 135 | { 136 | assert(npts != 0); 137 | index = new index_t( static_cast(dim), *this /* adaptor */, nanoflann::KDTreeSingleIndexAdaptorParams(leaf_max_size ) ); 138 | index->buildIndex(); 139 | } 140 | 141 | ~KDTreeTableAdaptor() { 142 | delete index; 143 | } 144 | 145 | 146 | /** Query for the \a num_closest closest points to a given point (entered as query_point[0:dim-1]). 147 | * Note that this is a short-cut method for index->findNeighbors(). 148 | * The user can also call index->... methods as desired. 149 | * \note nChecks_IGNORED is ignored but kept for compatibility with the original FLANN interface. 150 | */ 151 | inline void query(const num_t *query_point, const size_t num_closest, IndexType *out_indices, num_t *out_distances_sq, const int nChecks_IGNORED = 10) const 152 | { 153 | nanoflann::KNNResultSet resultSet(num_closest); 154 | resultSet.init(out_indices, out_distances_sq); 155 | index->findNeighbors(resultSet, query_point, nanoflann::SearchParams()); 156 | } 157 | 158 | /** @name Interface expected by KDTreeSingleIndexAdaptor 159 | * @{ */ 160 | 161 | const self_t & derived() const { 162 | return *this; 163 | } 164 | self_t & derived() { 165 | return *this; 166 | } 167 | 168 | // Must return the number of data points 169 | inline size_t kdtree_get_point_count() const { 170 | return this->npts; 171 | } 172 | 173 | // Returns the dim'th component of the idx'th point in the class: 174 | inline num_t kdtree_get_pt(const size_t pts_id, const size_t coord_id) const { 175 | return m_data[pts_id*this->dim + coord_id]; 176 | } 177 | 178 | // Optional bounding-box computation: return false to default to a standard bbox computation loop. 179 | // Return true if the BBOX was already computed by the class and returned in "bb" so it can be avoided to redo it again. 180 | // Look at bb.size() to find out the expected dimensionality (e.g. 2 or 3 for point clouds) 181 | template 182 | bool kdtree_get_bbox(BBOX & /*bb*/) const { 183 | return false; 184 | } 185 | 186 | /** @} */ 187 | 188 | }; // end of KDTreeVectorOfVectorsAdaptor 189 | 190 | -------------------------------------------------------------------------------- /utils/nearest_neighbors/knn.pyx: -------------------------------------------------------------------------------- 1 | # distutils: language = c++ 2 | # distutils: sources = knn.cxx 3 | 4 | import numpy as np 5 | cimport numpy as np 6 | import cython 7 | 8 | cdef extern from "knn_.h": 9 | void cpp_knn(const float* points, const size_t npts, const size_t dim, 10 | const float* queries, const size_t nqueries, 11 | const size_t K, long* indices) 12 | 13 | void cpp_knn_omp(const float* points, const size_t npts, const size_t dim, 14 | const float* queries, const size_t nqueries, 15 | const size_t K, long* indices) 16 | 17 | void cpp_knn_batch(const float* batch_data, const size_t batch_size, const size_t npts, const size_t dim, 18 | const float* queries, const size_t nqueries, 19 | const size_t K, long* batch_indices) 20 | 21 | void cpp_knn_batch_omp(const float* batch_data, const size_t batch_size, const size_t npts, const size_t dim, 22 | const float* queries, const size_t nqueries, 23 | const size_t K, long* batch_indices) 24 | 25 | void cpp_knn_batch_distance_pick(const float* batch_data, const size_t batch_size, const size_t npts, const size_t dim, 26 | float* queries, const size_t nqueries, 27 | const size_t K, long* batch_indices) 28 | 29 | void cpp_knn_batch_distance_pick_omp(const float* batch_data, const size_t batch_size, const size_t npts, const size_t dim, 30 | float* batch_queries, const size_t nqueries, 31 | const size_t K, long* batch_indices) 32 | 33 | def knn(pts, queries, K, omp=False): 34 | 35 | # define shape parameters 36 | cdef int npts 37 | cdef int dim 38 | cdef int K_cpp 39 | cdef int nqueries 40 | 41 | # define tables 42 | cdef np.ndarray[np.float32_t, ndim=2] pts_cpp 43 | cdef np.ndarray[np.float32_t, ndim=2] queries_cpp 44 | cdef np.ndarray[np.int64_t, ndim=2] indices_cpp 45 | 46 | # set shape values 47 | npts = pts.shape[0] 48 | nqueries = queries.shape[0] 49 | dim = pts.shape[1] 50 | K_cpp = K 51 | 52 | # create indices tensor 53 | indices = np.zeros((queries.shape[0], K), dtype=np.int64) 54 | 55 | pts_cpp = np.ascontiguousarray(pts, dtype=np.float32) 56 | queries_cpp = np.ascontiguousarray(queries, dtype=np.float32) 57 | indices_cpp = indices 58 | 59 | # normal estimation 60 | if omp: 61 | cpp_knn_omp( pts_cpp.data, npts, dim, 62 | queries_cpp.data, nqueries, 63 | K_cpp, indices_cpp.data) 64 | else: 65 | cpp_knn( pts_cpp.data, npts, dim, 66 | queries_cpp.data, nqueries, 67 | K_cpp, indices_cpp.data) 68 | 69 | return indices 70 | 71 | def knn_batch(pts, queries, K, omp=False): 72 | 73 | # define shape parameters 74 | cdef int batch_size 75 | cdef int npts 76 | cdef int nqueries 77 | cdef int K_cpp 78 | cdef int dim 79 | 80 | # define tables 81 | cdef np.ndarray[np.float32_t, ndim=3] pts_cpp 82 | cdef np.ndarray[np.float32_t, ndim=3] queries_cpp 83 | cdef np.ndarray[np.int64_t, ndim=3] indices_cpp 84 | 85 | # set shape values 86 | batch_size = pts.shape[0] 87 | npts = pts.shape[1] 88 | dim = pts.shape[2] 89 | nqueries = queries.shape[1] 90 | K_cpp = K 91 | 92 | # create indices tensor 93 | indices = np.zeros((pts.shape[0], queries.shape[1], K), dtype=np.int64) 94 | 95 | pts_cpp = np.ascontiguousarray(pts, dtype=np.float32) 96 | queries_cpp = np.ascontiguousarray(queries, dtype=np.float32) 97 | indices_cpp = indices 98 | 99 | # normal estimation 100 | if omp: 101 | cpp_knn_batch_omp( pts_cpp.data, batch_size, npts, dim, 102 | queries_cpp.data, nqueries, 103 | K_cpp, indices_cpp.data) 104 | else: 105 | cpp_knn_batch( pts_cpp.data, batch_size, npts, dim, 106 | queries_cpp.data, nqueries, 107 | K_cpp, indices_cpp.data) 108 | 109 | return indices 110 | 111 | def knn_batch_distance_pick(pts, nqueries, K, omp=False): 112 | 113 | # define shape parameters 114 | cdef int batch_size 115 | cdef int npts 116 | cdef int nqueries_cpp 117 | cdef int K_cpp 118 | cdef int dim 119 | 120 | # define tables 121 | cdef np.ndarray[np.float32_t, ndim=3] pts_cpp 122 | cdef np.ndarray[np.float32_t, ndim=3] queries_cpp 123 | cdef np.ndarray[np.int64_t, ndim=3] indices_cpp 124 | 125 | # set shape values 126 | batch_size = pts.shape[0] 127 | npts = pts.shape[1] 128 | dim = pts.shape[2] 129 | nqueries_cpp = nqueries 130 | K_cpp = K 131 | 132 | # create indices tensor 133 | indices = np.zeros((pts.shape[0], nqueries, K), dtype=np.long) 134 | queries = np.zeros((pts.shape[0], nqueries, dim), dtype=np.float32) 135 | 136 | pts_cpp = np.ascontiguousarray(pts, dtype=np.float32) 137 | queries_cpp = np.ascontiguousarray(queries, dtype=np.float32) 138 | indices_cpp = indices 139 | 140 | if omp: 141 | cpp_knn_batch_distance_pick_omp( pts_cpp.data, batch_size, npts, dim, 142 | queries_cpp.data, nqueries, 143 | K_cpp, indices_cpp.data) 144 | else: 145 | cpp_knn_batch_distance_pick( pts_cpp.data, batch_size, npts, dim, 146 | queries_cpp.data, nqueries, 147 | K_cpp, indices_cpp.data) 148 | 149 | return indices, queries -------------------------------------------------------------------------------- /utils/nearest_neighbors/knn_.cxx: -------------------------------------------------------------------------------- 1 | 2 | #include "knn_.h" 3 | #include "nanoflann.hpp" 4 | using namespace nanoflann; 5 | 6 | #include "KDTreeTableAdaptor.h" 7 | 8 | #include 9 | #include 10 | #include 11 | #include 12 | 13 | #include 14 | #include 15 | #include 16 | #include 17 | 18 | using namespace std; 19 | 20 | 21 | 22 | void cpp_knn(const float* points, const size_t npts, const size_t dim, 23 | const float* queries, const size_t nqueries, 24 | const size_t K, long* indices){ 25 | 26 | // create the kdtree 27 | typedef KDTreeTableAdaptor< float, float> KDTree; 28 | KDTree mat_index(npts, dim, points, 10); 29 | mat_index.index->buildIndex(); 30 | 31 | std::vector out_dists_sqr(K); 32 | std::vector out_ids(K); 33 | 34 | // iterate over the points 35 | for(size_t i=0; i resultSet(K); 38 | resultSet.init(&out_ids[0], &out_dists_sqr[0] ); 39 | mat_index.index->findNeighbors(resultSet, &queries[i*dim], nanoflann::SearchParams(10)); 40 | for(size_t j=0; j KDTree; 52 | KDTree mat_index(npts, dim, points, 10); 53 | mat_index.index->buildIndex(); 54 | 55 | 56 | // iterate over the points 57 | # pragma omp parallel for 58 | for(size_t i=0; i out_ids(K); 60 | std::vector out_dists_sqr(K); 61 | 62 | nanoflann::KNNResultSet resultSet(K); 63 | resultSet.init(&out_ids[0], &out_dists_sqr[0] ); 64 | mat_index.index->findNeighbors(resultSet, &queries[i*dim], nanoflann::SearchParams(10)); 65 | for(size_t j=0; j KDTree; 83 | KDTree mat_index(npts, dim, points, 10); 84 | 85 | mat_index.index->buildIndex(); 86 | 87 | std::vector out_dists_sqr(K); 88 | std::vector out_ids(K); 89 | 90 | // iterate over the points 91 | for(size_t i=0; i resultSet(K); 93 | resultSet.init(&out_ids[0], &out_dists_sqr[0] ); 94 | mat_index.index->findNeighbors(resultSet, &queries[bid*nqueries*dim + i*dim], nanoflann::SearchParams(10)); 95 | for(size_t j=0; j KDTree; 116 | KDTree mat_index(npts, dim, points, 10); 117 | 118 | mat_index.index->buildIndex(); 119 | 120 | std::vector out_dists_sqr(K); 121 | std::vector out_ids(K); 122 | 123 | // iterate over the points 124 | for(size_t i=0; i resultSet(K); 126 | resultSet.init(&out_ids[0], &out_dists_sqr[0] ); 127 | mat_index.index->findNeighbors(resultSet, &queries[bid*nqueries*dim + i*dim], nanoflann::SearchParams(10)); 128 | for(size_t j=0; j KDTree; 153 | KDTree tree(npts, dim, points, 10); 154 | tree.index->buildIndex(); 155 | 156 | vector used(npts, 0); 157 | int current_id = 0; 158 | for(size_t ptid=0; ptid possible_ids; 162 | while(possible_ids.size() == 0){ 163 | for(size_t i=0; i query(3); 178 | for(size_t i=0; i dists(K); 183 | std::vector ids(K); 184 | nanoflann::KNNResultSet resultSet(K); 185 | resultSet.init(&ids[0], &dists[0] ); 186 | tree.index->findNeighbors(resultSet, &query[0], nanoflann::SearchParams(10)); 187 | 188 | for(size_t i=0; i KDTree; 221 | KDTree tree(npts, dim, points, 10); 222 | tree.index->buildIndex(); 223 | 224 | vector used(npts, 0); 225 | int current_id = 0; 226 | for(size_t ptid=0; ptid possible_ids; 230 | while(possible_ids.size() == 0){ 231 | for(size_t i=0; i query(3); 246 | for(size_t i=0; i dists(K); 251 | std::vector ids(K); 252 | nanoflann::KNNResultSet resultSet(K); 253 | resultSet.init(&ids[0], &dists[0] ); 254 | tree.index->findNeighbors(resultSet, &query[0], nanoflann::SearchParams(10)); 255 | 256 | for(size_t i=0; i 4 | void cpp_knn(const float* points, const size_t npts, const size_t dim, 5 | const float* queries, const size_t nqueries, 6 | const size_t K, long* indices); 7 | 8 | void cpp_knn_omp(const float* points, const size_t npts, const size_t dim, 9 | const float* queries, const size_t nqueries, 10 | const size_t K, long* indices); 11 | 12 | 13 | void cpp_knn_batch(const float* batch_data, const size_t batch_size, const size_t npts, const size_t dim, 14 | const float* queries, const size_t nqueries, 15 | const size_t K, long* batch_indices); 16 | 17 | void cpp_knn_batch_omp(const float* batch_data, const size_t batch_size, const size_t npts, const size_t dim, 18 | const float* queries, const size_t nqueries, 19 | const size_t K, long* batch_indices); 20 | 21 | void cpp_knn_batch_distance_pick(const float* batch_data, const size_t batch_size, const size_t npts, const size_t dim, 22 | float* queries, const size_t nqueries, 23 | const size_t K, long* batch_indices); 24 | 25 | void cpp_knn_batch_distance_pick_omp(const float* batch_data, const size_t batch_size, const size_t npts, const size_t dim, 26 | float* batch_queries, const size_t nqueries, 27 | const size_t K, long* batch_indices); -------------------------------------------------------------------------------- /utils/nearest_neighbors/setup.py: -------------------------------------------------------------------------------- 1 | from distutils.core import setup 2 | from distutils.extension import Extension 3 | from Cython.Distutils import build_ext 4 | import numpy 5 | 6 | 7 | 8 | ext_modules = [Extension( 9 | "nearest_neighbors", 10 | sources=["knn.pyx", "knn_.cxx",], # source file(s) 11 | include_dirs=["./", numpy.get_include()], 12 | language="c++", 13 | extra_compile_args = [ "-std=c++11", "-fopenmp",], 14 | extra_link_args=["-std=c++11", '-fopenmp'], 15 | )] 16 | 17 | setup( 18 | name = "KNN NanoFLANN", 19 | ext_modules = ext_modules, 20 | cmdclass = {'build_ext': build_ext}, 21 | ) 22 | -------------------------------------------------------------------------------- /utils/nearest_neighbors/test.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import lib.python.nearest_neighbors as nearest_neighbors 3 | import time 4 | 5 | batch_size = 16 6 | num_points = 81920 7 | K = 16 8 | pc = np.random.rand(batch_size, num_points, 3).astype(np.float32) 9 | 10 | # nearest neighbours 11 | start = time.time() 12 | neigh_idx = nearest_neighbors.knn_batch(pc, pc, K, omp=True) 13 | print(time.time() - start) 14 | 15 | 16 | -------------------------------------------------------------------------------- /utils/np_ioueval.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # This file is covered by the LICENSE file in the root of this project. 3 | 4 | import numpy as np 5 | 6 | 7 | class iouEval: 8 | def __init__(self, n_classes, ignore=None): 9 | # classes 10 | self.n_classes = n_classes 11 | 12 | # What to include and ignore from the means 13 | self.ignore = np.array(ignore, dtype=np.int64) 14 | self.include = np.array( 15 | [n for n in range(self.n_classes) if n not in self.ignore], dtype=np.int64) 16 | print("[IOU EVAL] IGNORE: ", self.ignore) 17 | print("[IOU EVAL] INCLUDE: ", self.include) 18 | 19 | # reset the class counters 20 | self.reset() 21 | 22 | def num_classes(self): 23 | return self.n_classes 24 | 25 | def reset(self): 26 | self.conf_matrix = np.zeros((self.n_classes, 27 | self.n_classes), 28 | dtype=np.int64) 29 | 30 | def addBatch(self, x, y): # x=preds, y=targets 31 | # sizes should be matching 32 | x_row = x.reshape(-1) # de-batchify 33 | y_row = y.reshape(-1) # de-batchify 34 | 35 | # check 36 | assert(x_row.shape == x_row.shape) 37 | 38 | # create indexes 39 | idxs = tuple(np.stack((x_row, y_row), axis=0)) 40 | 41 | # make confusion matrix (cols = gt, rows = pred) 42 | np.add.at(self.conf_matrix, idxs, 1) 43 | 44 | def getStats(self): 45 | # remove fp from confusion on the ignore classes cols 46 | conf = self.conf_matrix.copy() 47 | conf[:, self.ignore] = 0 48 | 49 | # get the clean stats 50 | tp = np.diag(conf) 51 | fp = conf.sum(axis=1) - tp 52 | fn = conf.sum(axis=0) - tp 53 | return tp, fp, fn 54 | 55 | def getIoU(self): 56 | tp, fp, fn = self.getStats() 57 | intersection = tp 58 | union = tp + fp + fn + 1e-15 59 | iou = intersection / union 60 | iou_mean = (intersection[self.include] / union[self.include]).mean() 61 | return iou_mean, iou # returns "iou mean", "iou per class" ALL CLASSES 62 | 63 | def getacc(self): 64 | tp, fp, fn = self.getStats() 65 | total_tp = tp.sum() 66 | total = tp[self.include].sum() + fp[self.include].sum() + 1e-15 67 | acc_mean = total_tp / total 68 | return acc_mean # returns "acc mean" 69 | 70 | def get_confusion(self): 71 | return self.conf_matrix.copy() 72 | -------------------------------------------------------------------------------- /utils/semantic-kitti.yaml: -------------------------------------------------------------------------------- 1 | # This file is covered by the LICENSE file in the root of this project. 2 | labels: 3 | 0 : "unlabeled" 4 | 1 : "outlier" 5 | 10: "car" 6 | 11: "bicycle" 7 | 13: "bus" 8 | 15: "motorcycle" 9 | 16: "on-rails" 10 | 18: "truck" 11 | 20: "other-vehicle" 12 | 30: "person" 13 | 31: "bicyclist" 14 | 32: "motorcyclist" 15 | 40: "road" 16 | 44: "parking" 17 | 48: "sidewalk" 18 | 49: "other-ground" 19 | 50: "building" 20 | 51: "fence" 21 | 52: "other-structure" 22 | 60: "lane-marking" 23 | 70: "vegetation" 24 | 71: "trunk" 25 | 72: "terrain" 26 | 80: "pole" 27 | 81: "traffic-sign" 28 | 99: "other-object" 29 | 252: "moving-car" 30 | 253: "moving-bicyclist" 31 | 254: "moving-person" 32 | 255: "moving-motorcyclist" 33 | 256: "moving-on-rails" 34 | 257: "moving-bus" 35 | 258: "moving-truck" 36 | 259: "moving-other-vehicle" 37 | color_map: # bgr 38 | 0 : [0, 0, 0] 39 | 1 : [0, 0, 255] 40 | 10: [245, 150, 100] 41 | 11: [245, 230, 100] 42 | 13: [250, 80, 100] 43 | 15: [150, 60, 30] 44 | 16: [255, 0, 0] 45 | 18: [180, 30, 80] 46 | 20: [255, 0, 0] 47 | 30: [30, 30, 255] 48 | 31: [200, 40, 255] 49 | 32: [90, 30, 150] 50 | 40: [255, 0, 255] 51 | 44: [255, 150, 255] 52 | 48: [75, 0, 75] 53 | 49: [75, 0, 175] 54 | 50: [0, 200, 255] 55 | 51: [50, 120, 255] 56 | 52: [0, 150, 255] 57 | 60: [170, 255, 150] 58 | 70: [0, 175, 0] 59 | 71: [0, 60, 135] 60 | 72: [80, 240, 150] 61 | 80: [150, 240, 255] 62 | 81: [0, 0, 255] 63 | 99: [255, 255, 50] 64 | 252: [245, 150, 100] 65 | 256: [255, 0, 0] 66 | 253: [200, 40, 255] 67 | 254: [30, 30, 255] 68 | 255: [90, 30, 150] 69 | 257: [250, 80, 100] 70 | 258: [180, 30, 80] 71 | 259: [255, 0, 0] 72 | content: # as a ratio with the total number of points 73 | 0: 0.018889854628292943 74 | 1: 0.0002937197336781505 75 | 10: 0.040818519255974316 76 | 11: 0.00016609538710764618 77 | 13: 2.7879693665067774e-05 78 | 15: 0.00039838616015114444 79 | 16: 0.0 80 | 18: 0.0020633612104619787 81 | 20: 0.0016218197275284021 82 | 30: 0.00017698551338515307 83 | 31: 1.1065903904919655e-08 84 | 32: 5.532951952459828e-09 85 | 40: 0.1987493871255525 86 | 44: 0.014717169549888214 87 | 48: 0.14392298360372 88 | 49: 0.0039048553037472045 89 | 50: 0.1326861944777486 90 | 51: 0.0723592229456223 91 | 52: 0.002395131480328884 92 | 60: 4.7084144280367186e-05 93 | 70: 0.26681502148037506 94 | 71: 0.006035012012626033 95 | 72: 0.07814222006271769 96 | 80: 0.002855498193863172 97 | 81: 0.0006155958086189918 98 | 99: 0.009923127583046915 99 | 252: 0.001789309418528068 100 | 253: 0.00012709999297008662 101 | 254: 0.00016059776092534436 102 | 255: 3.745553104802113e-05 103 | 256: 0.0 104 | 257: 0.00011351574470342043 105 | 258: 0.00010157861367183268 106 | 259: 4.3840131989471124e-05 107 | # classes that are indistinguishable from single scan or inconsistent in 108 | # ground truth are mapped to their closest equivalent 109 | learning_map: 110 | 0 : 0 # "unlabeled" 111 | 1 : 0 # "outlier" mapped to "unlabeled" --------------------------mapped 112 | 10: 1 # "car" 113 | 11: 2 # "bicycle" 114 | 13: 5 # "bus" mapped to "other-vehicle" --------------------------mapped 115 | 15: 3 # "motorcycle" 116 | 16: 5 # "on-rails" mapped to "other-vehicle" ---------------------mapped 117 | 18: 4 # "truck" 118 | 20: 5 # "other-vehicle" 119 | 30: 6 # "person" 120 | 31: 7 # "bicyclist" 121 | 32: 8 # "motorcyclist" 122 | 40: 9 # "road" 123 | 44: 10 # "parking" 124 | 48: 11 # "sidewalk" 125 | 49: 12 # "other-ground" 126 | 50: 13 # "building" 127 | 51: 14 # "fence" 128 | 52: 0 # "other-structure" mapped to "unlabeled" ------------------mapped 129 | 60: 9 # "lane-marking" to "road" ---------------------------------mapped 130 | 70: 15 # "vegetation" 131 | 71: 16 # "trunk" 132 | 72: 17 # "terrain" 133 | 80: 18 # "pole" 134 | 81: 19 # "traffic-sign" 135 | 99: 0 # "other-object" to "unlabeled" ----------------------------mapped 136 | 252: 1 # "moving-car" to "car" ------------------------------------mapped 137 | 253: 7 # "moving-bicyclist" to "bicyclist" ------------------------mapped 138 | 254: 6 # "moving-person" to "person" ------------------------------mapped 139 | 255: 8 # "moving-motorcyclist" to "motorcyclist" ------------------mapped 140 | 256: 5 # "moving-on-rails" mapped to "other-vehicle" --------------mapped 141 | 257: 5 # "moving-bus" mapped to "other-vehicle" -------------------mapped 142 | 258: 4 # "moving-truck" to "truck" --------------------------------mapped 143 | 259: 5 # "moving-other"-vehicle to "other-vehicle" ----------------mapped 144 | learning_map_inv: # inverse of previous map 145 | 0: 0 # "unlabeled", and others ignored 146 | 1: 10 # "car" 147 | 2: 11 # "bicycle" 148 | 3: 15 # "motorcycle" 149 | 4: 18 # "truck" 150 | 5: 20 # "other-vehicle" 151 | 6: 30 # "person" 152 | 7: 31 # "bicyclist" 153 | 8: 32 # "motorcyclist" 154 | 9: 40 # "road" 155 | 10: 44 # "parking" 156 | 11: 48 # "sidewalk" 157 | 12: 49 # "other-ground" 158 | 13: 50 # "building" 159 | 14: 51 # "fence" 160 | 15: 70 # "vegetation" 161 | 16: 71 # "trunk" 162 | 17: 72 # "terrain" 163 | 18: 80 # "pole" 164 | 19: 81 # "traffic-sign" 165 | learning_ignore: # Ignore classes 166 | 0: True # "unlabeled", and others ignored 167 | 1: False # "car" 168 | 2: False # "bicycle" 169 | 3: False # "motorcycle" 170 | 4: False # "truck" 171 | 5: False # "other-vehicle" 172 | 6: False # "person" 173 | 7: False # "bicyclist" 174 | 8: False # "motorcyclist" 175 | 9: False # "road" 176 | 10: False # "parking" 177 | 11: False # "sidewalk" 178 | 12: False # "other-ground" 179 | 13: False # "building" 180 | 14: False # "fence" 181 | 15: False # "vegetation" 182 | 16: False # "trunk" 183 | 17: False # "terrain" 184 | 18: False # "pole" 185 | 19: False # "traffic-sign" 186 | split: # sequence numbers 187 | train: 188 | - 0 189 | - 1 190 | - 2 191 | - 3 192 | - 4 193 | - 5 194 | - 6 195 | - 7 196 | - 9 197 | - 10 198 | valid: 199 | - 8 200 | test: 201 | - 11 202 | - 12 203 | - 13 204 | - 14 205 | - 15 206 | - 16 207 | - 17 208 | - 18 209 | - 19 210 | - 20 211 | - 21 212 | -------------------------------------------------------------------------------- /utils/semantic-poss.yaml: -------------------------------------------------------------------------------- 1 | # This file is covered by the LICENSE file in the root of this project. 2 | labels: 3 | 0: "unlabeled" 4 | 4: "1 person" 5 | 5: "2+ person" 6 | 6: "rider" 7 | 7: "car" 8 | 8: "trunk" 9 | 9: "plants" 10 | 10: "traffic sign 1" # standing sign 11 | 11: "traffic sign 2" # hanging sign 12 | 12: "traffic sign 3" # high/big hanging sign 13 | 13: "pole" 14 | 14: "trashcan" 15 | 15: "building" 16 | 17: "fence" 17 | 16: "cone/stone" 18 | 21: "bike" 19 | 22: "ground" # class definition 20 | 21 | color_map: # bgr 22 | 23 | 0 : [0, 0, 0] # 0: "unlabeled" 24 | 4 : [30, 30, 255] # 4: "1 person" 25 | 5 : [30, 30, 255] # 5: "2+ person" 26 | 6 : [200, 40, 255] # 6: "rider" 27 | 7 : [245, 150, 100] # 7: "car" 28 | 8 : [0,60,135] # 8: "trunk" 29 | 9 : [0, 175, 0] # 9: "plants" 30 | 10: [0, 0, 255] # 10: "traffic sign 1" 31 | 11: [0, 0, 255] # 11: "traffic sign 2" 32 | 12: [0, 0, 255] # 12: "traffic sign 3" 33 | 13: [150, 240, 255] # 13: "pole" 34 | 14: [0, 255, 125] # 14: "trashcan" 35 | 15: [0, 200, 255] # 15: "building" 36 | 16: [255, 255, 50] # 16: "cone/stone" 37 | 17: [50, 120, 255] # 17: "fence" 38 | 21: [245, 230, 100] # 21: "bike" 39 | 22: [128, 128, 128] # 22: "ground" 40 | 41 | 42 | content: # as a ratio with the total number of points 43 | 0: 0.018889854628292943 44 | 1: 0.0002937197336781505 45 | 10: 0.040818519255974316 46 | 11: 0.00016609538710764618 47 | 13: 2.7879693665067774e-05 48 | 15: 0.00039838616015114444 49 | 16: 0.0 50 | 18: 0.0020633612104619787 51 | 20: 0.0016218197275284021 52 | 30: 0.00017698551338515307 53 | 31: 1.1065903904919655e-08 54 | 32: 5.532951952459828e-09 55 | 40: 0.1987493871255525 56 | 44: 0.014717169549888214 57 | 48: 0.14392298360372 58 | 49: 0.0039048553037472045 59 | 50: 0.1326861944777486 60 | 51: 0.0723592229456223 61 | 52: 0.002395131480328884 62 | 60: 4.7084144280367186e-05 63 | 70: 0.26681502148037506 64 | 71: 0.006035012012626033 65 | 72: 0.07814222006271769 66 | 80: 0.002855498193863172 67 | 81: 0.0006155958086189918 68 | 99: 0.009923127583046915 69 | 252: 0.001789309418528068 70 | 253: 0.00012709999297008662 71 | 254: 0.00016059776092534436 72 | 255: 3.745553104802113e-05 73 | 256: 0.0 74 | 257: 0.00011351574470342043 75 | 258: 0.00010157861367183268 76 | 259: 4.3840131989471124e-05 77 | # classes that are indistinguishable from single scan or inconsistent in 78 | # ground truth are mapped to their closest equivalent 79 | 80 | # 11 CLASSES 81 | learning_map: 82 | 0: 0 #"unlabeled" 83 | 4: 1 # "1 person" --> "people" ----------------mapped 84 | 5: 1 # "2+ person" --> "people" ---------------mapped 85 | 6: 2 #"rider" 86 | 7: 3 #"car" 87 | 8: 4 #"trunk" 88 | 9: 5 #"plants" 89 | 10: 6 # "traffic sign 1" # standing sign -->traffic sign----------------mapped 90 | 11: 6 #"traffic sign 2" # hanging sign-->traffic sign----------------mapped 91 | 12: 6 #"traffic sign 3" # high/big hanging sign-->traffic sign----------------mapped 92 | 13: 7 #"pole" 93 | 14: 0 #"trashcan" --> "unlabeled" ----------------mapped 94 | 15: 8 #"building" 95 | 16: 0 # "cone/stone" --> "unlabeled" ----------------mapped 96 | 17: 9 # "fence" 97 | 21: 10 #"bike" 98 | 22: 11 #"ground" # class definition 99 | 100 | learning_map_inv: # inverse of previous map 101 | 0: 0 # "unlabeled" 102 | 1: 4 # "people" 103 | 2: 6 # "rider" 104 | 3: 7 # "car" 105 | 4: 8 # "trunk" 106 | 5: 9 # "plants" 107 | 6: 10 # "traffic sign" 108 | 7: 13 # "pole" 109 | 8: 15 # "building" 110 | 9: 17 # "fence" 111 | 10: 21 # "bike" 112 | 11: 22 # "ground" 113 | 114 | learning_ignore: # Ignore classes 115 | 0: True # "unlabeled", and others ignored 116 | 1: False # "people" 117 | 2: False # "rider" 118 | 3: False # "car" 119 | 4: False # "trunk" 120 | 5: False # "plants" 121 | 6: False # "traffic sign" 122 | 7: False # "pole" 123 | 8: False # "building" 124 | 9: False # "fence" 125 | 10: False # "bike" 126 | 11: False # "ground" 127 | 128 | 129 | split: # sequence numbers 130 | train: 131 | - 0 132 | - 1 133 | - 2 134 | - 4 135 | - 5 136 | valid: 137 | - 3 138 | test: 139 | - 3 -------------------------------------------------------------------------------- /utils/semkitti_vis/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/huixiancheng/PCB-RandNet/304c84b67380a23ab66516581b348c7db2f9f0df/utils/semkitti_vis/__init__.py -------------------------------------------------------------------------------- /utils/semkitti_vis/laserscan.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | import numpy as np 3 | 4 | 5 | class LaserScan: 6 | """Class that contains LaserScan with x,y,z,r""" 7 | EXTENSIONS_SCAN = ['.npy'] 8 | 9 | def __init__(self, project=False, H=64, W=1024, fov_up=3.0, fov_down=-25.0): 10 | self.project = project 11 | self.proj_H = H 12 | self.proj_W = W 13 | self.proj_fov_up = fov_up 14 | self.proj_fov_down = fov_down 15 | self.reset() 16 | 17 | def reset(self): 18 | """ Reset scan members. """ 19 | self.points = np.zeros((0, 3), dtype=np.float32) # [m, 3]: x, y, z 20 | self.remissions = np.zeros((0, 1), dtype=np.float32) # [m ,1]: remission 21 | 22 | # projected range image - [H,W] range (-1 is no data) 23 | self.proj_range = np.full((self.proj_H, self.proj_W), -1, 24 | dtype=np.float32) 25 | 26 | # unprojected range (list of depths for each point) 27 | self.unproj_range = np.zeros((0, 1), dtype=np.float32) 28 | 29 | # projected point cloud xyz - [H,W,3] xyz coord (-1 is no data) 30 | self.proj_xyz = np.full((self.proj_H, self.proj_W, 3), -1, 31 | dtype=np.float32) 32 | 33 | # projected remission - [H,W] intensity (-1 is no data) 34 | self.proj_remission = np.full((self.proj_H, self.proj_W), -1, 35 | dtype=np.float32) 36 | 37 | self.proj_idx = np.full((self.proj_H, self.proj_W), -1, 38 | dtype=np.int32) 39 | 40 | # for each point, where it is in the range image 41 | self.proj_x = np.zeros((0, 1), dtype=np.float32) # [m, 1]: x 42 | self.proj_y = np.zeros((0, 1), dtype=np.float32) # [m, 1]: y 43 | 44 | # mask containing for each pixel, if it contains a point or not 45 | self.proj_mask = np.zeros((self.proj_H, self.proj_W), 46 | dtype=np.int32) # [H,W] mask 47 | 48 | def size(self): 49 | """ Return the size of the point cloud. """ 50 | return self.points.shape[0] 51 | 52 | def __len__(self): 53 | return self.size() 54 | 55 | def open_scan(self, filename): 56 | """ Open raw scan and fill in attributes 57 | """ 58 | # reset just in case there was an open structure 59 | self.reset() 60 | 61 | # check filename is string 62 | if not isinstance(filename, str): 63 | raise TypeError("Filename should be string type, " 64 | "but was {type}".format(type=str(type(filename)))) 65 | 66 | # check extension is a laserscan 67 | if not any(filename.endswith(ext) for ext in self.EXTENSIONS_SCAN): 68 | raise RuntimeError("Filename extension is not valid scan file.") 69 | 70 | # if all goes well, open pointcloud 71 | scan = np.load(filename) 72 | 73 | # put in attribute 74 | points = scan # get xyz 75 | remissions = None # get remission 76 | self.set_points(points, remissions) 77 | 78 | def set_points(self, points, remissions=None): 79 | """ Set scan attributes (instead of opening from file) 80 | """ 81 | # reset just in case there was an open structure 82 | self.reset() 83 | 84 | # check scan makes sense 85 | if not isinstance(points, np.ndarray): 86 | raise TypeError("Scan should be numpy array") 87 | 88 | # check remission makes sense 89 | if remissions is not None and not isinstance(remissions, np.ndarray): 90 | raise TypeError("Remissions should be numpy array") 91 | 92 | # put in attribute 93 | self.points = points # get xyz 94 | if remissions is not None: 95 | self.remissions = remissions # get remission 96 | else: 97 | self.remissions = np.zeros((points.shape[0]), dtype=np.float32) 98 | 99 | # if projection is wanted, then do it and fill in the structure 100 | if self.project: 101 | self.do_range_projection() 102 | 103 | def do_range_projection(self): 104 | """ Project a pointcloud into a spherical projection image.projection. 105 | Function takes no arguments because it can be also called externally 106 | if the value of the constructor was not set (in case you change your 107 | mind about wanting the projection) 108 | """ 109 | # laser parameters 110 | fov_up = self.proj_fov_up / 180.0 * np.pi # field of view up in rad 111 | fov_down = self.proj_fov_down / 180.0 * np.pi # field of view down in rad 112 | fov = abs(fov_down) + abs(fov_up) # get field of view total in rad 113 | 114 | # get depth of all points 115 | depth = np.linalg.norm(self.points, 2, axis=1) 116 | 117 | # get scan components 118 | scan_x = self.points[:, 0] 119 | scan_y = self.points[:, 1] 120 | scan_z = self.points[:, 2] 121 | 122 | # get angles of all points 123 | yaw = -np.arctan2(scan_y, scan_x) 124 | pitch = np.arcsin(scan_z / depth) 125 | 126 | # get projections in image coords 127 | proj_x = 0.5 * (yaw / np.pi + 1.0) # in [0.0, 1.0] 128 | proj_y = 1.0 - (pitch + abs(fov_down)) / fov # in [0.0, 1.0] 129 | 130 | # scale to image size using angular resolution 131 | proj_x *= self.proj_W # in [0.0, W] 132 | proj_y *= self.proj_H # in [0.0, H] 133 | 134 | # round and clamp for use as index 135 | proj_x = np.floor(proj_x) 136 | proj_x = np.minimum(self.proj_W - 1, proj_x) 137 | proj_x = np.maximum(0, proj_x).astype(np.int32) # in [0,W-1] 138 | self.proj_x = np.copy(proj_x) # store a copy in orig order 139 | 140 | proj_y = np.floor(proj_y) 141 | proj_y = np.minimum(self.proj_H - 1, proj_y) 142 | proj_y = np.maximum(0, proj_y).astype(np.int32) # in [0,H-1] 143 | self.proj_y = np.copy(proj_y) # stope a copy in original order 144 | 145 | # copy of depth in original order 146 | self.unproj_range = np.copy(depth) 147 | 148 | # order in decreasing depth 149 | indices = np.arange(depth.shape[0]) 150 | order = np.argsort(depth)[::-1] 151 | depth = depth[order] 152 | indices = indices[order] 153 | points = self.points[order] 154 | remission = self.remissions[order] 155 | proj_y = proj_y[order] 156 | proj_x = proj_x[order] 157 | 158 | # assing to images 159 | self.proj_range[proj_y, proj_x] = depth 160 | self.proj_xyz[proj_y, proj_x] = points 161 | self.proj_remission[proj_y, proj_x] = remission 162 | self.proj_idx[proj_y, proj_x] = indices 163 | self.proj_mask = (self.proj_idx > 0).astype(np.float32) 164 | 165 | 166 | class SemLaserScan(LaserScan): 167 | """Class that contains LaserScan with x,y,z,r,sem_label,sem_color_label,inst_label,inst_color_label""" 168 | EXTENSIONS_LABEL = ['.label'] 169 | 170 | def __init__(self, nclasses, sem_color_dict=None, project=False, H=64, W=1024, fov_up=3.0, fov_down=-25.0): 171 | super(SemLaserScan, self).__init__(project, H, W, fov_up, fov_down) 172 | self.reset() 173 | self.nclasses = nclasses # number of classes 174 | 175 | # make semantic colors 176 | max_sem_key = 0 177 | for key, data in sem_color_dict.items(): 178 | if key + 1 > max_sem_key: 179 | max_sem_key = key + 1 180 | self.sem_color_lut = np.zeros((max_sem_key + 100, 3), dtype=np.float32) 181 | for key, value in sem_color_dict.items(): 182 | self.sem_color_lut[key] = np.array(value, np.float32) / 255.0 183 | 184 | # make instance colors 185 | max_inst_id = 100000 186 | self.inst_color_lut = np.random.uniform(low=0.0, high=1.0, size=(max_inst_id, 3)) 187 | # force zero to a gray-ish color 188 | self.inst_color_lut[0] = np.full((3), 0.1) 189 | 190 | def reset(self): 191 | """ Reset scan members. """ 192 | super(SemLaserScan, self).reset() 193 | 194 | # semantic labels 195 | self.sem_label = np.zeros((0, 1), dtype=np.uint32) # [m, 1]: label 196 | self.sem_label_color = np.zeros((0, 3), dtype=np.float32) # [m ,3]: color 197 | 198 | # instance labels 199 | self.inst_label = np.zeros((0, 1), dtype=np.uint32) # [m, 1]: label 200 | self.inst_label_color = np.zeros((0, 3), dtype=np.float32) # [m ,3]: color 201 | 202 | # projection color with semantic labels 203 | self.proj_sem_label = np.zeros((self.proj_H, self.proj_W), 204 | dtype=np.int32) # [H,W] label 205 | self.proj_sem_color = np.zeros((self.proj_H, self.proj_W, 3), 206 | dtype=np.float) # [H,W,3] color 207 | 208 | # projection color with instance labels 209 | self.proj_inst_label = np.zeros((self.proj_H, self.proj_W), 210 | dtype=np.int32) # [H,W] label 211 | self.proj_inst_color = np.zeros((self.proj_H, self.proj_W, 3), 212 | dtype=np.float) # [H,W,3] color 213 | 214 | def open_label(self, filename): 215 | """ Open raw scan and fill in attributes 216 | """ 217 | # check filename is string 218 | if not isinstance(filename, str): 219 | raise TypeError("Filename should be string type, " 220 | "but was {type}".format(type=str(type(filename)))) 221 | 222 | # check extension is a laserscan 223 | if not any(filename.endswith(ext) for ext in self.EXTENSIONS_LABEL): 224 | raise RuntimeError("Filename extension is not valid label file.") 225 | 226 | # if all goes well, open label 227 | label = np.fromfile(filename, dtype=np.uint32) 228 | 229 | # set it 230 | self.set_label(label) 231 | 232 | def set_label(self, label): 233 | """ Set points for label not from file but from np 234 | """ 235 | # check label makes sense 236 | if not isinstance(label, np.ndarray): 237 | raise TypeError("Label should be numpy array") 238 | 239 | # only fill in attribute if the right size 240 | if label.shape[0] == self.points.shape[0]: 241 | self.sem_label = label & 0xFFFF # semantic label in lower half 242 | self.inst_label = label >> 16 # instance id in upper half 243 | else: 244 | print("Points shape: ", self.points.shape) 245 | print("Label shape: ", label.shape) 246 | raise ValueError("Scan and Label don't contain same number of points") 247 | 248 | # sanity check 249 | assert((self.sem_label + (self.inst_label << 16) == label).all()) 250 | 251 | if self.project: 252 | self.do_label_projection() 253 | 254 | def colorize(self): 255 | """ Colorize pointcloud with the color of each semantic label 256 | """ 257 | self.sem_label_color = self.sem_color_lut[self.sem_label] 258 | self.sem_label_color = self.sem_label_color.reshape((-1, 3)) 259 | 260 | self.inst_label_color = self.inst_color_lut[self.inst_label] 261 | self.inst_label_color = self.inst_label_color.reshape((-1, 3)) 262 | 263 | def do_label_projection(self): 264 | # only map colors to labels that exist 265 | mask = self.proj_idx >= 0 266 | 267 | # semantics 268 | self.proj_sem_label[mask] = self.sem_label[self.proj_idx[mask]] 269 | self.proj_sem_color[mask] = self.sem_color_lut[self.sem_label[self.proj_idx[mask]]] 270 | 271 | # instances 272 | self.proj_inst_label[mask] = self.inst_label[self.proj_idx[mask]] 273 | self.proj_inst_color[mask] = self.inst_color_lut[self.inst_label[self.proj_idx[mask]]] 274 | -------------------------------------------------------------------------------- /utils/semkitti_vis/laserscanvis.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # This file is covered by the LICENSE file in the root of this project. 3 | 4 | import vispy 5 | from vispy.scene import visuals, SceneCanvas 6 | import numpy as np 7 | from matplotlib import pyplot as plt 8 | 9 | 10 | class LaserScanVis: 11 | """Class that creates and handles a visualizer for a pointcloud""" 12 | 13 | def __init__(self, scan, scan_names, label_names, offset=0, 14 | semantics=True, instances=False): 15 | self.scan = scan 16 | self.scan_names = scan_names 17 | self.label_names = label_names 18 | self.offset = offset 19 | self.total = len(self.scan_names) 20 | self.semantics = semantics 21 | self.instances = instances 22 | # sanity check 23 | if not self.semantics and self.instances: 24 | print("Instances are only allowed in when semantics=True") 25 | raise ValueError 26 | 27 | self.reset() 28 | self.update_scan() 29 | 30 | def reset(self): 31 | """ Reset. """ 32 | # last key press (it should have a mutex, but visualization is not 33 | # safety critical, so let's do things wrong) 34 | self.action = "no" # no, next, back, quit are the possibilities 35 | 36 | # new canvas prepared for visualizing data 37 | self.canvas = SceneCanvas(keys='interactive', show=True) 38 | # interface (n next, b back, q quit, very simple) 39 | self.canvas.events.key_press.connect(self.key_press) 40 | self.canvas.events.draw.connect(self.draw) 41 | # grid 42 | self.grid = self.canvas.central_widget.add_grid() 43 | 44 | # laserscan part 45 | self.scan_view = vispy.scene.widgets.ViewBox( 46 | border_color='white', parent=self.canvas.scene) 47 | self.grid.add_widget(self.scan_view, 0, 0) 48 | self.scan_vis = visuals.Markers() 49 | self.scan_view.camera = 'turntable' 50 | self.scan_view.add(self.scan_vis) 51 | visuals.XYZAxis(parent=self.scan_view.scene) 52 | # add semantics 53 | print("Using semantics in visualizer") 54 | self.sem_view = vispy.scene.widgets.ViewBox( 55 | border_color='white', parent=self.canvas.scene) 56 | self.grid.add_widget(self.sem_view, 0, 1) 57 | self.sem_vis = visuals.Markers() 58 | self.sem_view.camera = 'turntable' 59 | self.sem_view.add(self.sem_vis) 60 | visuals.XYZAxis(parent=self.sem_view.scene) 61 | # self.sem_view.camera.link(self.scan_view.camera) 62 | 63 | # img canvas size 64 | self.multiplier = 1 65 | self.canvas_W = 1024 66 | self.canvas_H = 64 67 | if self.semantics: 68 | self.multiplier += 1 69 | if self.instances: 70 | self.multiplier += 1 71 | 72 | # new canvas for img 73 | self.img_canvas = SceneCanvas(keys='interactive', show=True, 74 | size=(self.canvas_W, self.canvas_H * self.multiplier)) 75 | # grid 76 | self.img_grid = self.img_canvas.central_widget.add_grid() 77 | # interface (n next, b back, q quit, very simple) 78 | self.img_canvas.events.key_press.connect(self.key_press) 79 | self.img_canvas.events.draw.connect(self.draw) 80 | 81 | # add a view for the depth 82 | self.img_view = vispy.scene.widgets.ViewBox( 83 | border_color='white', parent=self.img_canvas.scene) 84 | self.img_grid.add_widget(self.img_view, 0, 0) 85 | self.img_vis = visuals.Image(cmap='viridis') 86 | self.img_view.add(self.img_vis) 87 | 88 | # add semantics 89 | self.sem_img_view = vispy.scene.widgets.ViewBox( 90 | border_color='white', parent=self.img_canvas.scene) 91 | self.img_grid.add_widget(self.sem_img_view, 1, 0) 92 | self.sem_img_vis = visuals.Image(cmap='viridis') 93 | self.sem_img_view.add(self.sem_img_vis) 94 | 95 | def get_mpl_colormap(self, cmap_name): 96 | cmap = plt.get_cmap(cmap_name) 97 | 98 | # Initialize the matplotlib color map 99 | sm = plt.cm.ScalarMappable(cmap=cmap) 100 | 101 | # Obtain linear color range 102 | color_range = sm.to_rgba(np.linspace(0, 1, 256), bytes=True)[:, 2::-1] 103 | 104 | return color_range.reshape(256, 3).astype(np.float32) / 255.0 105 | 106 | def update_scan(self): 107 | # first open data 108 | self.scan.open_scan(self.scan_names[self.offset]) 109 | self.scan.open_label(self.label_names[self.offset]) 110 | self.scan.colorize() 111 | 112 | # then change names 113 | title = "scan " + str(self.offset) 114 | self.canvas.title = title 115 | self.img_canvas.title = title 116 | 117 | # then do all the point cloud stuff 118 | 119 | # plot scan 120 | power = 16 121 | # print() 122 | range_data = np.copy(self.scan.unproj_range) 123 | # print(range_data.max(), range_data.min()) 124 | range_data = range_data**(1 / power) 125 | # print(range_data.max(), range_data.min()) 126 | viridis_range = ((range_data - range_data.min()) / 127 | (range_data.max() - range_data.min()) * 128 | 255).astype(np.uint8) 129 | viridis_map = self.get_mpl_colormap("viridis") 130 | viridis_colors = viridis_map[viridis_range] 131 | self.scan_vis.set_data(self.scan.points, 132 | face_color=viridis_colors[..., ::-1], 133 | edge_color=viridis_colors[..., ::-1], 134 | size=1) 135 | 136 | # plot semantics 137 | self.sem_vis.set_data(self.scan.points, 138 | face_color=self.scan.sem_label_color[..., ::-1], 139 | edge_color=self.scan.sem_label_color[..., ::-1], 140 | size=1) 141 | 142 | # now do all the range image stuff 143 | # plot range image 144 | data = np.copy(self.scan.proj_range) 145 | # print(data[data > 0].max(), data[data > 0].min()) 146 | data[data > 0] = data[data > 0]**(1 / power) 147 | data[data < 0] = data[data > 0].min() 148 | # print(data.max(), data.min()) 149 | data = (data - data[data > 0].min()) / \ 150 | (data.max() - data[data > 0].min()) 151 | # print(data.max(), data.min()) 152 | self.img_vis.set_data(data) 153 | self.img_vis.update() 154 | 155 | self.sem_img_vis.set_data(self.scan.proj_sem_color[..., ::-1]) 156 | self.sem_img_vis.update() 157 | 158 | # interface 159 | def key_press(self, event): 160 | self.canvas.events.key_press.block() 161 | self.img_canvas.events.key_press.block() 162 | if event.key == 'N': 163 | self.offset += 1 164 | if self.offset >= self.total: 165 | self.offset = 0 166 | self.update_scan() 167 | elif event.key == 'B': 168 | self.offset -= 1 169 | if self.offset < 0: 170 | self.offset = self.total - 1 171 | self.update_scan() 172 | elif event.key == 'Q' or event.key == 'Escape': 173 | self.destroy() 174 | 175 | def draw(self, event): 176 | if self.canvas.events.key_press.blocked(): 177 | self.canvas.events.key_press.unblock() 178 | if self.img_canvas.events.key_press.blocked(): 179 | self.img_canvas.events.key_press.unblock() 180 | 181 | def destroy(self): 182 | # destroy the visualization 183 | self.canvas.close() 184 | self.img_canvas.close() 185 | vispy.app.quit() 186 | 187 | def run(self): 188 | vispy.app.run() 189 | -------------------------------------------------------------------------------- /visualize_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 argparse 5 | import os 6 | import yaml 7 | from utils.semkitti_vis.laserscan import SemLaserScan 8 | from utils.semkitti_vis.laserscanvis import LaserScanVis 9 | 10 | if __name__ == '__main__': 11 | parser = argparse.ArgumentParser("./visualize.py") 12 | parser.add_argument( 13 | '--dataset', '-d', 14 | type=str, 15 | required=True, 16 | help='Dataset to visualize. No Default', 17 | ) 18 | parser.add_argument( 19 | '--config', '-c', 20 | type=str, 21 | required=False, 22 | default="utils/semantic-kitti.yaml", 23 | help='Dataset config file. Defaults to %(default)s', 24 | ) 25 | parser.add_argument( 26 | '--sequence', '-s', 27 | type=str, 28 | default="00", 29 | required=False, 30 | help='Sequence to visualize. Defaults to %(default)s', 31 | ) 32 | parser.add_argument( 33 | '--predictions', '-p', 34 | type=str, 35 | default=None, 36 | required=False, 37 | help='Alternate location for labels, to use predictions folder. ' 38 | 'Must point to directory containing the predictions in the proper format ' 39 | ' (see readme)' 40 | 'Defaults to %(default)s', 41 | ) 42 | FLAGS, unparsed = parser.parse_known_args() 43 | 44 | # print summary of what we will do 45 | print("*" * 80) 46 | print("INTERFACE:") 47 | print("Dataset", FLAGS.dataset) 48 | print("Config", FLAGS.config) 49 | print("Sequence", FLAGS.sequence) 50 | print("Predictions", FLAGS.predictions) 51 | print("*" * 80) 52 | 53 | # open config file 54 | print("Opening config file %s" % FLAGS.config) 55 | CFG = yaml.safe_load(open(FLAGS.config, 'r')) 56 | 57 | # fix sequence name 58 | FLAGS.sequence = '{0:02d}'.format(int(FLAGS.sequence)) 59 | 60 | # does sequence folder exist? 61 | scan_paths = os.path.join(FLAGS.dataset, FLAGS.sequence, "velodyne") 62 | if os.path.isdir(scan_paths): 63 | print("Sequence folder exists! Using sequence from %s" % scan_paths) 64 | else: 65 | print("Sequence folder doesn't exist! Exiting...") 66 | quit() 67 | 68 | # populate the pointclouds 69 | scan_names = [os.path.join(dp, f) for dp, dn, fn in os.walk( 70 | os.path.expanduser(scan_paths)) for f in fn] 71 | scan_names.sort() 72 | 73 | # does sequence folder exist? 74 | if FLAGS.predictions is not None: 75 | label_paths = os.path.join(FLAGS.predictions, FLAGS.sequence, "predictions") 76 | else: 77 | label_paths = os.path.join(FLAGS.dataset, FLAGS.sequence, "labels") 78 | if os.path.isdir(label_paths): 79 | print("Labels folder exists! Using labels from %s" % label_paths) 80 | else: 81 | print("Labels folder doesn't exist! Exiting...") 82 | quit() 83 | # populate the pointclouds 84 | label_names = [os.path.join(dp, f) for dp, dn, fn in os.walk( 85 | os.path.expanduser(label_paths)) for f in fn] 86 | label_names.sort() 87 | 88 | color_dict = CFG["color_map"] 89 | nclasses = len(color_dict) 90 | scan = SemLaserScan(nclasses, color_dict, project=True) 91 | 92 | vis = LaserScanVis(scan=scan, 93 | scan_names=scan_names, 94 | label_names=label_names, 95 | offset=0, 96 | semantics=True, instances=False) 97 | 98 | # print instructions 99 | print("To navigate:") 100 | print("\tb: back (previous scan)") 101 | print("\tn: next (next scan)") 102 | print("\tq: quit (exit program)") 103 | 104 | # run the visualizer 105 | vis.run() 106 | --------------------------------------------------------------------------------