├── .gitignore ├── LICENSE ├── README.md ├── README_zh.md ├── data.py ├── image ├── DGCNN.jpg ├── partseg_colors.png ├── partseg_visu.png ├── semseg_s3dis_colors.png ├── semseg_s3dis_visu.png ├── semseg_scannet_colors.png └── semseg_scannet_visu.png ├── main_cls.py ├── main_partseg.py ├── main_semseg_s3dis.py ├── main_semseg_scannet.py ├── model.py ├── prepare_data ├── collect_indoor3d_data.py ├── data_prep_util.py ├── gen_indoor3d_h5.py ├── indoor3d_util.py ├── meta │ ├── all_data_label.txt │ ├── anno_paths.txt │ ├── class_names.txt │ ├── partseg_colors.txt │ └── semseg_colors.txt ├── scannetv2_seg_dataset_rgb21c_pointid.py ├── scannetv2_test.txt ├── scannetv2_train.txt ├── scannetv2_val.txt └── util.py ├── pretrained ├── model.cls.1024.t7 ├── model.cls.2048.t7 ├── model.partseg.airplane.t7 ├── model.partseg.car.t7 ├── model.partseg.t7 ├── semseg_s3dis │ ├── model_1.t7 │ ├── model_2.t7 │ ├── model_3.t7 │ ├── model_4.t7 │ ├── model_5.t7 │ └── model_6.t7 └── semseg_scannet │ └── models │ └── model_200.pth ├── sync_bn ├── __init__.py ├── batchnorm.py ├── comm.py ├── replicate.py └── unittest.py └── util.py /.gitignore: -------------------------------------------------------------------------------- 1 | .DS_Store 2 | */.DS_Store 3 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2020 An Tao 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 | # DGCNN.pytorch 2 | [[中文版]](README_zh.md) 3 | 4 | This repo is a PyTorch implementation for **Dynamic Graph CNN for Learning on Point Clouds (DGCNN)**(https://arxiv.org/pdf/1801.07829). Our code skeleton is borrowed from [WangYueFt/dgcnn](https://github.com/WangYueFt/dgcnn/tree/master/pytorch). 5 | 6 | **Updates:** 7 | 8 | - [2022/10/22] Add semantic segmentation on the ScanNet dataset by [Ziyi Wu](https://wuziyi616.github.io/) (dazitu616@gmail.com) and [Yingqi Wang](https://github.com/HilbertWang2002) (yingqi-w19@mails.tsinghua.edu.cn). 9 | - [2021/07/20] Add visualization code for part segmentation and semantic segmentation by [Pengliang Ji](https://github.com/Ji-Pengliang) (jpl1723@buaa.edu.cn). 10 | 11 | The network structure (Fig. 3) for classification in DGCNN paper is not consistent with the corresponding description in section 4.1 of the paper. The author of DGCNN adopts the setting of classification network in section 4.1, not Fig. 3. We fixed this mistake in Fig. 3 using PS and present the revised figure below. 12 | 13 |   14 |

15 | 16 |

17 | 18 |   19 | 20 | **Tip:** The result of point cloud experiment usually faces greater randomness than 2D image. We suggest you run your experiment more than one time and select the best result. 21 | 22 |   23 | ## Requirements 24 | - Python >= 3.7 25 | - PyTorch >= 1.2 26 | - CUDA >= 10.0 27 | - Package: glob, h5py, sklearn, plyfile, torch_scatter 28 | 29 |   30 | ## Contents 31 | - [Point Cloud Classification](#point-cloud-classification) 32 | - [Point Cloud Part Segmentation](#point-cloud-part-segmentation) 33 | - [Point Cloud Semantic Segmentation on the S3DIS Dataset](#point-cloud-semantic-segmentation-on-the-s3dis-dataset) 34 | - [Point Cloud Semantic Segmentation on the ScanNet Dataset](#point-cloud-semantic-segmentation-on-the-scannet-dataset) 35 | 36 | **Note:** All following commands default use all GPU cards. To specify the cards to use, add `CUDA_VISIBLE_DEVICES=0,1,2,3` before each command, where the user uses 4 GPU cards with card index `0,1,2,3`. You can change the card number and indexes depending on your own needs. 37 | 38 |   39 | ## Point Cloud Classification 40 | ### Run the training script: 41 | 42 | - 1024 points 43 | 44 | ``` 45 | python main_cls.py --exp_name=cls_1024 --num_points=1024 --k=20 46 | ``` 47 | 48 | - 2048 points 49 | 50 | ``` 51 | python main_cls.py --exp_name=cls_2048 --num_points=2048 --k=40 52 | ``` 53 | 54 | ### Run the evaluation script after training finished: 55 | 56 | - 1024 points 57 | 58 | ``` 59 | python main_cls.py --exp_name=cls_1024_eval --num_points=1024 --k=20 --eval=True --model_path=outputs/cls_1024/models/model.t7 60 | ``` 61 | 62 | - 2048 points 63 | 64 | ``` 65 | python main_cls.py --exp_name=cls_2048_eval --num_points=2048 --k=40 --eval=True --model_path=outputs/cls_2048/models/model.t7 66 | ``` 67 | 68 | ### Run the evaluation script with pretrained models: 69 | 70 | - 1024 points 71 | 72 | ``` 73 | python main_cls.py --exp_name=cls_1024_eval --num_points=1024 --k=20 --eval=True --model_path=pretrained/model.cls.1024.t7 74 | ``` 75 | 76 | - 2048 points 77 | 78 | ``` 79 | python main_cls.py --exp_name=cls_2048_eval --num_points=2048 --k=40 --eval=True --model_path=pretrained/model.cls.2048.t7 80 | ``` 81 | 82 | ### Performance: 83 | ModelNet40 dataset 84 | 85 | | | Mean Class Acc | Overall Acc | 86 | | :---: | :---: | :---: | 87 | | Paper (1024 points) | 90.2 | 92.9 | 88 | | This repo (1024 points) | **90.9** | **93.3** | 89 | | Paper (2048 points) | 90.7 | 93.5 | 90 | | This repo (2048 points) | **91.2** | **93.6** | 91 | 92 |   93 | ## Point Cloud Part Segmentation 94 | **Note:** The training modes **'full dataset'** and **'with class choice'** are different. 95 | 96 | - In **'full dataset'**, the model is trained and evaluated in all 16 classes and outputs mIoU 85.2% in this repo. The prediction of points in each shape can be any part of all 16 classes. 97 | - In **'with class choice'**, the model is trained and evaluated in one class, for example airplane, and outputs mIoU 84.5% for airplane in this repo. The prediction of points in each shape can only be one of the parts in this chosen class. 98 | 99 | ### Run the training script: 100 | 101 | - Full dataset 102 | 103 | ``` 104 | python main_partseg.py --exp_name=partseg 105 | ``` 106 | 107 | - With class choice, for example airplane 108 | 109 | ``` 110 | python main_partseg.py --exp_name=partseg_airplane --class_choice=airplane 111 | ``` 112 | 113 | ### Run the evaluation script after training finished: 114 | 115 | - Full dataset 116 | 117 | ``` 118 | python main_partseg.py --exp_name=partseg_eval --eval=True --model_path=outputs/partseg/models/model.t7 119 | ``` 120 | 121 | - With class choice, for example airplane 122 | 123 | ``` 124 | python main_partseg.py --exp_name=partseg_airplane_eval --class_choice=airplane --eval=True --model_path=outputs/partseg_airplane/models/model.t7 125 | ``` 126 | 127 | ### Run the evaluation script with pretrained models: 128 | 129 | - Full dataset 130 | 131 | ``` 132 | python main_partseg.py --exp_name=partseg_eval --eval=True --model_path=pretrained/model.partseg.t7 133 | ``` 134 | 135 | - With class choice, for example airplane 136 | 137 | ``` 138 | python main_partseg.py --exp_name=partseg_airplane_eval --class_choice=airplane --eval=True --model_path=pretrained/model.partseg.airplane.t7 139 | ``` 140 | 141 | ### Performance: 142 | 143 | ShapeNet part dataset 144 | 145 | | | Mean IoU | Airplane | Bag | Cap | Car | Chair | Earphone | Guitar | Knife | Lamp | Laptop | Motor | Mug | Pistol | Rocket | Skateboard | Table 146 | | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | 147 | | Shapes | | 2690 | 76 | 55 | 898 | 3758 | 69 | 787 | 392 | 1547 | 451 | 202 | 184 | 283 | 66 | 152 | 5271 | 148 | | Paper | **85.2** | 84.0 | **83.4** | **86.7** | 77.8 | 90.6 | 74.7 | 91.2 | **87.5** | 82.8 | **95.7** | 66.3 | **94.9** | 81.1 | **63.5** | 74.5 | 82.6 | 149 | | This repo | **85.2** | **84.5** | 80.3 | 84.7 | **79.8** | **91.1** | **76.8** | **92.0** | 87.3 | **83.8** | **95.7** | **69.6** | 94.3 | **83.7** | 51.5 | **76.1** | **82.8** | 150 | 151 | ### Visualization: 152 | #### Usage: 153 | 154 | Use `--visu` to control visualization file. 155 | 156 | - To visualize a single shape, for example the 0-th airplane (the shape index starts from 0), use `--visu=airplane_0`. 157 | - To visualize all shapes in a class, for example airplane, use `--visu=airplane`. 158 | - To visualize all shapes in all classes, use `--visu=all`. 159 | 160 | Use `--visu_format` to control visualization file format. 161 | 162 | - To output .txt file, use `--visu_format=txt`. 163 | - To output .ply format, use `--visu_format=ply`. 164 | 165 | Both .txt and .ply file can be loaded into [MeshLab](https://www.meshlab.net) for visualization. For the usage of MeshLab on .txt file, see issue [#8](https://github.com/AnTao97/dgcnn.pytorch/issues/8) for details. The .ply file can be directly loaded into MeshLab by dragging. 166 | 167 | The visualization file name follows the format `shapename_pred_miou.FILE_TYPE` for prediction output or `shapename_gt.FILE_TYPE` for ground-truth, where `miou` shows the mIoU prediction for this shape. 168 | 169 | #### Full dataset: 170 | 171 | - Output the visualization file of the 0-th airplane with .ply format 172 | 173 | ``` 174 | # Use trained model 175 | python main_partseg.py --exp_name=partseg_eval --eval=True --model_path=outputs/partseg/models/model.t7 --visu=airplane_0 --visu_format=ply 176 | 177 | # Use pretrained model 178 | python main_partseg.py --exp_name=partseg_eval --eval=True --model_path=pretrained/model.partseg.t7 --visu=airplane_0 --visu_format=ply 179 | ``` 180 | 181 | - Output the visualization files of all shapes in airplane class with .ply format 182 | 183 | ``` 184 | # Use trained model 185 | python main_partseg.py --exp_name=partseg_eval --eval=True --model_path=outputs/partseg/models/model.t7 --visu=airplane --visu_format=ply 186 | 187 | # Use pretrained model 188 | python main_partseg.py --exp_name=partseg_eval --eval=True --model_path=pretrained/model.partseg.t7 --visu=airplane --visu_format=ply 189 | ``` 190 | 191 | - Output the visualization files of all shapes in all classes with .ply format 192 | 193 | ``` 194 | # Use trained model 195 | python main_partseg.py --exp_name=partseg_eval --eval=True --model_path=outputs/partseg/models/model.t7 --visu=all --visu_format=ply 196 | 197 | # Use pretrained model 198 | python main_partseg.py --exp_name=partseg_eval --eval=True --model_path=pretrained/model.partseg.t7 --visu=all --visu_format=ply 199 | ``` 200 | 201 | #### With class choice, for example airplane: 202 | 203 | - Output the visualization file of the 0-th airplane with .ply format 204 | 205 | ``` 206 | # Use trained model 207 | python main_partseg.py --exp_name=partseg_airplane_eval --class_choice=airplane --eval=True --model_path=outputs/partseg_airplane/models/model.t7 --visu=airplane_0 --visu_format=ply 208 | 209 | # Use pretrained model 210 | python main_partseg.py --exp_name=partseg_airplane_eval --class_choice=airplane --eval=True --model_path=pretrained/model.partseg.airplane.t7 --visu=airplane_0 --visu_format=ply 211 | ``` 212 | 213 | - Output the visualization files of all shapes in airplane class with .ply format 214 | 215 | ``` 216 | # Use trained model 217 | python main_partseg.py --exp_name=partseg_airplane_eval --class_choice=airplane --eval=True --model_path=outputs/partseg_airplane/models/model.t7 --visu=airplane --visu_format=ply 218 | 219 | # Use pretrained model 220 | python main_partseg.py --exp_name=partseg_airplane_eval --class_choice=airplane --eval=True --model_path=pretrained/model.partseg.airplane.t7 --visu=airplane --visu_format=ply 221 | ``` 222 | 223 | #### Results: 224 | The visualization result of the airplane 0: 225 | 226 |

227 | 228 |

229 | 230 | Color map: 231 |

232 | 233 |

234 | 235 |   236 | ## Point Cloud Semantic Segmentation on the S3DIS Dataset 237 | 238 | The network structure for this task is slightly different with part segmentation, without spatial transform and categorical vector. The MLP in the end is changed into (512, 256, 13) and only one dropout is used after 256. 239 | 240 | You have to download `Stanford3dDataset_v1.2_Aligned_Version.zip` manually from https://goo.gl/forms/4SoGp4KtH1jfRqEj2 and place it under `data/` 241 | 242 | ### Run the training script: 243 | 244 | This task uses 6-fold training, such that 6 models are trained leaving 1 of 6 areas as the testing area for each model. 245 | 246 | - Train in area 1-5 247 | 248 | ``` 249 | python main_semseg_s3dis.py --exp_name=semseg_s3dis_6 --test_area=6 250 | ``` 251 | 252 | ### Run the evaluation script after training finished: 253 | 254 | - Evaluate in area 6 after the model is trained in area 1-5 255 | 256 | ``` 257 | python main_semseg_s3dis.py --exp_name=semseg_s3dis_eval_6 --test_area=6 --eval=True --model_root=outputs/semseg_s3dis/models/ 258 | ``` 259 | 260 | - Evaluate in all areas after 6 models are trained 261 | 262 | ``` 263 | python main_semseg_s3dis.py --exp_name=semseg_s3dis_eval --test_area=all --eval=True --model_root=outputs/semseg_s3dis/models/ 264 | ``` 265 | 266 | ### Run the evaluation script with pretrained models: 267 | 268 | - Evaluate in area 6 269 | 270 | ``` 271 | python main_semseg_s3dis.py --exp_name=semseg_s3dis_eval_6 --test_area=6 --eval=True --model_root=pretrained/semseg_s3dis/ 272 | ``` 273 | 274 | - Evaluate in all areas 275 | 276 | ``` 277 | python main_semseg_s3dis.py --exp_name=semseg_s3dis_eval --test_area=all --eval=True --model_root=pretrained/semseg_s3dis/ 278 | ``` 279 | 280 | ### Performance: 281 | 282 | Stanford Large-Scale 3D Indoor Spaces Dataset (S3DIS) dataset 283 | 284 | | | Mean IoU | Overall Acc | 285 | | :---: | :---: | :---: | 286 | | Paper | 56.1 | 84.1 | 287 | | This repo | **59.2** | **85.0** | 288 | 289 | ### Visualization: 290 | #### Usage: 291 | 292 | Use `--visu` to control visualization file. 293 | 294 | - To visualize a single room, for example the office room 1 in area 6 (the room index starts from 1), use `--visu=area_6_office_1`. 295 | - To visualize all rooms in an area, for example area 6, use `--visu=area_6`. 296 | - To visualize all rooms in all areas, use `--visu=all`. 297 | 298 | Use `--visu_format` to control visualization file format. 299 | 300 | - To output .txt file, use `--visu_format=txt`. 301 | - To output .ply format, use `--visu_format=ply`. 302 | 303 | Both .txt and .ply file can be loaded into [MeshLab](https://www.meshlab.net) for visualization. For the usage of MeshLab on .txt file, see issue [#8](https://github.com/AnTao97/dgcnn.pytorch/issues/8) for details. The .ply file can be directly loaded into MeshLab by dragging. 304 | 305 | The visualization file name follows the format `roomname_pred_.FILE_TYPE` for prediction output or `roomname_gt.FILE_TYPE` for ground-truth, where `` shows the mIoU prediction for this room. 306 | 307 | **Note:** In semantic segmentation, you need to first run a training or evaluation command without visualization in the above sections to preprocess dataset. With the dataset well prepared, you can run a command with visualization in the following sections. 308 | 309 | #### Evaluate in area 6 after the model is trained in area 1-5: 310 | 311 | - Output the visualization file of office room 1 in area 6 with .ply format 312 | 313 | ``` 314 | # Use trained model 315 | python main_semseg_s3dis.py --exp_name=semseg_s3dis_eval_6 --test_area=6 --eval=True --model_root=outputs/semseg_s3dis/models/ --visu=area_6_office_1 --visu_format=ply 316 | 317 | # Use pretrained model 318 | python main_semseg_s3dis.py --exp_name=semseg_s3dis_eval_6 --test_area=6 --eval=True --model_root=pretrained/semseg_s3dis/ --visu=area_6_office_1 --visu_format=ply 319 | ``` 320 | 321 | - Output the visualization files of all rooms in area 6 with .ply format 322 | 323 | ``` 324 | # Use trained model 325 | python main_semseg_s3dis.py --exp_name=semseg_s3dis_eval_6 --test_area=6 --eval=True --model_root=outputs/semseg_s3dis/models/ --visu=area_6 --visu_format=ply 326 | 327 | # Use pretrained model 328 | python main_semseg_s3dis.py --exp_name=semseg_s3dis_eval_6 --test_area=6 --eval=True --model_root=pretrained/semseg_s3dis/ --visu=area_6 --visu_format=ply 329 | ``` 330 | 331 | #### Evaluate in all areas after 6 models are trained: 332 | 333 | - Output the visualization file of office room 1 in area 6 with .ply format 334 | 335 | 336 | ``` 337 | # Use trained model 338 | python main_semseg_s3dis.py --exp_name=semseg_s3dis_eval --test_area=all --eval=True --model_root=outputs/semseg_s3dis/models/ --visu=area_6_office_1 --visu_format=ply 339 | 340 | # Use pretrained model 341 | python main_semseg_s3dis.py --exp_name=semseg_s3dis_eval --test_area=all --eval=True --model_root=pretrained/semseg_s3dis/ --visu=area_6_office_1 --visu_format=ply 342 | ``` 343 | 344 | - Output the visualization files of all rooms in area 6 with .ply format 345 | 346 | ``` 347 | # Use trained model 348 | python main_semseg_s3dis.py --exp_name=semseg_s3dis_eval --test_area=all --eval=True --model_root=outputs/semseg_s3dis/models/ --visu=area_6 --visu_format=ply 349 | 350 | # Use pretrained model 351 | python main_semseg_s3dis.py --exp_name=semseg_s3dis_eval --test_area=all --eval=True --model_root=pretrained/semseg_s3dis/ --visu=area_6 --visu_format=ply 352 | ``` 353 | 354 | - Output the visualization files of all rooms in all areas with .ply format 355 | 356 | ``` 357 | # Use trained model 358 | python main_semseg_s3dis.py --exp_name=semseg_s3dis_eval --test_area=all --eval=True --model_root=outputs/semseg_s3dis/models/ --visu=all --visu_format=ply 359 | 360 | # Use pretrained model 361 | python main_semseg_s3dis.py --exp_name=semseg_s3dis_eval --test_area=all --eval=True --model_root=pretrained/semseg_s3dis/ --visu=all --visu_format=ply 362 | ``` 363 | 364 | #### Results: 365 | The visualization result of the office room 1 in area 6: 366 | 367 |

368 | 369 |

370 | 371 | Color map: 372 |

373 | 374 |

375 | 376 |   377 | ## Point Cloud Semantic Segmentation on the ScanNet Dataset 378 | 379 | The DGCNN authors do not test on the ScanNet dataset. We try our best to implement the DGCNN model on the dataset. 380 | 381 | ### Prepare dataset: 382 | 383 | You need to change the directory to the `prepare_data/` folder. 384 | 385 | ``` 386 | cd prepare_data/ 387 | ``` 388 | 389 | Please download original dataset from [website](http://www.scan-net.org/). You need to place the dataset under `data/ScanNet/`. The path `data/ScanNet` includes `data/ScanNet/scans/` and `data/ScanNet/scans_test/` folder. 390 | 391 | To prepare the Scannet dataset for training and evaluation, run 392 | 393 | ``` 394 | python scannetv2_seg_dataset_rgb21c_pointid.py 395 | ``` 396 | 397 | This will generate four pickle files: `scannet_train_rgb21c_pointid.pickle`, `scannet_val_rgb21c_pointid.pickle`, `scannet_val_rgb21c_pointid_keep_unanno.pickle`, and `scannet_test_rgb21c_pointid_keep_unanno.pickle`. 398 | 399 | Return to the root directory: 400 | 401 | ``` 402 | cd .. 403 | ``` 404 | 405 | ### Run the training script: 406 | 407 | ``` 408 | python main_semseg_scannet.py --exp_name=semseg_scannet 409 | ``` 410 | 411 | To train with both the training split and the validation split, use `--train_val=True`. 412 | 413 | You can use [TensorBoard](https://tensorflow.google.cn/tensorboard) to view the training log under `outputs/semseg_scannet/logs/`. 414 | 415 | ### Run the evaluation script after training finished: 416 | 417 | - Evaluate on the validation set 418 | 419 | ``` 420 | python main_semseg_scannet.py --eval=True --model_path=outputs/semseg_scannet/models/model_200.pth --exp_name=semseg_scannet_val --split=val 421 | ``` 422 | 423 | - Evaluate on the testing set 424 | 425 | ``` 426 | python main_semseg_scannet.py --eval=True --model_path=outputs/semseg_scannet/models/model_200.pth --exp_name=semseg_scannet_test --split=test 427 | ``` 428 | 429 | ### Run the evaluation script with pretrained models: 430 | 431 | - Evaluate on the validation set 432 | 433 | ``` 434 | python main_semseg_scannet.py --eval=True --model_path=pretrained/semseg_scannet/models/model_200.pth --exp_name=semseg_scannet_val --split=val 435 | ``` 436 | 437 | - Evaluate on the testing set 438 | 439 | ``` 440 | python main_semseg_scannet.py --eval=True --model_path=pretrained/semseg_scannet/models/model_200.pth --exp_name=semseg_scannet_test --split=test 441 | ``` 442 | 443 | Since there are no ground-truth labels on the testing set, this script will directly save prediction result. You need to upload your prediction results to the [website](https://kaldir.vc.in.tum.de/scannet_benchmark/semantic_label_3d) for evaluation. 444 | 445 | ### Performance: 446 | 447 | The validation set of the ScanNet Dataset 448 | 449 | | | Mean IoU | wall | floor | cabinet | bed | chair | sofa | table | door | window | bookshelf | picture | counter | desk | curtain | refrigerator | shower curtain | toilet | sink | bathtub | otherfurniture | 450 | | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | 451 | | This repo | 49.6 | 73.2 | 93.6 | 44.9 | 64.7 | 70.0 | 50.5 | 55.7 | 35.7 | 47.7 | 69.1 | 14.6 | 41.8 | 45.3 | 33.8 | 29.2 | 35.7 | 55.9 | 40.2 | 56.5 | 32.9 | 452 | 453 | The testing set of the ScanNet Dataset 454 | 455 | | | Mean IoU | wall | floor | cabinet | bed | chair | sofa | table | door | window | bookshelf | picture | counter | desk | curtain | refrigerator | shower curtain | toilet | sink | bathtub | other furniture | 456 | | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | 457 | | This repo | 44.6 | 72.3 | 93.7 | 36.6 | 62.3 | 65.1 | 57.7 | 44.5 | 33.0 | 39.4 | 46.3 | 12.6 | 31.0 | 34.9 | 38.9 | 28.5 | 22.4 | 62.5 | 35.0 | 47.4 | 27.1 | 458 | 459 | These is no official results of DGCNN on the ScanNet dataset. You can find our results on the [website](https://kaldir.vc.in.tum.de/scannet_benchmark/semantic_label_3d) as `DGCNN_reproduce`. 460 | 461 | ### Visualization: 462 | #### Usage: 463 | 464 | Use `--visu` to control visualization file. 465 | 466 | - To visualization a single scene, for example `scene0568_00`, use `--visu=scene0568_00`. 467 | - To visualization multiple scenes, please input these scenes separated by commas. e.g. `--visu=scene0568_00,scene0568_01`. 468 | - To visualization the scenes of a split, use `--visu=train`, `--visu=val` or `--visu=test`. 469 | - To visualization all the scenes in the dataset, use `--visu=all`. 470 | 471 | Use `--visu_format` to control visualization file format. 472 | 473 | - To output .txt file, use `--visu_format=txt`. 474 | - To output .ply format, use `--visu_format=ply`. 475 | 476 | Both .txt and .ply file can be loaded into [MeshLab](https://www.meshlab.net) for visualization. For the usage of MeshLab on .txt file, see issue [#8](https://github.com/AnTao97/dgcnn.pytorch/issues/8) for details. The .ply file can be directly loaded into MeshLab by dragging. 477 | 478 | For example, if you want to visualize `scene0568_00` on the pretrained model, run: 479 | 480 | ``` 481 | python main_semseg_scannet.py --eval=True --model_path=pretrained/semseg_scannet/models/model_200.pth --exp_name=semseg_scannet_val_visu --visu=scene0568_00 482 | ``` 483 | 484 | #### Results: 485 | The visualization result of `scene0568_00`: 486 | 487 |

488 | 489 |

490 | 491 | Color map: 492 |

493 | 494 |

-------------------------------------------------------------------------------- /README_zh.md: -------------------------------------------------------------------------------- 1 | # DGCNN.pytorch 2 | [[English]](README.md) 3 | 4 | 本仓库提供了一份PyTorch版本的 **Dynamic Graph CNN for Learning on Point Clouds (DGCNN)**（ https://arxiv.org/pdf/1801.07829 ）代码实现，代码框架来源于[WangYueFt/dgcnn](https://github.com/WangYueFt/dgcnn/tree/master/pytorch)。 5 | 6 | **更新：** 7 | 8 | - [2022/10/22] 增加ScanNet数据集上语义分割代码，作者：[吴紫屹](https://wuziyi616.github.io/) （dazitu616@gmail.com）、[王英琦](https://github.com/HilbertWang2002)（yingqi-w19@mails.tsinghua.edu.cn）。 9 | - [2021/07/20] 增加了可视化代码，作者：[纪鹏亮](https://github.com/Ji-Pengliang) (jpl1723@buaa.edu.cn). 10 | 11 | 在DGCNN文章中网络结构图（图3）中的分类网络和文章中对应的网络结构描述（第4.1节）并不吻合，原作者实际沿用了网络结构描述（第4.1节）中的网络结构，我们使用PS修复了网络结构图（图3）中不吻合的地方，修改后的图如下： 12 | 13 |   14 |

15 | 16 |

17 | 18 |   19 | 20 | **建议：** 3D点云实验结果往往比2D图像实验结果面临更大的随机性，因此我们建议您多跑几次实验，然后选择最佳的结果。 21 | 22 |   23 | ## 运行需求 24 | - Python 3.7 25 | - PyTorch 1.2 26 | - CUDA 10.0 27 | - Python包：glob, h5py, sklearn, plyfile 28 | 29 |   30 | ## 内容目录 31 | - [点云分类](#_3) 32 | - [点云局部分割](#_8) 33 | - [S3DIS数据集上点云场景语义分割](#s3dis) 34 | - [ScanNet数据集上点云场景语义分割](#scannet) 35 | 36 | **注意：** 以下所有命令默认使用所有显卡，如果需要明确使用几张显卡，比如4张显卡索引为`0,1,2,3`，那么在每个命令前需要添加`CUDA_VISIBLE_DEVICES=0,1,2,3`。你可以根据你的需要调整使用显卡的数量和索引。 37 | 38 |   39 | ## 点云分类 40 | ### 运行训练脚本： 41 | 42 | - 1024点 43 | 44 | ``` 45 | python main_cls.py --exp_name=cls_1024 --num_points=1024 --k=20 46 | ``` 47 | 48 | - 2048点 49 | 50 | ``` 51 | python main_cls.py --exp_name=cls_2048 --num_points=2048 --k=40 52 | ``` 53 | 54 | ### 训练结束后运行评估脚本： 55 | 56 | - 1024点 57 | 58 | ``` 59 | python main_cls.py --exp_name=cls_1024_eval --num_points=1024 --k=20 --eval=True --model_path=outputs/cls_1024/models/model.t7 60 | ``` 61 | 62 | - 2048点 63 | 64 | ``` 65 | python main_cls.py --exp_name=cls_2048_eval --num_points=2048 --k=40 --eval=True --model_path=outputs/cls_2048/models/model.t7 66 | ``` 67 | 68 | ### 使用提供的已训练模型运行评估脚本： 69 | 70 | - 1024点 71 | 72 | ``` 73 | python main_cls.py --exp_name=cls_1024_eval --num_points=1024 --k=20 --eval=True --model_path=pretrained/model.cls.1024.t7 74 | ``` 75 | 76 | - 2048点 77 | 78 | ``` 79 | python main_cls.py --exp_name=cls_2048_eval --num_points=2048 --k=40 --eval=True --model_path=pretrained/model.cls.2048.t7 80 | ``` 81 | 82 | ### 模型性能： 83 | ModelNet40数据集 84 | 85 | | | 平均类别Acc | 整体Acc | 86 | | :---: | :---: | :---: | 87 | | 原文章（1024点） | 90.2 | 92.9 | 88 | | 本仓库（1024点） | **90.9** | **93.3** | 89 | | 原文章（2048点） | 90.7 | 93.5 | 90 | | 本仓库（2048点） | **91.2** | **93.6** | 91 | 92 |   93 | ## 点云局部分割 94 | **注意：** 训练模式中 **“全部类别”** 和 **“特定类别”** 有区别。 95 | 96 | - 在 **“全部类别”** 中，模型使用所有16个类别的数据进行训练和评估，并在这份代码中得出平均IoU为85.2%。每个点云形状中点的预测标签可以是16个类别中的任意一个局部类别。 97 | - 在 **“特定类别”** 中，模型仅使用1个类别的数据进行训练和评估，比如飞机类别，并在这份代码中得出飞机类别的IoU为84.5%。每个点云形状中点的预测标签只能是这1个类别中的任意一个局部类别。 98 | 99 | ### 运行训练脚本： 100 | 101 | - 使用数据集内全部类别 102 | 103 | ``` 104 | python main_partseg.py --exp_name=partseg 105 | ``` 106 | 107 | - 选择数据集内特定类别，例如飞机 108 | 109 | ``` 110 | python main_partseg.py --exp_name=partseg_airplane --class_choice=airplane 111 | ``` 112 | 113 | ### 训练结束后运行评估脚本： 114 | 115 | - 使用数据集内全部类别 116 | 117 | ``` 118 | python main_partseg.py --exp_name=partseg_eval --eval=True --model_path=outputs/partseg/models/model.t7 119 | ``` 120 | 121 | - 选择数据集内特定类别，例如飞机 122 | 123 | ``` 124 | python main_partseg.py --exp_name=partseg_airplane_eval --class_choice=airplane --eval=True --model_path=outputs/partseg_airplane/models/model.t7 125 | ``` 126 | 127 | ### 使用提供的已训练模型运行评估脚本： 128 | 129 | - 使用数据集内全部类别 130 | 131 | ``` 132 | python main_partseg.py --exp_name=partseg_eval --eval=True --model_path=pretrained/model.partseg.t7 133 | ``` 134 | 135 | - 选择数据集内特定类别，例如飞机 136 | 137 | ``` 138 | python main_partseg.py --exp_name=partseg_airplane_eval --class_choice=airplane --eval=True --model_path=pretrained/model.partseg.airplane.t7 139 | ``` 140 | 141 | ### 模型性能 142 | ShapeNet part 数据集 143 | 144 | |      | 平均IoU | 飞机 | 包 | 帽子 | 车 | 椅子 | 耳机 | 吉他 | 刀 | 灯 | 笔记本电脑 | 摩托车 | 杯子 | 手枪 | 火箭 | 滑板 | 桌子 | 145 | | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | 146 | |      | mIoU | airplane | bag | cap | car | chair | earphone | guitar | knife | lamp | laptop | motor | mug | pistol | rocket | skateboard | table | 147 | | 形状数量 |      | 2690 | 76 | 55 | 898 | 3758 | 69 | 787 | 392 | 1547 | 451 | 202 | 184 | 283 | 66 | 152 | 5271 | 148 | | 原文章 | **85.2** | 84.0 | **83.4** | **86.7** | 77.8 | 90.6 | 74.7 | 91.2 | **87.5** | 82.8 | **95.7** | 66.3 | **94.9** | 81.1 | **63.5** | 74.5 | 82.6 | 149 | | 本仓库 | **85.2** | **84.5** | 80.3 | 84.7 | **79.8** | **91.1** | **76.8** | **92.0** | 87.3 | **83.8** | **95.7** | **69.6** | 94.3 | **83.7** | 51.5 | **76.1** | **82.8** | 150 | 151 | ### 可视化： 152 | #### 使用说明： 153 | 154 | 使用`--visu`控制需要可视化的文件。 155 | 156 | - 对于可视化一个点云形状，比如类别飞机索引为0的点云形状（点云形状的索引为从0开始），使用`--visu=airplane_0`。 157 | - 对于可视化一个类别中的所有点云形状，比如飞机类别，使用`--visu=airplane`。 158 | - 对于可视化所有类别的所有点云类别，使用`--visu=all`。 159 | 160 | 使用`--visu_format`控制可视化文件的格式。 161 | 162 | - 输出.txt格式，使用`--visu_format=txt`。 163 | - 输出.ply格式，使用`--visu_format=ply`。 164 | 165 | 所有格式均可通过[MeshLab](https://www.meshlab.net)来加载进行可视化。对于使用MeshLab可视化.txt格式，参考问题[#8](https://github.com/AnTao97/dgcnn.pytorch/issues/8)中的介绍，而.ply格式可以直接拖入MeshLab进行可视化。 166 | 167 | 可视化文件名遵循统一的命名格式。对于预测结果，文件名格式为`点云形状名称_pred_miou.格式后缀`；对于真实标签，文件名格式为`点云形状名称_gt.格式后缀`。文件名中`miou`指该点云形状的平均IoU。 168 | 169 | #### 全部类别： 170 | 171 | - 输出飞机类别索引为0的点云形状的.ply格式可视化结果 172 | 173 | ``` 174 | # 使用训练后的模型 175 | python main_partseg.py --exp_name=partseg_eval --eval=True --model_path=outputs/partseg/models/model.t7 --visu=airplane_0 --visu_format=ply 176 | 177 | # 使用提供的已训练模型 178 | python main_partseg.py --exp_name=partseg_eval --eval=True --model_path=pretrained/model.partseg.t7 --visu=airplane_0 --visu_format=ply 179 | ``` 180 | 181 | - 输出飞机类别的所有点云形状的.ply格式可视化结果 182 | 183 | ``` 184 | # 使用训练后的模型 185 | python main_partseg.py --exp_name=partseg_eval --eval=True --model_path=outputs/partseg/models/model.t7 --visu=airplane --visu_format=ply 186 | 187 | # 使用提供的已训练模型 188 | python main_partseg.py --exp_name=partseg_eval --eval=True --model_path=pretrained/model.partseg.t7 --visu=airplane --visu_format=ply 189 | ``` 190 | 191 | - 输出所有类别的所有点云形状的.ply格式可视化结果 192 | 193 | ``` 194 | # 使用训练后的模型 195 | python main_partseg.py --exp_name=partseg_eval --eval=True --model_path=outputs/partseg/models/model.t7 --visu=all --visu_format=ply 196 | 197 | # 使用提供的已训练模型 198 | python main_partseg.py --exp_name=partseg_eval --eval=True --model_path=pretrained/model.partseg.t7 --visu=all --visu_format=ply 199 | ``` 200 | 201 | #### 选择数据集内特定类别，例如飞机： 202 | 203 | - 输出飞机类别索引为0的点云形状的.ply格式可视化结果 204 | 205 | ``` 206 | # 使用训练后的模型 207 | python main_partseg.py --exp_name=partseg_airplane_eval --class_choice=airplane --eval=True --model_path=outputs/partseg_airplane/models/model.t7 --visu=airplane_0 --visu_format=ply 208 | 209 | # 使用提供的已训练模型 210 | python main_partseg.py --exp_name=partseg_airplane_eval --class_choice=airplane --eval=True --model_path=pretrained/model.partseg.airplane.t7 --visu=airplane_0 --visu_format=ply 211 | ``` 212 | 213 | - 输出飞机类别的所有点云形状的.ply格式可视化结果 214 | 215 | ``` 216 | # 使用训练后的模型 217 | python main_partseg.py --exp_name=partseg_airplane_eval --class_choice=airplane --eval=True --model_path=outputs/partseg_airplane/models/model.t7 --visu=airplane --visu_format=ply 218 | 219 | # 使用提供的已训练模型 220 | python main_partseg.py --exp_name=partseg_airplane_eval --class_choice=airplane --eval=True --model_path=pretrained/model.partseg.airplane.t7 --visu=airplane --visu_format=ply 221 | ``` 222 | 223 | #### 可视化结果： 224 | 属于飞机类别索引为0的点云形状的可视化结果： 225 | 226 |

227 | 228 |

229 | 230 | 颜色对应图： 231 |

232 | 233 |

234 | 235 |   236 | ## 在S3DIS数据集上点云场景语义分割 237 | 238 | 在此任务中网络结构和点云局部分割有细微不同，最后的一个MLP尺寸改为（512, 256, 13），而且在256后只使用一个dropout。 239 | 240 | 您必须从 https://goo.gl/forms/4SoGp4KtH1jfRqEj2 手动下载数据集`Stanford3dDataset_v1.2_Aligned_Version.zip`，然后将其放在`data/`目录下。 241 | 242 | ### 运行训练脚本： 243 | 244 | 此任务使用6折训练，因此需要训练6个模型，轮流选择数据集中6个区域中的1个作为这个模型的测试区域。 245 | 246 | - 在区域1-5上训练 247 | 248 | ``` 249 | python main_semseg_s3dis.py --exp_name=semseg_s3dis_6 --test_area=6 250 | ``` 251 | 252 | ### 训练结束后运行评估脚本： 253 | 254 | - 当模型在区域1-5训练完成后，在区域6中评估 255 | 256 | ``` 257 | python main_semseg_s3dis.py --exp_name=semseg_s3dis_eval_6 --test_area=6 --eval=True --model_root=outputs/semseg_s3dis_s3dis/models/ 258 | ``` 259 | 260 | - 当6个模型训练完成后，在所有区域上评估 261 | 262 | ``` 263 | python main_semseg_s3dis.py --exp_name=semseg_s3dis_eval --test_area=all --eval=True --model_root=outputs/semseg_s3dis/models/ 264 | ``` 265 | 266 | ### 使用提供的已训练模型运行评估脚本： 267 | 268 | - 使用提供的在区域1-5上已训练模型，在区域6中评估 269 | 270 | ``` 271 | python main_semseg_s3dis.py --exp_name=semseg_s3dis_eval_6 --test_area=6 --eval=True --model_root=pretrained/semseg_s3dis/ 272 | ``` 273 | 274 | - 使用提供的6个已训练模型，在所有区域上评估 275 | 276 | ``` 277 | python main_semseg_s3dis.py --exp_name=semseg_s3dis_eval --test_area=all --eval=True --model_root=pretrained/semseg_s3dis/ 278 | ``` 279 | 280 | ### 模型性能： 281 | 斯坦福大学大型3D室内空间数据集（S3DIS） 282 | 283 | | | 平均IoU | 整体Acc | 284 | | :---: | :---: | :---: | 285 | | 原文章 | 56.1 | 84.1 | 286 | | 本仓库 | **59.2** | **85.0** | 287 | 288 | ### 可视化： 289 | #### 使用说明： 290 | 291 | 使用`--visu`控制需要可视化的文件。 292 | 293 | - 对于可视化一个房间，比如第区域6的办公室1（房间索引为从1开始），使用`--visu=area_6_office_1`。 294 | - 对于可视化一个区域中所有的房间，比如区域6，使用`--visu=area_6`。 295 | - 对于可视化所有区域的所有房间，使用`--visu=all`。 296 | 297 | 使用`--visu_format`控制可视化文件的格式。 298 | 299 | - 输出.txt格式，使用`--visu_format=txt`。 300 | - 输出.ply格式，使用`--visu_format=ply`。 301 | 302 | 所有格式均可通过[MeshLab](https://www.meshlab.net)来加载进行可视化。对于使用MeshLab可视化.txt格式，参考问题[#8](https://github.com/AnTao97/dgcnn.pytorch/issues/8)中的介绍，而.ply格式可以直接拖入MeshLab进行可视化。 303 | 304 | 可视化文件名遵循统一的命名格式。对于预测结果，文件名格式为`房间名称_pred_.格式后缀`；对于真实标签，文件名格式为`房间名称_gt.格式后缀`。文件名中``指该房间的平均IoU。 305 | 306 | **注意：**对于语义分割，需要首先运行一个不带可视化的命令来预处理数据集，比如可视化部分之前的训练命令和评估命令。在数据集处理完成后，可以运行下面带有可视化的命令。 307 | 308 | #### 当模型在区域1-5训练完成后，在区域6中评估： 309 | 310 | - 输出区域6的办公室1的.ply格式可视化结果 311 | 312 | ``` 313 | # 使用训练后的模型 314 | python main_semseg_s3dis.py --exp_name=semseg_s3dis_eval_6 --test_area=6 --eval=True --model_root=outputs/semseg_s3dis/models/ --visu=area_6_office_1 --visu_format=ply 315 | 316 | # 使用提供的已训练模型 317 | python main_semseg_s3dis.py --exp_name=semseg_s3dis_eval_6 --test_area=6 --eval=True --model_root=pretrained/semseg_s3dis/ --visu=area_6_office_1 --visu_format=ply 318 | ``` 319 | 320 | - 输出区域6的所有房间的.ply格式可视化结果 321 | 322 | ``` 323 | # 使用训练后的模型 324 | python main_semseg_s3dis.py --exp_name=semseg_s3dis_eval_6 --test_area=6 --eval=True --model_root=outputs/semseg_s3dis/models/ --visu=area_6 --visu_format=ply 325 | 326 | # 使用提供的已训练模型 327 | python main_semseg_s3dis.py --exp_name=semseg_s3dis_eval_6 --test_area=6 --eval=True --model_root=pretrained/semseg_s3dis/ --visu=area_6 --visu_format=ply 328 | ``` 329 | 330 | #### 当6个模型训练完成后，在所有区域上评估： 331 | 332 | - 输出区域6的办公室1的.ply格式可视化结果 333 | 334 | 335 | ``` 336 | # 使用训练后的模型 337 | python main_semseg_s3dis.py --exp_name=semseg_s3dis_eval --test_area=all --eval=True --model_root=outputs/semseg/models_s3dis/ --visu=area_6_office_1 --visu_format=ply 338 | 339 | # 使用提供的已训练模型 340 | python main_semseg_s3dis.py --exp_name=semseg_s3dis_eval --test_area=all --eval=True --model_root=pretrained/semseg_s3dis/ --visu=area_6_office_1 --visu_format=ply 341 | ``` 342 | 343 | - 输出区域6的所有房间的.ply格式可视化结果 344 | 345 | ``` 346 | # 使用训练后的模型 347 | python main_semseg_s3dis.py --exp_name=semseg_s3dis_eval --test_area=all --eval=True --model_root=outputs/semseg_s3dis/models/ --visu=area_6 --visu_format=ply 348 | 349 | # 使用提供的已训练模型 350 | python main_semseg_s3dis.py --exp_name=semseg_s3dis_eval --test_area=all --eval=True --model_root=pretrained/semseg_s3dis/ --visu=area_6 --visu_format=ply 351 | ``` 352 | 353 | - 输出所有区域的所有房间的.ply格式可视化结果 354 | 355 | ``` 356 | # 使用训练后的模型 357 | python main_semseg_s3dis.py --exp_name=semseg_s3dis_eval --test_area=all --eval=True --model_root=outputs/semseg_s3dis/models/ --visu=all --visu_format=ply 358 | 359 | # 使用提供的已训练模型 360 | python main_semseg_s3dis.py --exp_name=semseg_s3dis_eval --test_area=all --eval=True --model_root=pretrained/semseg_s3dis/ --visu=all --visu_format=ply 361 | ``` 362 | 363 | #### 可视化结果： 364 | 区域6的办公室1的可视化结果： 365 | 366 |

367 | 368 |

369 | 370 | 颜色对应图： 371 |

372 | 373 |

374 | 375 |   376 | ## 在ScanNet数据集上点云场景语义分割 377 | 378 | DGCNN的作者并没有在ScanNet数据集上测试， 我们尽可能实现了一个DGCNN模型在ScanNet数据集上的实现版本。 379 | 380 | ### 准备数据集： 381 | 382 | 你需要把程序运行目录改变至`prepare_data/`文件夹下。 383 | 384 | ``` 385 | cd prepare_data/ 386 | ``` 387 | 388 | 请从[ScanNet网站](http://www.scan-net.org/)下载原始数据集，并放置其于`data/ScanNet/`，注意路径`data/ScanNet`内需要包括`data/ScanNet/scans/`文件夹和`data/ScanNet/scans_test/`文件夹。 389 | 390 | 使用下面命令准备ScanNet数据集 391 | 392 | ``` 393 | python scannetv2_seg_dataset_rgb21c_pointid.py 394 | ``` 395 | 396 | 本命令会产生四个pickle文件：`scannet_train_rgb21c_pointid.pickle`、`scannet_val_rgb21c_pointid.pickle`、`scannet_val_rgb21c_pointid_keep_unanno.pickle`和`scannet_test_rgb21c_pointid_keep_unanno.pickle`。 397 | 398 | 返回主目录： 399 | 400 | ``` 401 | cd .. 402 | ``` 403 | 404 | ### 运行训练脚本： 405 | 406 | ``` 407 | python main_semseg_scannet.py --exp_name=semseg_scannet 408 | ``` 409 | 410 | 如果训练时同时使用训练集和验证集，使用`--train_val=True`。 411 | 412 | 你可以使用[TensorBoard](https://tensorflow.google.cn/tensorboard)查看路径`outputs/semseg_scannet/logs/`保存的训练记录。 413 | 414 | ### 训练结束后运行评估脚本： 415 | 416 | - 在验证集上评估 417 | 418 | ``` 419 | python main_semseg_scannet.py --eval=True --model_path=outputs/semseg_scannet/models/model_200.pth --exp_name=semseg_scannet_val --split=val 420 | ``` 421 | 422 | - 在测试集上评估 423 | 424 | ``` 425 | python main_semseg_scannet.py --eval=True --model_path=outputs/semseg_scannet/models/model_200.pth --exp_name=semseg_scannet_test --split=test 426 | ``` 427 | 428 | ### 使用提供的已训练模型运行评估脚本： 429 | 430 | - 在验证集上评估 431 | 432 | ``` 433 | python main_semseg_scannet.py --eval=True --model_path=pretrained/semseg_scannet/models/model_200.pth --exp_name=semseg_scannet_val --split=val 434 | ``` 435 | 436 | - 在测试集上评估 437 | 438 | ``` 439 | python main_semseg_scannet.py --eval=True --model_path=pretrained/semseg_scannet/models/model_200.pth --exp_name=semseg_scannet_test --split=test 440 | ``` 441 | 442 | 由于测试集内没有真实标签，所以测试脚本会直接保存预测结果，你需要将预测结果上传至[ScanNet评估网站](https://kaldir.vc.in.tum.de/scannet_benchmark/semantic_label_3d) 获取评估结果。 443 | 444 | ### 模型性能： 445 | 446 | ScanNet数据集验证集 447 | 448 | | | Mean IoU | wall | floor | cabinet | bed | chair | sofa | table | door | window | bookshelf | picture | counter | desk | curtain | refrigerator | shower curtain | toilet | sink | bathtub | otherfurniture | 449 | | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | 450 | | This repo | 49.7 | 72.6 | 93.8 | 42.1 | 61.4 | 70.0 | 48.5 | 54.9 | 36.9 | 47.6 | 67.4 | 13.3 | 42.8 | 45.1 | 39.5 | 27.5 | 38.1 | 60.0 | 41.5 | 61.5 | 28.9 | 451 | 452 | ScanNet数据集测试集 453 | 454 | | | Mean IoU | wall | floor | cabinet | bed | chair | sofa | table | door | window | bookshelf | picture | counter | desk | curtain | refrigerator | shower curtain | toilet | sink | bathtub | other furniture | 455 | | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | 456 | | This repo | 44.6 | 72.3 | 93.7 | 36.6 | 62.3 | 65.1 | 57.7 | 44.5 | 33.0 | 39.4 | 46.3 | 12.6 | 31.0 | 34.9 | 38.9 | 28.5 | 22.4 | 62.5 | 35.0 | 47.4 | 27.1 | 457 | 458 | ScanNet数据集上没有DGCNN的官方结果，你可以在[ScanNet评估网站](https://kaldir.vc.in.tum.de/scannet_benchmark/semantic_label_3d) 上找到我们的上传记录`DGCNN_reproduce`。 459 | 460 | ### 可视化： 461 | #### 使用说明： 462 | 463 | 使用`--visu`控制需要可视化的文件。 464 | 465 | - 对于可视化一个场景，比如`scene0568_00`，使用`--visu=scene0568_00`。 466 | - 对于可视化多个场景，请将场景之间用英文逗号区分，使用`--visu=scene0568_00,scene0568_01`。 467 | - 对于可视化一个划分集合内的场景，使用`--visu=train`或者`--visu=val`或者`--visu=test`。 468 | - 对于可视化数据集内所有场景，使用`--visu=all`。 469 | 470 | 使用`--visu_format`控制可视化文件的格式。 471 | 472 | - 输出.txt格式，使用`--visu_format=txt`。 473 | - 输出.ply格式，使用`--visu_format=ply`。 474 | 475 | 所有格式均可通过[MeshLab](https://www.meshlab.net)来加载进行可视化。对于使用MeshLab可视化.txt格式，参考问题[#8](https://github.com/AnTao97/dgcnn.pytorch/issues/8)中的介绍，而.ply格式可以直接拖入MeshLab进行可视化。 476 | 477 | 例如，如果你想可视化预训练模型在`scene0568_00`的结果，运行： 478 | 479 | ``` 480 | python main_semseg_scannet.py --eval=True --model_path=pretrained/semseg_scannet/models/model_200.pth --exp_name=semseg_scannet_val_visu --visu=scene0568_00 481 | ``` 482 | 483 | #### 可视化结果： 484 | 场景`scene0568_00`的可视化结果： 485 | 486 |

487 | 488 |

489 | 490 | 颜色对应图： 491 |

492 | 493 |

-------------------------------------------------------------------------------- /data.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | """ 4 | @Author: Yue Wang 5 | @Contact: yuewangx@mit.edu 6 | @File: data.py 7 | @Time: 2018/10/13 6:21 PM 8 | 9 | Modified by 10 | @Author: An Tao, Pengliang Ji, Ziyi Wu 11 | @Contact: ta19@mails.tsinghua.edu.cn, jpl1723@buaa.edu.cn, dazitu616@gmail.com 12 | @Time: 2022/7/30 7:49 PM 13 | """ 14 | 15 | 16 | import os 17 | import sys 18 | import glob 19 | import h5py 20 | import numpy as np 21 | import torch 22 | import json 23 | import cv2 24 | import pickle 25 | from torch.utils.data import Dataset 26 | 27 | 28 | def download_modelnet40(): 29 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 30 | DATA_DIR = os.path.join(BASE_DIR, 'data') 31 | if not os.path.exists(DATA_DIR): 32 | os.mkdir(DATA_DIR) 33 | if not os.path.exists(os.path.join(DATA_DIR, 'modelnet40_ply_hdf5_2048')): 34 | www = 'https://shapenet.cs.stanford.edu/media/modelnet40_ply_hdf5_2048.zip' 35 | zipfile = os.path.basename(www) 36 | os.system('wget --no-check-certificate %s; unzip %s' % (www, zipfile)) 37 | os.system('mv %s %s' % ('modelnet40_ply_hdf5_2048', DATA_DIR)) 38 | os.system('rm %s' % (zipfile)) 39 | 40 | 41 | def download_shapenetpart(): 42 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 43 | DATA_DIR = os.path.join(BASE_DIR, 'data') 44 | if not os.path.exists(DATA_DIR): 45 | os.mkdir(DATA_DIR) 46 | if not os.path.exists(os.path.join(DATA_DIR, 'shapenet_part_seg_hdf5_data')): 47 | www = 'https://shapenet.cs.stanford.edu/media/shapenet_part_seg_hdf5_data.zip' 48 | zipfile = os.path.basename(www) 49 | os.system('wget --no-check-certificate %s; unzip %s' % (www, zipfile)) 50 | os.system('mv %s %s' % ('hdf5_data', os.path.join(DATA_DIR, 'shapenet_part_seg_hdf5_data'))) 51 | os.system('rm %s' % (zipfile)) 52 | 53 | 54 | def download_S3DIS(): 55 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 56 | DATA_DIR = os.path.join(BASE_DIR, 'data') 57 | if not os.path.exists(DATA_DIR): 58 | os.mkdir(DATA_DIR) 59 | if not os.path.exists(os.path.join(DATA_DIR, 'indoor3d_sem_seg_hdf5_data')): 60 | www = 'https://shapenet.cs.stanford.edu/media/indoor3d_sem_seg_hdf5_data.zip' 61 | zipfile = os.path.basename(www) 62 | os.system('wget %s --no-check-certificate; unzip %s' % (www, zipfile)) 63 | os.system('mv %s %s' % ('indoor3d_sem_seg_hdf5_data', DATA_DIR)) 64 | os.system('rm %s' % (zipfile)) 65 | if not os.path.exists(os.path.join(DATA_DIR, 'Stanford3dDataset_v1.2_Aligned_Version')): 66 | if not os.path.exists(os.path.join(DATA_DIR, 'Stanford3dDataset_v1.2_Aligned_Version.zip')): 67 | print('Please download Stanford3dDataset_v1.2_Aligned_Version.zip \ 68 | from https://goo.gl/forms/4SoGp4KtH1jfRqEj2 and place it under data/') 69 | sys.exit(0) 70 | else: 71 | zippath = os.path.join(DATA_DIR, 'Stanford3dDataset_v1.2_Aligned_Version.zip') 72 | os.system('unzip %s' % (zippath)) 73 | os.system('mv %s %s' % ('Stanford3dDataset_v1.2_Aligned_Version', DATA_DIR)) 74 | os.system('rm %s' % (zippath)) 75 | 76 | 77 | def load_data_cls(partition): 78 | download_modelnet40() 79 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 80 | DATA_DIR = os.path.join(BASE_DIR, 'data') 81 | all_data = [] 82 | all_label = [] 83 | for h5_name in glob.glob(os.path.join(DATA_DIR, 'modelnet40_ply_hdf5_2048', '*%s*.h5'%partition)): 84 | f = h5py.File(h5_name, 'r+') 85 | data = f['data'][:].astype('float32') 86 | label = f['label'][:].astype('int64') 87 | f.close() 88 | all_data.append(data) 89 | all_label.append(label) 90 | all_data = np.concatenate(all_data, axis=0) 91 | all_label = np.concatenate(all_label, axis=0) 92 | return all_data, all_label 93 | 94 | 95 | def load_data_partseg(partition): 96 | download_shapenetpart() 97 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 98 | DATA_DIR = os.path.join(BASE_DIR, 'data') 99 | all_data = [] 100 | all_label = [] 101 | all_seg = [] 102 | if partition == 'trainval': 103 | file = glob.glob(os.path.join(DATA_DIR, 'shapenet_part_seg_hdf5_data', '*train*.h5')) \ 104 | + glob.glob(os.path.join(DATA_DIR, 'shapenet_part_seg_hdf5_data', '*val*.h5')) 105 | else: 106 | file = glob.glob(os.path.join(DATA_DIR, 'shapenet_part_seg_hdf5_data', '*%s*.h5'%partition)) 107 | for h5_name in file: 108 | f = h5py.File(h5_name, 'r+') 109 | data = f['data'][:].astype('float32') 110 | label = f['label'][:].astype('int64') 111 | seg = f['pid'][:].astype('int64') 112 | f.close() 113 | all_data.append(data) 114 | all_label.append(label) 115 | all_seg.append(seg) 116 | all_data = np.concatenate(all_data, axis=0) 117 | all_label = np.concatenate(all_label, axis=0) 118 | all_seg = np.concatenate(all_seg, axis=0) 119 | return all_data, all_label, all_seg 120 | 121 | 122 | def prepare_test_data_semseg(): 123 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 124 | DATA_DIR = os.path.join(BASE_DIR, 'data') 125 | if not os.path.exists(os.path.join(DATA_DIR, 'stanford_indoor3d')): 126 | os.system('python prepare_data/collect_indoor3d_data.py') 127 | if not os.path.exists(os.path.join(DATA_DIR, 'indoor3d_sem_seg_hdf5_data_test')): 128 | os.system('python prepare_data/gen_indoor3d_h5.py') 129 | 130 | 131 | def load_data_semseg(partition, test_area): 132 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 133 | DATA_DIR = os.path.join(BASE_DIR, 'data') 134 | download_S3DIS() 135 | prepare_test_data_semseg() 136 | if partition == 'train': 137 | data_dir = os.path.join(DATA_DIR, 'indoor3d_sem_seg_hdf5_data') 138 | else: 139 | data_dir = os.path.join(DATA_DIR, 'indoor3d_sem_seg_hdf5_data_test') 140 | with open(os.path.join(data_dir, "all_files.txt")) as f: 141 | all_files = [line.rstrip() for line in f] 142 | with open(os.path.join(data_dir, "room_filelist.txt")) as f: 143 | room_filelist = [line.rstrip() for line in f] 144 | data_batchlist, label_batchlist = [], [] 145 | for f in all_files: 146 | file = h5py.File(os.path.join(DATA_DIR, f), 'r+') 147 | data = file["data"][:] 148 | label = file["label"][:] 149 | data_batchlist.append(data) 150 | label_batchlist.append(label) 151 | data_batches = np.concatenate(data_batchlist, 0) 152 | seg_batches = np.concatenate(label_batchlist, 0) 153 | test_area_name = "Area_" + test_area 154 | train_idxs, test_idxs = [], [] 155 | for i, room_name in enumerate(room_filelist): 156 | if test_area_name in room_name: 157 | test_idxs.append(i) 158 | else: 159 | train_idxs.append(i) 160 | if partition == 'train': 161 | all_data = data_batches[train_idxs, ...] 162 | all_seg = seg_batches[train_idxs, ...] 163 | else: 164 | all_data = data_batches[test_idxs, ...] 165 | all_seg = seg_batches[test_idxs, ...] 166 | return all_data, all_seg 167 | 168 | 169 | def load_color_partseg(): 170 | colors = [] 171 | labels = [] 172 | f = open("prepare_data/meta/partseg_colors.txt") 173 | for line in json.load(f): 174 | colors.append(line['color']) 175 | labels.append(line['label']) 176 | partseg_colors = np.array(colors) 177 | partseg_colors = partseg_colors[:, [2, 1, 0]] 178 | partseg_labels = np.array(labels) 179 | font = cv2.FONT_HERSHEY_SIMPLEX 180 | img_size = 1350 181 | img = np.zeros((1350, 1890, 3), dtype="uint8") 182 | cv2.rectangle(img, (0, 0), (1900, 1900), [255, 255, 255], thickness=-1) 183 | column_numbers = [4, 2, 2, 4, 4, 3, 3, 2, 4, 2, 6, 2, 3, 3, 3, 3] 184 | column_gaps = [320, 320, 300, 300, 285, 285] 185 | color_size = 64 186 | color_index = 0 187 | label_index = 0 188 | row_index = 16 189 | for row in range(0, img_size): 190 | column_index = 32 191 | for column in range(0, img_size): 192 | color = partseg_colors[color_index] 193 | label = partseg_labels[label_index] 194 | length = len(str(label)) 195 | cv2.rectangle(img, (column_index, row_index), (column_index + color_size, row_index + color_size), 196 | color=(int(color[0]), int(color[1]), int(color[2])), thickness=-1) 197 | img = cv2.putText(img, label, (column_index + int(color_size * 1.15), row_index + int(color_size / 2)), 198 | font, 199 | 0.76, (0, 0, 0), 2) 200 | column_index = column_index + column_gaps[column] 201 | color_index = color_index + 1 202 | label_index = label_index + 1 203 | if color_index >= 50: 204 | cv2.imwrite("prepare_data/meta/partseg_colors.png", img, [cv2.IMWRITE_PNG_COMPRESSION, 0]) 205 | return np.array(colors) 206 | elif (column + 1 >= column_numbers[row]): 207 | break 208 | row_index = row_index + int(color_size * 1.3) 209 | if (row_index >= img_size): 210 | break 211 | 212 | 213 | def load_color_semseg(): 214 | colors = [] 215 | labels = [] 216 | f = open("prepare_data/meta/semseg_colors.txt") 217 | for line in json.load(f): 218 | colors.append(line['color']) 219 | labels.append(line['label']) 220 | semseg_colors = np.array(colors) 221 | semseg_colors = semseg_colors[:, [2, 1, 0]] 222 | partseg_labels = np.array(labels) 223 | font = cv2.FONT_HERSHEY_SIMPLEX 224 | img_size = 1500 225 | img = np.zeros((500, img_size, 3), dtype="uint8") 226 | cv2.rectangle(img, (0, 0), (img_size, 750), [255, 255, 255], thickness=-1) 227 | color_size = 64 228 | color_index = 0 229 | label_index = 0 230 | row_index = 16 231 | for _ in range(0, img_size): 232 | column_index = 32 233 | for _ in range(0, img_size): 234 | color = semseg_colors[color_index] 235 | label = partseg_labels[label_index] 236 | length = len(str(label)) 237 | cv2.rectangle(img, (column_index, row_index), (column_index + color_size, row_index + color_size), 238 | color=(int(color[0]), int(color[1]), int(color[2])), thickness=-1) 239 | img = cv2.putText(img, label, (column_index + int(color_size * 1.15), row_index + int(color_size / 2)), 240 | font, 241 | 0.7, (0, 0, 0), 2) 242 | column_index = column_index + 200 243 | color_index = color_index + 1 244 | label_index = label_index + 1 245 | if color_index >= 13: 246 | cv2.imwrite("prepare_data/meta/semseg_colors.png", img, [cv2.IMWRITE_PNG_COMPRESSION, 0]) 247 | return np.array(colors) 248 | elif (column_index >= 1280): 249 | break 250 | row_index = row_index + int(color_size * 1.3) 251 | if (row_index >= img_size): 252 | break 253 | 254 | 255 | def translate_pointcloud(pointcloud): 256 | xyz1 = np.random.uniform(low=2./3., high=3./2., size=[3]) 257 | xyz2 = np.random.uniform(low=-0.2, high=0.2, size=[3]) 258 | 259 | translated_pointcloud = np.add(np.multiply(pointcloud, xyz1), xyz2).astype('float32') 260 | return translated_pointcloud 261 | 262 | 263 | def jitter_pointcloud(pointcloud, sigma=0.01, clip=0.02): 264 | N, C = pointcloud.shape 265 | pointcloud += np.clip(sigma * np.random.randn(N, C), -1*clip, clip) 266 | return pointcloud 267 | 268 | 269 | def rotate_pointcloud(pointcloud): 270 | theta = np.pi*2 * np.random.uniform() 271 | rotation_matrix = np.array([[np.cos(theta), -np.sin(theta)],[np.sin(theta), np.cos(theta)]]) 272 | pointcloud[:,[0,2]] = pointcloud[:,[0,2]].dot(rotation_matrix) # random rotation (x,z) 273 | return pointcloud 274 | 275 | 276 | class ModelNet40(Dataset): 277 | def __init__(self, num_points, partition='train'): 278 | self.data, self.label = load_data_cls(partition) 279 | self.num_points = num_points 280 | self.partition = partition 281 | 282 | def __getitem__(self, item): 283 | pointcloud = self.data[item][:self.num_points] 284 | label = self.label[item] 285 | if self.partition == 'train': 286 | pointcloud = translate_pointcloud(pointcloud) 287 | np.random.shuffle(pointcloud) 288 | return pointcloud, label 289 | 290 | def __len__(self): 291 | return self.data.shape[0] 292 | 293 | 294 | class ShapeNetPart(Dataset): 295 | def __init__(self, num_points, partition='train', class_choice=None): 296 | self.data, self.label, self.seg = load_data_partseg(partition) 297 | self.cat2id = {'airplane': 0, 'bag': 1, 'cap': 2, 'car': 3, 'chair': 4, 298 | 'earphone': 5, 'guitar': 6, 'knife': 7, 'lamp': 8, 'laptop': 9, 299 | 'motor': 10, 'mug': 11, 'pistol': 12, 'rocket': 13, 'skateboard': 14, 'table': 15} 300 | self.seg_num = [4, 2, 2, 4, 4, 3, 3, 2, 4, 2, 6, 2, 3, 3, 3, 3] 301 | self.index_start = [0, 4, 6, 8, 12, 16, 19, 22, 24, 28, 30, 36, 38, 41, 44, 47] 302 | self.num_points = num_points 303 | self.partition = partition 304 | self.class_choice = class_choice 305 | self.partseg_colors = load_color_partseg() 306 | 307 | if self.class_choice != None: 308 | id_choice = self.cat2id[self.class_choice] 309 | indices = (self.label == id_choice).squeeze() 310 | self.data = self.data[indices] 311 | self.label = self.label[indices] 312 | self.seg = self.seg[indices] 313 | self.seg_num_all = self.seg_num[id_choice] 314 | self.seg_start_index = self.index_start[id_choice] 315 | else: 316 | self.seg_num_all = 50 317 | self.seg_start_index = 0 318 | 319 | 320 | def __getitem__(self, item): 321 | pointcloud = self.data[item][:self.num_points] 322 | label = self.label[item] 323 | seg = self.seg[item][:self.num_points] 324 | if self.partition == 'trainval': 325 | # pointcloud = translate_pointcloud(pointcloud) 326 | indices = list(range(pointcloud.shape[0])) 327 | np.random.shuffle(indices) 328 | pointcloud = pointcloud[indices] 329 | seg = seg[indices] 330 | return pointcloud, label, seg 331 | 332 | def __len__(self): 333 | return self.data.shape[0] 334 | 335 | 336 | class S3DIS(Dataset): 337 | def __init__(self, num_points=4096, partition='train', test_area='1'): 338 | self.data, self.seg = load_data_semseg(partition, test_area) 339 | self.num_points = num_points 340 | self.partition = partition 341 | self.semseg_colors = load_color_semseg() 342 | 343 | def __getitem__(self, item): 344 | pointcloud = self.data[item][:self.num_points] 345 | seg = self.seg[item][:self.num_points] 346 | if self.partition == 'train': 347 | indices = list(range(pointcloud.shape[0])) 348 | np.random.shuffle(indices) 349 | pointcloud = pointcloud[indices] 350 | seg = seg[indices] 351 | seg = torch.LongTensor(seg) 352 | return pointcloud, seg 353 | 354 | def __len__(self): 355 | return self.data.shape[0] 356 | 357 | 358 | class ScanNet(Dataset): 359 | def __init__(self, num_point=8192, partition='train', 360 | data_root='scannet', classes=20, block_size=1.5, 361 | sample_rate=1.0, transform=None, use_rgb=False): 362 | self.partition = partition 363 | self.num_point = num_point 364 | self.block_size = block_size 365 | self.transform = transform 366 | xyz_all = [] 367 | label_all = [] 368 | if not isinstance(partition, list): 369 | partition = [partition] 370 | for i in partition: 371 | data_file = os.path.join( 372 | data_root, 'scannet_{}_rgb21c_pointid.pickle'.format(i)) 373 | file_pickle = open(data_file, 'rb') 374 | _xyz_all = pickle.load(file_pickle) 375 | _label_all = pickle.load(file_pickle) 376 | file_pickle.close() 377 | xyz_all.append(_xyz_all) 378 | label_all.append(_label_all) 379 | xyz_all = np.hstack(xyz_all) 380 | label_all = np.hstack(label_all) 381 | self.label_all = [] # for change 0-20 to 0-19 + 255 382 | self.room_coord_min, self.room_coord_max = [], [] 383 | num_point_all = [] 384 | for index in range(len(xyz_all)): 385 | xyz, label = xyz_all[index], label_all[index] # xyzrgb, N*6; l, N 386 | coord_min, coord_max = np.amin( 387 | xyz, axis=0)[:3], np.amax(xyz, axis=0)[:3] 388 | self.room_coord_min.append(coord_min) 389 | self.room_coord_max.append(coord_max) 390 | num_point_all.append(label.size) 391 | # we have set all ignore_class to 0 392 | # class 0 is also ignored 393 | # so we set all them as 255 394 | label_new = label - 1 395 | label_new[label == 0] = 255 396 | self.label_all.append(label_new.astype(np.uint8)) 397 | sample_prob = num_point_all / np.sum(num_point_all) 398 | num_iter = int(np.sum(num_point_all) * sample_rate / num_point) 399 | room_idxs = [] 400 | for index in range(len(xyz_all)): 401 | room_idxs.extend( 402 | [index] * int(round(sample_prob[index] * num_iter))) 403 | self.room_idxs = np.array(room_idxs) 404 | self.xyz_all = xyz_all 405 | 406 | # whether load RGB information 407 | self.use_rgb = use_rgb 408 | 409 | print("Totally {} samples in {} set.".format(len(self.room_idxs), partition)) 410 | 411 | def __getitem__(self, idx): 412 | room_idx = self.room_idxs[idx] 413 | points = self.xyz_all[room_idx] # N * 6 414 | if not self.use_rgb: 415 | points = points[:, :3] 416 | labels = self.label_all[room_idx] # N 417 | N_points = points.shape[0] 418 | 419 | for i in range(10): 420 | center = points[np.random.choice(N_points)][:3] 421 | block_min = center - [self.block_size / 422 | 2.0, self.block_size / 2.0, 0] 423 | block_max = center + [self.block_size / 424 | 2.0, self.block_size / 2.0, 0] 425 | block_min[2] = self.room_coord_min[room_idx][2] 426 | block_max[2] = self.room_coord_max[room_idx][2] 427 | point_idxs = np.where((points[:, 0] >= block_min[0]) & 428 | (points[:, 0] <= block_max[0]) & 429 | (points[:, 1] >= block_min[1]) & 430 | (points[:, 1] <= block_max[1]))[0] 431 | if point_idxs.size == 0: 432 | continue 433 | vidx = np.ceil((points[point_idxs, :3] - block_min) / 434 | (block_max - block_min) * [31.0, 31.0, 62.0]) 435 | vidx = np.unique(vidx[:, 0] * 31.0 * 62.0 + 436 | vidx[:, 1] * 62.0 + vidx[:, 2]) 437 | if ((labels[point_idxs] != 255).sum() / point_idxs.size >= 0.7) and (vidx.size/31.0/31.0/62.0 >= 0.02): 438 | break 439 | 440 | if point_idxs.size >= self.num_point: 441 | selected_point_idxs = np.random.choice( 442 | point_idxs, self.num_point, replace=False) 443 | else: 444 | selected_point_idxs = np.random.choice( 445 | point_idxs, self.num_point, replace=True) 446 | # normalize 447 | selected_points = points[selected_point_idxs, :] # num_point * 3/6 448 | num_feats = 9 if self.use_rgb else 6 449 | current_points = np.zeros( 450 | (self.num_point, num_feats)) # num_point * 6/9 451 | current_points[:, -3] = selected_points[:, 0] / \ 452 | self.room_coord_max[room_idx][0] 453 | current_points[:, -2] = selected_points[:, 1] / \ 454 | self.room_coord_max[room_idx][1] 455 | current_points[:, -1] = selected_points[:, 2] / \ 456 | self.room_coord_max[room_idx][2] 457 | selected_points[:, 0] = selected_points[:, 0] - center[0] 458 | selected_points[:, 1] = selected_points[:, 1] - center[1] 459 | current_points[:, 0:3] = selected_points[:, 0:3] 460 | if self.use_rgb: 461 | current_points[:, 3:6] = selected_points[:, 3:6] / 255. 462 | current_labels = labels[selected_point_idxs] 463 | if self.transform is not None: 464 | current_points, current_labels = self.transform( 465 | current_points, current_labels) 466 | return current_points, current_labels 467 | 468 | def __len__(self): 469 | return len(self.room_idxs) 470 | 471 | if __name__ == '__main__': 472 | train = ModelNet40(1024) 473 | test = ModelNet40(1024, 'test') 474 | data, label = train[0] 475 | print(data.shape) 476 | print(label.shape) 477 | 478 | trainval = ShapeNetPart(2048, 'trainval') 479 | test = ShapeNetPart(2048, 'test') 480 | data, label, seg = trainval[0] 481 | print(data.shape) 482 | print(label.shape) 483 | print(seg.shape) 484 | 485 | train = S3DIS(4096) 486 | test = S3DIS(4096, 'test') 487 | data, seg = train[0] 488 | print(data.shape) 489 | print(seg.shape) 490 | 491 | train = ScanNet(8192) 492 | test = ScanNet(8192, 'test') 493 | data, seg = train[0] 494 | print(data.shape) 495 | print(seg.shape) 496 | -------------------------------------------------------------------------------- /image/DGCNN.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/antao97/dgcnn.pytorch/07d534c2702905010ec9991619f552d8cacae45b/image/DGCNN.jpg -------------------------------------------------------------------------------- /image/partseg_colors.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/antao97/dgcnn.pytorch/07d534c2702905010ec9991619f552d8cacae45b/image/partseg_colors.png -------------------------------------------------------------------------------- /image/partseg_visu.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/antao97/dgcnn.pytorch/07d534c2702905010ec9991619f552d8cacae45b/image/partseg_visu.png -------------------------------------------------------------------------------- /image/semseg_s3dis_colors.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/antao97/dgcnn.pytorch/07d534c2702905010ec9991619f552d8cacae45b/image/semseg_s3dis_colors.png -------------------------------------------------------------------------------- /image/semseg_s3dis_visu.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/antao97/dgcnn.pytorch/07d534c2702905010ec9991619f552d8cacae45b/image/semseg_s3dis_visu.png -------------------------------------------------------------------------------- /image/semseg_scannet_colors.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/antao97/dgcnn.pytorch/07d534c2702905010ec9991619f552d8cacae45b/image/semseg_scannet_colors.png -------------------------------------------------------------------------------- /image/semseg_scannet_visu.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/antao97/dgcnn.pytorch/07d534c2702905010ec9991619f552d8cacae45b/image/semseg_scannet_visu.png -------------------------------------------------------------------------------- /main_cls.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | """ 4 | @Author: Yue Wang 5 | @Contact: yuewangx@mit.edu 6 | @File: main_cls.py 7 | @Time: 2018/10/13 10:39 PM 8 | 9 | Modified by 10 | @Author: An Tao 11 | @Contact: ta19@mails.tsinghua.edu.cn 12 | @Time: 2019/12/30 9:32 PM 13 | """ 14 | 15 | 16 | from __future__ import print_function 17 | import os 18 | import argparse 19 | import torch 20 | import torch.nn as nn 21 | import torch.nn.functional as F 22 | import torch.optim as optim 23 | from torch.optim.lr_scheduler import CosineAnnealingLR, StepLR 24 | from data import ModelNet40 25 | from model import PointNet, DGCNN_cls 26 | import numpy as np 27 | from torch.utils.data import DataLoader 28 | from util import cal_loss, IOStream 29 | import sklearn.metrics as metrics 30 | 31 | 32 | def _init_(): 33 | if not os.path.exists('outputs'): 34 | os.makedirs('outputs') 35 | if not os.path.exists('outputs/'+args.exp_name): 36 | os.makedirs('outputs/'+args.exp_name) 37 | if not os.path.exists('outputs/'+args.exp_name+'/'+'models'): 38 | os.makedirs('outputs/'+args.exp_name+'/'+'models') 39 | os.system('cp main_cls.py outputs'+'/'+args.exp_name+'/'+'main_cls.py.backup') 40 | os.system('cp model.py outputs' + '/' + args.exp_name + '/' + 'model.py.backup') 41 | os.system('cp util.py outputs' + '/' + args.exp_name + '/' + 'util.py.backup') 42 | os.system('cp data.py outputs' + '/' + args.exp_name + '/' + 'data.py.backup') 43 | 44 | def train(args, io): 45 | train_loader = DataLoader(ModelNet40(partition='train', num_points=args.num_points), num_workers=8, 46 | batch_size=args.batch_size, shuffle=True, drop_last=True) 47 | test_loader = DataLoader(ModelNet40(partition='test', num_points=args.num_points), num_workers=8, 48 | batch_size=args.test_batch_size, shuffle=True, drop_last=False) 49 | 50 | device = torch.device("cuda" if args.cuda else "cpu") 51 | 52 | #Try to load models 53 | if args.model == 'pointnet': 54 | model = PointNet(args).to(device) 55 | elif args.model == 'dgcnn': 56 | model = DGCNN_cls(args).to(device) 57 | else: 58 | raise Exception("Not implemented") 59 | 60 | print(str(model)) 61 | 62 | model = nn.DataParallel(model) 63 | print("Let's use", torch.cuda.device_count(), "GPUs!") 64 | 65 | if args.use_sgd: 66 | print("Use SGD") 67 | opt = optim.SGD(model.parameters(), lr=args.lr*100, momentum=args.momentum, weight_decay=1e-4) 68 | else: 69 | print("Use Adam") 70 | opt = optim.Adam(model.parameters(), lr=args.lr, weight_decay=1e-4) 71 | 72 | if args.scheduler == 'cos': 73 | scheduler = CosineAnnealingLR(opt, args.epochs, eta_min=1e-3) 74 | elif args.scheduler == 'step': 75 | scheduler = StepLR(opt, step_size=20, gamma=0.7) 76 | 77 | criterion = cal_loss 78 | 79 | best_test_acc = 0 80 | for epoch in range(args.epochs): 81 | #################### 82 | # Train 83 | #################### 84 | train_loss = 0.0 85 | count = 0.0 86 | model.train() 87 | train_pred = [] 88 | train_true = [] 89 | for data, label in train_loader: 90 | data, label = data.to(device), label.to(device).squeeze() 91 | data = data.permute(0, 2, 1) 92 | batch_size = data.size()[0] 93 | opt.zero_grad() 94 | logits = model(data) 95 | loss = criterion(logits, label) 96 | loss.backward() 97 | opt.step() 98 | preds = logits.max(dim=1)[1] 99 | count += batch_size 100 | train_loss += loss.item() * batch_size 101 | train_true.append(label.cpu().numpy()) 102 | train_pred.append(preds.detach().cpu().numpy()) 103 | if args.scheduler == 'cos': 104 | scheduler.step() 105 | elif args.scheduler == 'step': 106 | if opt.param_groups[0]['lr'] > 1e-5: 107 | scheduler.step() 108 | if opt.param_groups[0]['lr'] < 1e-5: 109 | for param_group in opt.param_groups: 110 | param_group['lr'] = 1e-5 111 | 112 | train_true = np.concatenate(train_true) 113 | train_pred = np.concatenate(train_pred) 114 | outstr = 'Train %d, loss: %.6f, train acc: %.6f, train avg acc: %.6f' % (epoch, 115 | train_loss*1.0/count, 116 | metrics.accuracy_score( 117 | train_true, train_pred), 118 | metrics.balanced_accuracy_score( 119 | train_true, train_pred)) 120 | io.cprint(outstr) 121 | 122 | #################### 123 | # Test 124 | #################### 125 | test_loss = 0.0 126 | count = 0.0 127 | model.eval() 128 | test_pred = [] 129 | test_true = [] 130 | for data, label in test_loader: 131 | data, label = data.to(device), label.to(device).squeeze() 132 | data = data.permute(0, 2, 1) 133 | batch_size = data.size()[0] 134 | logits = model(data) 135 | loss = criterion(logits, label) 136 | preds = logits.max(dim=1)[1] 137 | count += batch_size 138 | test_loss += loss.item() * batch_size 139 | test_true.append(label.cpu().numpy()) 140 | test_pred.append(preds.detach().cpu().numpy()) 141 | test_true = np.concatenate(test_true) 142 | test_pred = np.concatenate(test_pred) 143 | test_acc = metrics.accuracy_score(test_true, test_pred) 144 | avg_per_class_acc = metrics.balanced_accuracy_score(test_true, test_pred) 145 | outstr = 'Test %d, loss: %.6f, test acc: %.6f, test avg acc: %.6f' % (epoch, 146 | test_loss*1.0/count, 147 | test_acc, 148 | avg_per_class_acc) 149 | io.cprint(outstr) 150 | if test_acc >= best_test_acc: 151 | best_test_acc = test_acc 152 | torch.save(model.state_dict(), 'outputs/%s/models/model.t7' % args.exp_name) 153 | 154 | 155 | def test(args, io): 156 | test_loader = DataLoader(ModelNet40(partition='test', num_points=args.num_points), 157 | batch_size=args.test_batch_size, shuffle=True, drop_last=False) 158 | 159 | device = torch.device("cuda" if args.cuda else "cpu") 160 | 161 | #Try to load models 162 | if args.model == 'pointnet': 163 | model = PointNet(args).to(device) 164 | elif args.model == 'dgcnn': 165 | model = DGCNN_cls(args).to(device) 166 | else: 167 | raise Exception("Not implemented") 168 | 169 | model = nn.DataParallel(model) 170 | model.load_state_dict(torch.load(args.model_path)) 171 | model = model.eval() 172 | test_acc = 0.0 173 | count = 0.0 174 | test_true = [] 175 | test_pred = [] 176 | for data, label in test_loader: 177 | 178 | data, label = data.to(device), label.to(device).squeeze() 179 | data = data.permute(0, 2, 1) 180 | batch_size = data.size()[0] 181 | logits = model(data) 182 | preds = logits.max(dim=1)[1] 183 | test_true.append(label.cpu().numpy()) 184 | test_pred.append(preds.detach().cpu().numpy()) 185 | test_true = np.concatenate(test_true) 186 | test_pred = np.concatenate(test_pred) 187 | test_acc = metrics.accuracy_score(test_true, test_pred) 188 | avg_per_class_acc = metrics.balanced_accuracy_score(test_true, test_pred) 189 | outstr = 'Test :: test acc: %.6f, test avg acc: %.6f'%(test_acc, avg_per_class_acc) 190 | io.cprint(outstr) 191 | 192 | 193 | if __name__ == "__main__": 194 | # Training settings 195 | parser = argparse.ArgumentParser(description='Point Cloud Recognition') 196 | parser.add_argument('--exp_name', type=str, default='exp', metavar='N', 197 | help='Name of the experiment') 198 | parser.add_argument('--model', type=str, default='dgcnn', metavar='N', 199 | choices=['pointnet', 'dgcnn'], 200 | help='Model to use, [pointnet, dgcnn]') 201 | parser.add_argument('--dataset', type=str, default='modelnet40', metavar='N', 202 | choices=['modelnet40']) 203 | parser.add_argument('--batch_size', type=int, default=32, metavar='batch_size', 204 | help='Size of batch)') 205 | parser.add_argument('--test_batch_size', type=int, default=16, metavar='batch_size', 206 | help='Size of batch)') 207 | parser.add_argument('--epochs', type=int, default=250, metavar='N', 208 | help='number of episode to train ') 209 | parser.add_argument('--use_sgd', type=bool, default=True, 210 | help='Use SGD') 211 | parser.add_argument('--lr', type=float, default=0.001, metavar='LR', 212 | help='learning rate (default: 0.001, 0.1 if using sgd)') 213 | parser.add_argument('--momentum', type=float, default=0.9, metavar='M', 214 | help='SGD momentum (default: 0.9)') 215 | parser.add_argument('--scheduler', type=str, default='cos', metavar='N', 216 | choices=['cos', 'step'], 217 | help='Scheduler to use, [cos, step]') 218 | parser.add_argument('--no_cuda', type=bool, default=False, 219 | help='enables CUDA training') 220 | parser.add_argument('--seed', type=int, default=1, metavar='S', 221 | help='random seed (default: 1)') 222 | parser.add_argument('--eval', type=bool, default=False, 223 | help='evaluate the model') 224 | parser.add_argument('--num_points', type=int, default=1024, 225 | help='num of points to use') 226 | parser.add_argument('--dropout', type=float, default=0.5, 227 | help='initial dropout rate') 228 | parser.add_argument('--emb_dims', type=int, default=1024, metavar='N', 229 | help='Dimension of embeddings') 230 | parser.add_argument('--k', type=int, default=20, metavar='N', 231 | help='Num of nearest neighbors to use') 232 | parser.add_argument('--model_path', type=str, default='', metavar='N', 233 | help='Pretrained model path') 234 | args = parser.parse_args() 235 | 236 | _init_() 237 | 238 | io = IOStream('outputs/' + args.exp_name + '/run.log') 239 | io.cprint(str(args)) 240 | 241 | args.cuda = not args.no_cuda and torch.cuda.is_available() 242 | torch.manual_seed(args.seed) 243 | if args.cuda: 244 | io.cprint( 245 | 'Using GPU : ' + str(torch.cuda.current_device()) + ' from ' + str(torch.cuda.device_count()) + ' devices') 246 | torch.cuda.manual_seed(args.seed) 247 | else: 248 | io.cprint('Using CPU') 249 | 250 | if not args.eval: 251 | train(args, io) 252 | else: 253 | test(args, io) 254 | -------------------------------------------------------------------------------- /main_partseg.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | """ 4 | @Author: An Tao, Pengliang Ji 5 | @Contact: ta19@mails.tsinghua.edu.cn, jpl1723@buaa.edu.cn 6 | @File: main_partseg.py 7 | @Time: 2021/7/20 7:49 PM 8 | """ 9 | 10 | 11 | from __future__ import print_function 12 | import os 13 | import argparse 14 | import torch 15 | import torch.nn as nn 16 | import torch.nn.functional as F 17 | import torch.optim as optim 18 | from torch.optim.lr_scheduler import CosineAnnealingLR, StepLR 19 | from data import ShapeNetPart 20 | from model import DGCNN_partseg 21 | import numpy as np 22 | from torch.utils.data import DataLoader 23 | from util import cal_loss, IOStream 24 | import sklearn.metrics as metrics 25 | from plyfile import PlyData, PlyElement 26 | 27 | global class_cnts 28 | class_indexs = np.zeros((16,), dtype=int) 29 | global visual_warning 30 | visual_warning = True 31 | 32 | class_choices = ['airplane', 'bag', 'cap', 'car', 'chair', 'earphone', 'guitar', 'knife', 'lamp', 'laptop', 'motorbike', 'mug', 'pistol', 'rocket', 'skateboard', 'table'] 33 | seg_num = [4, 2, 2, 4, 4, 3, 3, 2, 4, 2, 6, 2, 3, 3, 3, 3] 34 | index_start = [0, 4, 6, 8, 12, 16, 19, 22, 24, 28, 30, 36, 38, 41, 44, 47] 35 | 36 | 37 | def _init_(): 38 | if not os.path.exists('outputs'): 39 | os.makedirs('outputs') 40 | if not os.path.exists('outputs/'+args.exp_name): 41 | os.makedirs('outputs/'+args.exp_name) 42 | if not os.path.exists('outputs/'+args.exp_name+'/'+'models'): 43 | os.makedirs('outputs/'+args.exp_name+'/'+'models') 44 | if not os.path.exists('outputs/'+args.exp_name+'/'+'visualization'): 45 | os.makedirs('outputs/'+args.exp_name+'/'+'visualization') 46 | os.system('cp main_partseg.py outputs'+'/'+args.exp_name+'/'+'main_partseg.py.backup') 47 | os.system('cp model.py outputs' + '/' + args.exp_name + '/' + 'model.py.backup') 48 | os.system('cp util.py outputs' + '/' + args.exp_name + '/' + 'util.py.backup') 49 | os.system('cp data.py outputs' + '/' + args.exp_name + '/' + 'data.py.backup') 50 | 51 | 52 | def calculate_shape_IoU(pred_np, seg_np, label, class_choice, visual=False): 53 | if not visual: 54 | label = label.squeeze() 55 | shape_ious = [] 56 | for shape_idx in range(seg_np.shape[0]): 57 | if not class_choice: 58 | start_index = index_start[label[shape_idx]] 59 | num = seg_num[label[shape_idx]] 60 | parts = range(start_index, start_index + num) 61 | else: 62 | parts = range(seg_num[label[0]]) 63 | part_ious = [] 64 | for part in parts: 65 | I = np.sum(np.logical_and(pred_np[shape_idx] == part, seg_np[shape_idx] == part)) 66 | U = np.sum(np.logical_or(pred_np[shape_idx] == part, seg_np[shape_idx] == part)) 67 | if U == 0: 68 | iou = 1 # If the union of groundtruth and prediction points is empty, then count part IoU as 1 69 | else: 70 | iou = I / float(U) 71 | part_ious.append(iou) 72 | shape_ious.append(np.mean(part_ious)) 73 | return shape_ious 74 | 75 | 76 | def visualization(visu, visu_format, data, pred, seg, label, partseg_colors, class_choice): 77 | global class_indexs 78 | global visual_warning 79 | visu = visu.split('_') 80 | for i in range(0, data.shape[0]): 81 | RGB = [] 82 | RGB_gt = [] 83 | skip = False 84 | classname = class_choices[int(label[i])] 85 | class_index = class_indexs[int(label[i])] 86 | if visu[0] != 'all': 87 | if len(visu) != 1: 88 | if visu[0] != classname or visu[1] != str(class_index): 89 | skip = True 90 | else: 91 | visual_warning = False 92 | elif visu[0] != classname: 93 | skip = True 94 | else: 95 | visual_warning = False 96 | elif class_choice != None: 97 | skip = True 98 | else: 99 | visual_warning = False 100 | if skip: 101 | class_indexs[int(label[i])] = class_indexs[int(label[i])] + 1 102 | else: 103 | if not os.path.exists('outputs/'+args.exp_name+'/'+'visualization'+'/'+classname): 104 | os.makedirs('outputs/'+args.exp_name+'/'+'visualization'+'/'+classname) 105 | for j in range(0, data.shape[2]): 106 | RGB.append(partseg_colors[int(pred[i][j])]) 107 | RGB_gt.append(partseg_colors[int(seg[i][j])]) 108 | pred_np = [] 109 | seg_np = [] 110 | pred_np.append(pred[i].cpu().numpy()) 111 | seg_np.append(seg[i].cpu().numpy()) 112 | xyz_np = data[i].cpu().numpy() 113 | xyzRGB = np.concatenate((xyz_np.transpose(1, 0), np.array(RGB)), axis=1) 114 | xyzRGB_gt = np.concatenate((xyz_np.transpose(1, 0), np.array(RGB_gt)), axis=1) 115 | IoU = calculate_shape_IoU(np.array(pred_np), np.array(seg_np), label[i].cpu().numpy(), class_choice, visual=True) 116 | IoU = str(round(IoU[0], 4)) 117 | filepath = 'outputs/'+args.exp_name+'/'+'visualization'+'/'+classname+'/'+classname+'_'+str(class_index)+'_pred_'+IoU+'.'+visu_format 118 | filepath_gt = 'outputs/'+args.exp_name+'/'+'visualization'+'/'+classname+'/'+classname+'_'+str(class_index)+'_gt.'+visu_format 119 | if visu_format=='txt': 120 | np.savetxt(filepath, xyzRGB, fmt='%s', delimiter=' ') 121 | np.savetxt(filepath_gt, xyzRGB_gt, fmt='%s', delimiter=' ') 122 | print('TXT visualization file saved in', filepath) 123 | print('TXT visualization file saved in', filepath_gt) 124 | elif visu_format=='ply': 125 | xyzRGB = [(xyzRGB[i, 0], xyzRGB[i, 1], xyzRGB[i, 2], xyzRGB[i, 3], xyzRGB[i, 4], xyzRGB[i, 5]) for i in range(xyzRGB.shape[0])] 126 | xyzRGB_gt = [(xyzRGB_gt[i, 0], xyzRGB_gt[i, 1], xyzRGB_gt[i, 2], xyzRGB_gt[i, 3], xyzRGB_gt[i, 4], xyzRGB_gt[i, 5]) for i in range(xyzRGB_gt.shape[0])] 127 | vertex = PlyElement.describe(np.array(xyzRGB, dtype=[('x', 'f4'), ('y', 'f4'), ('z', 'f4'), ('red', 'u1'), ('green', 'u1'), ('blue', 'u1')]), 'vertex') 128 | PlyData([vertex]).write(filepath) 129 | vertex = PlyElement.describe(np.array(xyzRGB_gt, dtype=[('x', 'f4'), ('y', 'f4'), ('z', 'f4'), ('red', 'u1'), ('green', 'u1'), ('blue', 'u1')]), 'vertex') 130 | PlyData([vertex]).write(filepath_gt) 131 | print('PLY visualization file saved in', filepath) 132 | print('PLY visualization file saved in', filepath_gt) 133 | else: 134 | print('ERROR!! Unknown visualization format: %s, please use txt or ply.' % \ 135 | (visu_format)) 136 | exit() 137 | class_indexs[int(label[i])] = class_indexs[int(label[i])] + 1 138 | 139 | 140 | def train(args, io): 141 | train_dataset = ShapeNetPart(partition='trainval', num_points=args.num_points, class_choice=args.class_choice) 142 | if (len(train_dataset) < 100): 143 | drop_last = False 144 | else: 145 | drop_last = True 146 | train_loader = DataLoader(train_dataset, num_workers=8, batch_size=args.batch_size, shuffle=True, drop_last=drop_last) 147 | test_loader = DataLoader(ShapeNetPart(partition='test', num_points=args.num_points, class_choice=args.class_choice), 148 | num_workers=8, batch_size=args.test_batch_size, shuffle=True, drop_last=False) 149 | 150 | device = torch.device("cuda" if args.cuda else "cpu") 151 | 152 | #Try to load models 153 | seg_num_all = train_loader.dataset.seg_num_all 154 | seg_start_index = train_loader.dataset.seg_start_index 155 | if args.model == 'dgcnn': 156 | model = DGCNN_partseg(args, seg_num_all).to(device) 157 | else: 158 | raise Exception("Not implemented") 159 | print(str(model)) 160 | 161 | model = nn.DataParallel(model) 162 | print("Let's use", torch.cuda.device_count(), "GPUs!") 163 | 164 | if args.use_sgd: 165 | print("Use SGD") 166 | opt = optim.SGD(model.parameters(), lr=args.lr*100, momentum=args.momentum, weight_decay=1e-4) 167 | else: 168 | print("Use Adam") 169 | opt = optim.Adam(model.parameters(), lr=args.lr, weight_decay=1e-4) 170 | 171 | if args.scheduler == 'cos': 172 | scheduler = CosineAnnealingLR(opt, args.epochs, eta_min=1e-3) 173 | elif args.scheduler == 'step': 174 | scheduler = StepLR(opt, step_size=20, gamma=0.5) 175 | 176 | criterion = cal_loss 177 | 178 | best_test_iou = 0 179 | for epoch in range(args.epochs): 180 | #################### 181 | # Train 182 | #################### 183 | train_loss = 0.0 184 | count = 0.0 185 | model.train() 186 | train_true_cls = [] 187 | train_pred_cls = [] 188 | train_true_seg = [] 189 | train_pred_seg = [] 190 | train_label_seg = [] 191 | for data, label, seg in train_loader: 192 | seg = seg - seg_start_index 193 | label_one_hot = np.zeros((label.shape[0], 16)) 194 | for idx in range(label.shape[0]): 195 | label_one_hot[idx, label[idx]] = 1 196 | label_one_hot = torch.from_numpy(label_one_hot.astype(np.float32)) 197 | data, label_one_hot, seg = data.to(device), label_one_hot.to(device), seg.to(device) 198 | data = data.permute(0, 2, 1) 199 | batch_size = data.size()[0] 200 | opt.zero_grad() 201 | seg_pred = model(data, label_one_hot) 202 | seg_pred = seg_pred.permute(0, 2, 1).contiguous() 203 | loss = criterion(seg_pred.view(-1, seg_num_all), seg.view(-1,1).squeeze()) 204 | loss.backward() 205 | opt.step() 206 | pred = seg_pred.max(dim=2)[1] # (batch_size, num_points) 207 | count += batch_size 208 | train_loss += loss.item() * batch_size 209 | seg_np = seg.cpu().numpy() # (batch_size, num_points) 210 | pred_np = pred.detach().cpu().numpy() # (batch_size, num_points) 211 | train_true_cls.append(seg_np.reshape(-1)) # (batch_size * num_points) 212 | train_pred_cls.append(pred_np.reshape(-1)) # (batch_size * num_points) 213 | train_true_seg.append(seg_np) 214 | train_pred_seg.append(pred_np) 215 | train_label_seg.append(label.reshape(-1)) 216 | if args.scheduler == 'cos': 217 | scheduler.step() 218 | elif args.scheduler == 'step': 219 | if opt.param_groups[0]['lr'] > 1e-5: 220 | scheduler.step() 221 | if opt.param_groups[0]['lr'] < 1e-5: 222 | for param_group in opt.param_groups: 223 | param_group['lr'] = 1e-5 224 | train_true_cls = np.concatenate(train_true_cls) 225 | train_pred_cls = np.concatenate(train_pred_cls) 226 | train_acc = metrics.accuracy_score(train_true_cls, train_pred_cls) 227 | avg_per_class_acc = metrics.balanced_accuracy_score(train_true_cls, train_pred_cls) 228 | train_true_seg = np.concatenate(train_true_seg, axis=0) 229 | train_pred_seg = np.concatenate(train_pred_seg, axis=0) 230 | train_label_seg = np.concatenate(train_label_seg) 231 | train_ious = calculate_shape_IoU(train_pred_seg, train_true_seg, train_label_seg, args.class_choice) 232 | outstr = 'Train %d, loss: %.6f, train acc: %.6f, train avg acc: %.6f, train iou: %.6f' % (epoch, 233 | train_loss*1.0/count, 234 | train_acc, 235 | avg_per_class_acc, 236 | np.mean(train_ious)) 237 | io.cprint(outstr) 238 | 239 | #################### 240 | # Test 241 | #################### 242 | test_loss = 0.0 243 | count = 0.0 244 | model.eval() 245 | test_true_cls = [] 246 | test_pred_cls = [] 247 | test_true_seg = [] 248 | test_pred_seg = [] 249 | test_label_seg = [] 250 | for data, label, seg in test_loader: 251 | seg = seg - seg_start_index 252 | label_one_hot = np.zeros((label.shape[0], 16)) 253 | for idx in range(label.shape[0]): 254 | label_one_hot[idx, label[idx]] = 1 255 | label_one_hot = torch.from_numpy(label_one_hot.astype(np.float32)) 256 | data, label_one_hot, seg = data.to(device), label_one_hot.to(device), seg.to(device) 257 | data = data.permute(0, 2, 1) 258 | batch_size = data.size()[0] 259 | seg_pred = model(data, label_one_hot) 260 | seg_pred = seg_pred.permute(0, 2, 1).contiguous() 261 | loss = criterion(seg_pred.view(-1, seg_num_all), seg.view(-1,1).squeeze()) 262 | pred = seg_pred.max(dim=2)[1] 263 | count += batch_size 264 | test_loss += loss.item() * batch_size 265 | seg_np = seg.cpu().numpy() 266 | pred_np = pred.detach().cpu().numpy() 267 | test_true_cls.append(seg_np.reshape(-1)) 268 | test_pred_cls.append(pred_np.reshape(-1)) 269 | test_true_seg.append(seg_np) 270 | test_pred_seg.append(pred_np) 271 | test_label_seg.append(label.reshape(-1)) 272 | test_true_cls = np.concatenate(test_true_cls) 273 | test_pred_cls = np.concatenate(test_pred_cls) 274 | test_acc = metrics.accuracy_score(test_true_cls, test_pred_cls) 275 | avg_per_class_acc = metrics.balanced_accuracy_score(test_true_cls, test_pred_cls) 276 | test_true_seg = np.concatenate(test_true_seg, axis=0) 277 | test_pred_seg = np.concatenate(test_pred_seg, axis=0) 278 | test_label_seg = np.concatenate(test_label_seg) 279 | test_ious = calculate_shape_IoU(test_pred_seg, test_true_seg, test_label_seg, args.class_choice) 280 | outstr = 'Test %d, loss: %.6f, test acc: %.6f, test avg acc: %.6f, test iou: %.6f' % (epoch, 281 | test_loss*1.0/count, 282 | test_acc, 283 | avg_per_class_acc, 284 | np.mean(test_ious)) 285 | io.cprint(outstr) 286 | if np.mean(test_ious) >= best_test_iou: 287 | best_test_iou = np.mean(test_ious) 288 | torch.save(model.state_dict(), 'outputs/%s/models/model.t7' % args.exp_name) 289 | 290 | 291 | def test(args, io): 292 | test_loader = DataLoader(ShapeNetPart(partition='test', num_points=args.num_points, class_choice=args.class_choice), 293 | batch_size=args.test_batch_size, shuffle=True, drop_last=False) 294 | device = torch.device("cuda" if args.cuda else "cpu") 295 | 296 | #Try to load models 297 | seg_num_all = test_loader.dataset.seg_num_all 298 | seg_start_index = test_loader.dataset.seg_start_index 299 | partseg_colors = test_loader.dataset.partseg_colors 300 | if args.model == 'dgcnn': 301 | model = DGCNN_partseg(args, seg_num_all).to(device) 302 | else: 303 | raise Exception("Not implemented") 304 | 305 | model = nn.DataParallel(model) 306 | model.load_state_dict(torch.load(args.model_path)) 307 | model = model.eval() 308 | test_acc = 0.0 309 | count = 0.0 310 | test_true_cls = [] 311 | test_pred_cls = [] 312 | test_true_seg = [] 313 | test_pred_seg = [] 314 | test_label_seg = [] 315 | for data, label, seg in test_loader: 316 | seg = seg - seg_start_index 317 | label_one_hot = np.zeros((label.shape[0], 16)) 318 | for idx in range(label.shape[0]): 319 | label_one_hot[idx, label[idx]] = 1 320 | label_one_hot = torch.from_numpy(label_one_hot.astype(np.float32)) 321 | data, label_one_hot, seg = data.to(device), label_one_hot.to(device), seg.to(device) 322 | data = data.permute(0, 2, 1) 323 | batch_size = data.size()[0] 324 | seg_pred = model(data, label_one_hot) 325 | seg_pred = seg_pred.permute(0, 2, 1).contiguous() 326 | pred = seg_pred.max(dim=2)[1] 327 | seg_np = seg.cpu().numpy() 328 | pred_np = pred.detach().cpu().numpy() 329 | test_true_cls.append(seg_np.reshape(-1)) 330 | test_pred_cls.append(pred_np.reshape(-1)) 331 | test_true_seg.append(seg_np) 332 | test_pred_seg.append(pred_np) 333 | test_label_seg.append(label.reshape(-1)) 334 | # visiualization 335 | visualization(args.visu, args.visu_format, data, pred, seg, label, partseg_colors, args.class_choice) 336 | if visual_warning and args.visu != '': 337 | print('Visualization Failed: You can only choose a point cloud shape to visualize within the scope of the test class') 338 | test_true_cls = np.concatenate(test_true_cls) 339 | test_pred_cls = np.concatenate(test_pred_cls) 340 | test_acc = metrics.accuracy_score(test_true_cls, test_pred_cls) 341 | avg_per_class_acc = metrics.balanced_accuracy_score(test_true_cls, test_pred_cls) 342 | test_true_seg = np.concatenate(test_true_seg, axis=0) 343 | test_pred_seg = np.concatenate(test_pred_seg, axis=0) 344 | test_label_seg = np.concatenate(test_label_seg) 345 | test_ious = calculate_shape_IoU(test_pred_seg, test_true_seg, test_label_seg, args.class_choice) 346 | outstr = 'Test :: test acc: %.6f, test avg acc: %.6f, test iou: %.6f' % (test_acc, 347 | avg_per_class_acc, 348 | np.mean(test_ious)) 349 | io.cprint(outstr) 350 | 351 | 352 | if __name__ == "__main__": 353 | # Training settings 354 | parser = argparse.ArgumentParser(description='Point Cloud Part Segmentation') 355 | parser.add_argument('--exp_name', type=str, default='exp', metavar='N', 356 | help='Name of the experiment') 357 | parser.add_argument('--model', type=str, default='dgcnn', metavar='N', 358 | choices=['dgcnn'], 359 | help='Model to use, [dgcnn]') 360 | parser.add_argument('--dataset', type=str, default='shapenetpart', metavar='N', 361 | choices=['shapenetpart']) 362 | parser.add_argument('--class_choice', type=str, default=None, metavar='N', 363 | choices=['airplane', 'bag', 'cap', 'car', 'chair', 364 | 'earphone', 'guitar', 'knife', 'lamp', 'laptop', 365 | 'motor', 'mug', 'pistol', 'rocket', 'skateboard', 'table']) 366 | parser.add_argument('--batch_size', type=int, default=32, metavar='batch_size', 367 | help='Size of batch)') 368 | parser.add_argument('--test_batch_size', type=int, default=16, metavar='batch_size', 369 | help='Size of batch)') 370 | parser.add_argument('--epochs', type=int, default=200, metavar='N', 371 | help='number of episode to train ') 372 | parser.add_argument('--use_sgd', type=bool, default=True, 373 | help='Use SGD') 374 | parser.add_argument('--lr', type=float, default=0.001, metavar='LR', 375 | help='learning rate (default: 0.001, 0.1 if using sgd)') 376 | parser.add_argument('--momentum', type=float, default=0.9, metavar='M', 377 | help='SGD momentum (default: 0.9)') 378 | parser.add_argument('--scheduler', type=str, default='cos', metavar='N', 379 | choices=['cos', 'step'], 380 | help='Scheduler to use, [cos, step]') 381 | parser.add_argument('--no_cuda', type=bool, default=False, 382 | help='enables CUDA training') 383 | parser.add_argument('--seed', type=int, default=1, metavar='S', 384 | help='random seed (default: 1)') 385 | parser.add_argument('--eval', type=bool, default=False, 386 | help='evaluate the model') 387 | parser.add_argument('--num_points', type=int, default=2048, 388 | help='num of points to use') 389 | parser.add_argument('--dropout', type=float, default=0.5, 390 | help='dropout rate') 391 | parser.add_argument('--emb_dims', type=int, default=1024, metavar='N', 392 | help='Dimension of embeddings') 393 | parser.add_argument('--k', type=int, default=40, metavar='N', 394 | help='Num of nearest neighbors to use') 395 | parser.add_argument('--model_path', type=str, default='', metavar='N', 396 | help='Pretrained model path') 397 | parser.add_argument('--visu', type=str, default='', 398 | help='visualize the model') 399 | parser.add_argument('--visu_format', type=str, default='ply', 400 | help='file format of visualization') 401 | args = parser.parse_args() 402 | 403 | _init_() 404 | 405 | io = IOStream('outputs/' + args.exp_name + '/run.log') 406 | io.cprint(str(args)) 407 | 408 | args.cuda = not args.no_cuda and torch.cuda.is_available() 409 | torch.manual_seed(args.seed) 410 | if args.cuda: 411 | io.cprint( 412 | 'Using GPU : ' + str(torch.cuda.current_device()) + ' from ' + str(torch.cuda.device_count()) + ' devices') 413 | torch.cuda.manual_seed(args.seed) 414 | else: 415 | io.cprint('Using CPU') 416 | 417 | if not args.eval: 418 | train(args, io) 419 | else: 420 | test(args, io) 421 | -------------------------------------------------------------------------------- /prepare_data/collect_indoor3d_data.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 4 | ROOT_DIR = os.path.dirname(BASE_DIR) 5 | DATA_PATH = os.path.join(ROOT_DIR, 'data/Stanford3dDataset_v1.2_Aligned_Version') 6 | import indoor3d_util 7 | 8 | anno_paths = [line.rstrip() for line in open(os.path.join(BASE_DIR, 'meta/anno_paths.txt'))] 9 | anno_paths = [os.path.join(DATA_PATH, p) for p in anno_paths] 10 | 11 | output_folder = os.path.join(ROOT_DIR, 'data/stanford_indoor3d') 12 | if not os.path.exists(output_folder): 13 | os.mkdir(output_folder) 14 | 15 | revise_file = os.path.join(DATA_PATH, "Area_5/hallway_6/Annotations/ceiling_1.txt") 16 | with open(revise_file, "r") as f: 17 | data = f.read() 18 | data = data[:5545347] + ' ' + data[5545348:] 19 | f.close() 20 | with open(revise_file, "w") as f: 21 | f.write(data) 22 | f.close() 23 | 24 | for anno_path in anno_paths: 25 | print(anno_path) 26 | elements = anno_path.split('/') 27 | out_filename = elements[-3]+'_'+elements[-2]+'.npy' # Area_1_hallway_1.npy 28 | indoor3d_util.collect_point_label(anno_path, os.path.join(output_folder, out_filename), 'numpy') -------------------------------------------------------------------------------- /prepare_data/data_prep_util.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 4 | sys.path.append(BASE_DIR) 5 | from plyfile import (PlyData, PlyElement, make2d, PlyParseError, PlyProperty) 6 | import numpy as np 7 | import h5py 8 | 9 | SAMPLING_BIN = os.path.join(BASE_DIR, 'third_party/mesh_sampling/build/pcsample') 10 | 11 | SAMPLING_POINT_NUM = 2048 12 | SAMPLING_LEAF_SIZE = 0.005 13 | 14 | MODELNET40_PATH = '../datasets/modelnet40' 15 | def export_ply(pc, filename): 16 | vertex = np.zeros(pc.shape[0], dtype=[('x', 'f4'), ('y', 'f4'), ('z', 'f4')]) 17 | for i in range(pc.shape[0]): 18 | vertex[i] = (pc[i][0], pc[i][1], pc[i][2]) 19 | ply_out = PlyData([PlyElement.describe(vertex, 'vertex', comments=['vertices'])]) 20 | ply_out.write(filename) 21 | 22 | # Sample points on the obj shape 23 | def get_sampling_command(obj_filename, ply_filename): 24 | cmd = SAMPLING_BIN + ' ' + obj_filename 25 | cmd += ' ' + ply_filename 26 | cmd += ' -n_samples %d ' % SAMPLING_POINT_NUM 27 | cmd += ' -leaf_size %f ' % SAMPLING_LEAF_SIZE 28 | return cmd 29 | 30 | # -------------------------------------------------------------- 31 | # Following are the helper functions to load MODELNET40 shapes 32 | # -------------------------------------------------------------- 33 | 34 | # Read in the list of categories in MODELNET40 35 | def get_category_names(): 36 | shape_names_file = os.path.join(MODELNET40_PATH, 'shape_names.txt') 37 | shape_names = [line.rstrip() for line in open(shape_names_file)] 38 | return shape_names 39 | 40 | # Return all the filepaths for the shapes in MODELNET40 41 | def get_obj_filenames(): 42 | obj_filelist_file = os.path.join(MODELNET40_PATH, 'filelist.txt') 43 | obj_filenames = [os.path.join(MODELNET40_PATH, line.rstrip()) for line in open(obj_filelist_file)] 44 | print('Got %d obj files in modelnet40.' % len(obj_filenames)) 45 | return obj_filenames 46 | 47 | # Helper function to create the father folder and all subdir folders if not exist 48 | def batch_mkdir(output_folder, subdir_list): 49 | if not os.path.exists(output_folder): 50 | os.mkdir(output_folder) 51 | for subdir in subdir_list: 52 | if not os.path.exists(os.path.join(output_folder, subdir)): 53 | os.mkdir(os.path.join(output_folder, subdir)) 54 | 55 | # ---------------------------------------------------------------- 56 | # Following are the helper functions to load save/load HDF5 files 57 | # ---------------------------------------------------------------- 58 | 59 | # Write numpy array data and label to h5_filename 60 | def save_h5_data_label_normal(h5_filename, data, label, normal, 61 | data_dtype='float32', label_dtype='uint8', normal_dtype='float32'): 62 | h5_fout = h5py.File(h5_filename) 63 | h5_fout.create_dataset( 64 | 'data', data=data, 65 | compression='gzip', compression_opts=4, 66 | dtype=data_dtype) 67 | h5_fout.create_dataset( 68 | 'normal', data=normal, 69 | compression='gzip', compression_opts=4, 70 | dtype=normal_dtype) 71 | h5_fout.create_dataset( 72 | 'label', data=label, 73 | compression='gzip', compression_opts=1, 74 | dtype=label_dtype) 75 | h5_fout.close() 76 | 77 | 78 | # Write numpy array data and label to h5_filename 79 | def save_h5(h5_filename, data, label, data_dtype='uint8', label_dtype='uint8'): 80 | h5_fout = h5py.File(h5_filename, "w") 81 | h5_fout.create_dataset( 82 | 'data', data=data, 83 | compression='gzip', compression_opts=4, 84 | dtype=data_dtype) 85 | h5_fout.create_dataset( 86 | 'label', data=label, 87 | compression='gzip', compression_opts=1, 88 | dtype=label_dtype) 89 | h5_fout.close() 90 | 91 | # Read numpy array data and label from h5_filename 92 | def load_h5_data_label_normal(h5_filename): 93 | f = h5py.File(h5_filename) 94 | data = f['data'][:] 95 | label = f['label'][:] 96 | normal = f['normal'][:] 97 | return (data, label, normal) 98 | 99 | # Read numpy array data and label from h5_filename 100 | def load_h5_data_label_seg(h5_filename): 101 | f = h5py.File(h5_filename) 102 | data = f['data'][:] 103 | label = f['label'][:] 104 | seg = f['pid'][:] 105 | return (data, label, seg) 106 | 107 | # Read numpy array data and label from h5_filename 108 | def load_h5(h5_filename): 109 | f = h5py.File(h5_filename) 110 | data = f['data'][:] 111 | label = f['label'][:] 112 | return (data, label) 113 | 114 | # ---------------------------------------------------------------- 115 | # Following are the helper functions to load save/load PLY files 116 | # ---------------------------------------------------------------- 117 | 118 | # Load PLY file 119 | def load_ply_data(filename, point_num): 120 | plydata = PlyData.read(filename) 121 | pc = plydata['vertex'].data[:point_num] 122 | pc_array = np.array([[x, y, z] for x,y,z in pc]) 123 | return pc_array 124 | 125 | # Load PLY file 126 | def load_ply_normal(filename, point_num): 127 | plydata = PlyData.read(filename) 128 | pc = plydata['normal'].data[:point_num] 129 | pc_array = np.array([[x, y, z] for x,y,z in pc]) 130 | return pc_array 131 | 132 | # Make up rows for Nxk array 133 | # Input Pad is 'edge' or 'constant' 134 | def pad_arr_rows(arr, row, pad='edge'): 135 | assert(len(arr.shape) == 2) 136 | assert(arr.shape[0] <= row) 137 | assert(pad == 'edge' or pad == 'constant') 138 | if arr.shape[0] == row: 139 | return arr 140 | if pad == 'edge': 141 | return np.lib.pad(arr, ((0, row-arr.shape[0]), (0, 0)), 'edge') 142 | if pad == 'constant': 143 | return np.lib.pad(arr, ((0, row-arr.shape[0]), (0, 0)), 'constant', (0, 0)) 144 | 145 | 146 | -------------------------------------------------------------------------------- /prepare_data/gen_indoor3d_h5.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | import sys 4 | import json 5 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 6 | ROOT_DIR = os.path.dirname(BASE_DIR) 7 | sys.path.append(BASE_DIR) 8 | import data_prep_util 9 | import indoor3d_util 10 | 11 | # Constants 12 | data_dir = os.path.join(ROOT_DIR, 'data') 13 | indoor3d_data_dir = os.path.join(data_dir, 'stanford_indoor3d') 14 | NUM_POINT = 4096 15 | H5_BATCH_SIZE = 1000 16 | data_dim = [NUM_POINT, 9] 17 | label_dim = [NUM_POINT] 18 | data_dtype = 'float32' 19 | label_dtype = 'uint8' 20 | 21 | # Set paths 22 | filelist = os.path.join(BASE_DIR, 'meta/all_data_label.txt') 23 | data_label_files = [os.path.join(indoor3d_data_dir, line.rstrip()) for line in open(filelist)] 24 | output_dir = os.path.join(data_dir, 'indoor3d_sem_seg_hdf5_data_test') 25 | if not os.path.exists(output_dir): 26 | os.mkdir(output_dir) 27 | output_filename_prefix = os.path.join(output_dir, 'ply_data_all') 28 | output_room_filelist = os.path.join(output_dir, 'room_filelist.txt') 29 | output_all_file = os.path.join(output_dir, 'all_files.txt') 30 | fout_room = open(output_room_filelist, 'w') 31 | all_file = open(output_all_file, 'w') 32 | 33 | # -------------------------------------- 34 | # ----- BATCH WRITE TO HDF5 ----- 35 | # -------------------------------------- 36 | batch_data_dim = [H5_BATCH_SIZE] + data_dim 37 | batch_label_dim = [H5_BATCH_SIZE] + label_dim 38 | h5_batch_data = np.zeros(batch_data_dim, dtype = np.float32) 39 | h5_batch_label = np.zeros(batch_label_dim, dtype = np.uint8) 40 | buffer_size = 0 # state: record how many samples are currently in buffer 41 | h5_index = 0 # state: the next h5 file to save 42 | 43 | def insert_batch(data, label, last_batch=False): 44 | global h5_batch_data, h5_batch_label 45 | global buffer_size, h5_index 46 | data_size = data.shape[0] 47 | # If there is enough space, just insert 48 | if buffer_size + data_size <= h5_batch_data.shape[0]: 49 | h5_batch_data[buffer_size:buffer_size+data_size, ...] = data 50 | h5_batch_label[buffer_size:buffer_size+data_size] = label 51 | buffer_size += data_size 52 | else: # not enough space 53 | capacity = h5_batch_data.shape[0] - buffer_size 54 | assert(capacity>=0) 55 | if capacity > 0: 56 | h5_batch_data[buffer_size:buffer_size+capacity, ...] = data[0:capacity, ...] 57 | h5_batch_label[buffer_size:buffer_size+capacity, ...] = label[0:capacity, ...] 58 | # Save batch data and label to h5 file, reset buffer_size 59 | h5_filename = output_filename_prefix + '_' + str(h5_index) + '.h5' 60 | data_prep_util.save_h5(h5_filename, h5_batch_data, h5_batch_label, data_dtype, label_dtype) 61 | print('Stored {0} with size {1}'.format(h5_filename, h5_batch_data.shape[0])) 62 | h5_index += 1 63 | buffer_size = 0 64 | # recursive call 65 | insert_batch(data[capacity:, ...], label[capacity:, ...], last_batch) 66 | if last_batch and buffer_size > 0: 67 | h5_filename = output_filename_prefix + '_' + str(h5_index) + '.h5' 68 | data_prep_util.save_h5(h5_filename, h5_batch_data[0:buffer_size, ...], h5_batch_label[0:buffer_size, ...], data_dtype, label_dtype) 69 | print('Stored {0} with size {1}'.format(h5_filename, buffer_size)) 70 | h5_index += 1 71 | buffer_size = 0 72 | return 73 | 74 | 75 | sample_cnt = 0 76 | for i, data_label_filename in enumerate(data_label_files): 77 | print(data_label_filename) 78 | data, label = indoor3d_util.room2blocks_wrapper_normalized(data_label_filename, NUM_POINT, block_size=1.0, stride=1, 79 | random_sample=False, sample_num=None) 80 | print('{0}, {1}'.format(data.shape, label.shape)) 81 | for _ in range(data.shape[0]): 82 | fout_room.write(os.path.basename(data_label_filename)[0:-4]+'\n') 83 | 84 | sample_cnt += data.shape[0] 85 | insert_batch(data, label, i == len(data_label_files)-1) 86 | 87 | fout_room.close() 88 | print("Total samples: {0}".format(sample_cnt)) 89 | 90 | for i in range(h5_index): 91 | all_file.write(os.path.join('indoor3d_sem_seg_hdf5_data_test', 'ply_data_all_') + str(i) +'.h5\n') 92 | all_file.close() 93 | -------------------------------------------------------------------------------- /prepare_data/meta/all_data_label.txt: -------------------------------------------------------------------------------- 1 | Area_1_conferenceRoom_1.npy 2 | Area_1_conferenceRoom_2.npy 3 | Area_1_copyRoom_1.npy 4 | Area_1_hallway_1.npy 5 | Area_1_hallway_2.npy 6 | Area_1_hallway_3.npy 7 | Area_1_hallway_4.npy 8 | Area_1_hallway_5.npy 9 | Area_1_hallway_6.npy 10 | Area_1_hallway_7.npy 11 | Area_1_hallway_8.npy 12 | Area_1_office_10.npy 13 | Area_1_office_11.npy 14 | Area_1_office_12.npy 15 | Area_1_office_13.npy 16 | Area_1_office_14.npy 17 | Area_1_office_15.npy 18 | Area_1_office_16.npy 19 | Area_1_office_17.npy 20 | Area_1_office_18.npy 21 | Area_1_office_19.npy 22 | Area_1_office_1.npy 23 | Area_1_office_20.npy 24 | Area_1_office_21.npy 25 | Area_1_office_22.npy 26 | Area_1_office_23.npy 27 | Area_1_office_24.npy 28 | Area_1_office_25.npy 29 | Area_1_office_26.npy 30 | Area_1_office_27.npy 31 | Area_1_office_28.npy 32 | Area_1_office_29.npy 33 | Area_1_office_2.npy 34 | Area_1_office_30.npy 35 | Area_1_office_31.npy 36 | Area_1_office_3.npy 37 | Area_1_office_4.npy 38 | Area_1_office_5.npy 39 | Area_1_office_6.npy 40 | Area_1_office_7.npy 41 | Area_1_office_8.npy 42 | Area_1_office_9.npy 43 | Area_1_pantry_1.npy 44 | Area_1_WC_1.npy 45 | Area_2_auditorium_1.npy 46 | Area_2_auditorium_2.npy 47 | Area_2_conferenceRoom_1.npy 48 | Area_2_hallway_10.npy 49 | Area_2_hallway_11.npy 50 | Area_2_hallway_12.npy 51 | Area_2_hallway_1.npy 52 | Area_2_hallway_2.npy 53 | Area_2_hallway_3.npy 54 | Area_2_hallway_4.npy 55 | Area_2_hallway_5.npy 56 | Area_2_hallway_6.npy 57 | Area_2_hallway_7.npy 58 | Area_2_hallway_8.npy 59 | Area_2_hallway_9.npy 60 | Area_2_office_10.npy 61 | Area_2_office_11.npy 62 | Area_2_office_12.npy 63 | Area_2_office_13.npy 64 | Area_2_office_14.npy 65 | Area_2_office_1.npy 66 | Area_2_office_2.npy 67 | Area_2_office_3.npy 68 | Area_2_office_4.npy 69 | Area_2_office_5.npy 70 | Area_2_office_6.npy 71 | Area_2_office_7.npy 72 | Area_2_office_8.npy 73 | Area_2_office_9.npy 74 | Area_2_storage_1.npy 75 | Area_2_storage_2.npy 76 | Area_2_storage_3.npy 77 | Area_2_storage_4.npy 78 | Area_2_storage_5.npy 79 | Area_2_storage_6.npy 80 | Area_2_storage_7.npy 81 | Area_2_storage_8.npy 82 | Area_2_storage_9.npy 83 | Area_2_WC_1.npy 84 | Area_2_WC_2.npy 85 | Area_3_conferenceRoom_1.npy 86 | Area_3_hallway_1.npy 87 | Area_3_hallway_2.npy 88 | Area_3_hallway_3.npy 89 | Area_3_hallway_4.npy 90 | Area_3_hallway_5.npy 91 | Area_3_hallway_6.npy 92 | Area_3_lounge_1.npy 93 | Area_3_lounge_2.npy 94 | Area_3_office_10.npy 95 | Area_3_office_1.npy 96 | Area_3_office_2.npy 97 | Area_3_office_3.npy 98 | Area_3_office_4.npy 99 | Area_3_office_5.npy 100 | Area_3_office_6.npy 101 | Area_3_office_7.npy 102 | Area_3_office_8.npy 103 | Area_3_office_9.npy 104 | Area_3_storage_1.npy 105 | Area_3_storage_2.npy 106 | Area_3_WC_1.npy 107 | Area_3_WC_2.npy 108 | Area_4_conferenceRoom_1.npy 109 | Area_4_conferenceRoom_2.npy 110 | Area_4_conferenceRoom_3.npy 111 | Area_4_hallway_10.npy 112 | Area_4_hallway_11.npy 113 | Area_4_hallway_12.npy 114 | Area_4_hallway_13.npy 115 | Area_4_hallway_14.npy 116 | Area_4_hallway_1.npy 117 | Area_4_hallway_2.npy 118 | Area_4_hallway_3.npy 119 | Area_4_hallway_4.npy 120 | Area_4_hallway_5.npy 121 | Area_4_hallway_6.npy 122 | Area_4_hallway_7.npy 123 | Area_4_hallway_8.npy 124 | Area_4_hallway_9.npy 125 | Area_4_lobby_1.npy 126 | Area_4_lobby_2.npy 127 | Area_4_office_10.npy 128 | Area_4_office_11.npy 129 | Area_4_office_12.npy 130 | Area_4_office_13.npy 131 | Area_4_office_14.npy 132 | Area_4_office_15.npy 133 | Area_4_office_16.npy 134 | Area_4_office_17.npy 135 | Area_4_office_18.npy 136 | Area_4_office_19.npy 137 | Area_4_office_1.npy 138 | Area_4_office_20.npy 139 | Area_4_office_21.npy 140 | Area_4_office_22.npy 141 | Area_4_office_2.npy 142 | Area_4_office_3.npy 143 | Area_4_office_4.npy 144 | Area_4_office_5.npy 145 | Area_4_office_6.npy 146 | Area_4_office_7.npy 147 | Area_4_office_8.npy 148 | Area_4_office_9.npy 149 | Area_4_storage_1.npy 150 | Area_4_storage_2.npy 151 | Area_4_storage_3.npy 152 | Area_4_storage_4.npy 153 | Area_4_WC_1.npy 154 | Area_4_WC_2.npy 155 | Area_4_WC_3.npy 156 | Area_4_WC_4.npy 157 | Area_5_conferenceRoom_1.npy 158 | Area_5_conferenceRoom_2.npy 159 | Area_5_conferenceRoom_3.npy 160 | Area_5_hallway_10.npy 161 | Area_5_hallway_11.npy 162 | Area_5_hallway_12.npy 163 | Area_5_hallway_13.npy 164 | Area_5_hallway_14.npy 165 | Area_5_hallway_15.npy 166 | Area_5_hallway_1.npy 167 | Area_5_hallway_2.npy 168 | Area_5_hallway_3.npy 169 | Area_5_hallway_4.npy 170 | Area_5_hallway_5.npy 171 | Area_5_hallway_6.npy 172 | Area_5_hallway_7.npy 173 | Area_5_hallway_8.npy 174 | Area_5_hallway_9.npy 175 | Area_5_lobby_1.npy 176 | Area_5_office_10.npy 177 | Area_5_office_11.npy 178 | Area_5_office_12.npy 179 | Area_5_office_13.npy 180 | Area_5_office_14.npy 181 | Area_5_office_15.npy 182 | Area_5_office_16.npy 183 | Area_5_office_17.npy 184 | Area_5_office_18.npy 185 | Area_5_office_19.npy 186 | Area_5_office_1.npy 187 | Area_5_office_20.npy 188 | Area_5_office_21.npy 189 | Area_5_office_22.npy 190 | Area_5_office_23.npy 191 | Area_5_office_24.npy 192 | Area_5_office_25.npy 193 | Area_5_office_26.npy 194 | Area_5_office_27.npy 195 | Area_5_office_28.npy 196 | Area_5_office_29.npy 197 | Area_5_office_2.npy 198 | Area_5_office_30.npy 199 | Area_5_office_31.npy 200 | Area_5_office_32.npy 201 | Area_5_office_33.npy 202 | Area_5_office_34.npy 203 | Area_5_office_35.npy 204 | Area_5_office_36.npy 205 | Area_5_office_37.npy 206 | Area_5_office_38.npy 207 | Area_5_office_39.npy 208 | Area_5_office_3.npy 209 | Area_5_office_40.npy 210 | Area_5_office_41.npy 211 | Area_5_office_42.npy 212 | Area_5_office_4.npy 213 | Area_5_office_5.npy 214 | Area_5_office_6.npy 215 | Area_5_office_7.npy 216 | Area_5_office_8.npy 217 | Area_5_office_9.npy 218 | Area_5_pantry_1.npy 219 | Area_5_storage_1.npy 220 | Area_5_storage_2.npy 221 | Area_5_storage_3.npy 222 | Area_5_storage_4.npy 223 | Area_5_WC_1.npy 224 | Area_5_WC_2.npy 225 | Area_6_conferenceRoom_1.npy 226 | Area_6_copyRoom_1.npy 227 | Area_6_hallway_1.npy 228 | Area_6_hallway_2.npy 229 | Area_6_hallway_3.npy 230 | Area_6_hallway_4.npy 231 | Area_6_hallway_5.npy 232 | Area_6_hallway_6.npy 233 | Area_6_lounge_1.npy 234 | Area_6_office_10.npy 235 | Area_6_office_11.npy 236 | Area_6_office_12.npy 237 | Area_6_office_13.npy 238 | Area_6_office_14.npy 239 | Area_6_office_15.npy 240 | Area_6_office_16.npy 241 | Area_6_office_17.npy 242 | Area_6_office_18.npy 243 | Area_6_office_19.npy 244 | Area_6_office_1.npy 245 | Area_6_office_20.npy 246 | Area_6_office_21.npy 247 | Area_6_office_22.npy 248 | Area_6_office_23.npy 249 | Area_6_office_24.npy 250 | Area_6_office_25.npy 251 | Area_6_office_26.npy 252 | Area_6_office_27.npy 253 | Area_6_office_28.npy 254 | Area_6_office_29.npy 255 | Area_6_office_2.npy 256 | Area_6_office_30.npy 257 | Area_6_office_31.npy 258 | Area_6_office_32.npy 259 | Area_6_office_33.npy 260 | Area_6_office_34.npy 261 | Area_6_office_35.npy 262 | Area_6_office_36.npy 263 | Area_6_office_37.npy 264 | Area_6_office_3.npy 265 | Area_6_office_4.npy 266 | Area_6_office_5.npy 267 | Area_6_office_6.npy 268 | Area_6_office_7.npy 269 | Area_6_office_8.npy 270 | Area_6_office_9.npy 271 | Area_6_openspace_1.npy 272 | Area_6_pantry_1.npy 273 | -------------------------------------------------------------------------------- /prepare_data/meta/anno_paths.txt: -------------------------------------------------------------------------------- 1 | Area_1/conferenceRoom_1/Annotations 2 | Area_1/conferenceRoom_2/Annotations 3 | Area_1/copyRoom_1/Annotations 4 | Area_1/hallway_1/Annotations 5 | Area_1/hallway_2/Annotations 6 | Area_1/hallway_3/Annotations 7 | Area_1/hallway_4/Annotations 8 | Area_1/hallway_5/Annotations 9 | Area_1/hallway_6/Annotations 10 | Area_1/hallway_7/Annotations 11 | Area_1/hallway_8/Annotations 12 | Area_1/office_10/Annotations 13 | Area_1/office_11/Annotations 14 | Area_1/office_12/Annotations 15 | Area_1/office_13/Annotations 16 | Area_1/office_14/Annotations 17 | Area_1/office_15/Annotations 18 | Area_1/office_16/Annotations 19 | Area_1/office_17/Annotations 20 | Area_1/office_18/Annotations 21 | Area_1/office_19/Annotations 22 | Area_1/office_1/Annotations 23 | Area_1/office_20/Annotations 24 | Area_1/office_21/Annotations 25 | Area_1/office_22/Annotations 26 | Area_1/office_23/Annotations 27 | Area_1/office_24/Annotations 28 | Area_1/office_25/Annotations 29 | Area_1/office_26/Annotations 30 | Area_1/office_27/Annotations 31 | Area_1/office_28/Annotations 32 | Area_1/office_29/Annotations 33 | Area_1/office_2/Annotations 34 | Area_1/office_30/Annotations 35 | Area_1/office_31/Annotations 36 | Area_1/office_3/Annotations 37 | Area_1/office_4/Annotations 38 | Area_1/office_5/Annotations 39 | Area_1/office_6/Annotations 40 | Area_1/office_7/Annotations 41 | Area_1/office_8/Annotations 42 | Area_1/office_9/Annotations 43 | Area_1/pantry_1/Annotations 44 | Area_1/WC_1/Annotations 45 | Area_2/auditorium_1/Annotations 46 | Area_2/auditorium_2/Annotations 47 | Area_2/conferenceRoom_1/Annotations 48 | Area_2/hallway_10/Annotations 49 | Area_2/hallway_11/Annotations 50 | Area_2/hallway_12/Annotations 51 | Area_2/hallway_1/Annotations 52 | Area_2/hallway_2/Annotations 53 | Area_2/hallway_3/Annotations 54 | Area_2/hallway_4/Annotations 55 | Area_2/hallway_5/Annotations 56 | Area_2/hallway_6/Annotations 57 | Area_2/hallway_7/Annotations 58 | Area_2/hallway_8/Annotations 59 | Area_2/hallway_9/Annotations 60 | Area_2/office_10/Annotations 61 | Area_2/office_11/Annotations 62 | Area_2/office_12/Annotations 63 | Area_2/office_13/Annotations 64 | Area_2/office_14/Annotations 65 | Area_2/office_1/Annotations 66 | Area_2/office_2/Annotations 67 | Area_2/office_3/Annotations 68 | Area_2/office_4/Annotations 69 | Area_2/office_5/Annotations 70 | Area_2/office_6/Annotations 71 | Area_2/office_7/Annotations 72 | Area_2/office_8/Annotations 73 | Area_2/office_9/Annotations 74 | Area_2/storage_1/Annotations 75 | Area_2/storage_2/Annotations 76 | Area_2/storage_3/Annotations 77 | Area_2/storage_4/Annotations 78 | Area_2/storage_5/Annotations 79 | Area_2/storage_6/Annotations 80 | Area_2/storage_7/Annotations 81 | Area_2/storage_8/Annotations 82 | Area_2/storage_9/Annotations 83 | Area_2/WC_1/Annotations 84 | Area_2/WC_2/Annotations 85 | Area_3/conferenceRoom_1/Annotations 86 | Area_3/hallway_1/Annotations 87 | Area_3/hallway_2/Annotations 88 | Area_3/hallway_3/Annotations 89 | Area_3/hallway_4/Annotations 90 | Area_3/hallway_5/Annotations 91 | Area_3/hallway_6/Annotations 92 | Area_3/lounge_1/Annotations 93 | Area_3/lounge_2/Annotations 94 | Area_3/office_10/Annotations 95 | Area_3/office_1/Annotations 96 | Area_3/office_2/Annotations 97 | Area_3/office_3/Annotations 98 | Area_3/office_4/Annotations 99 | Area_3/office_5/Annotations 100 | Area_3/office_6/Annotations 101 | Area_3/office_7/Annotations 102 | Area_3/office_8/Annotations 103 | Area_3/office_9/Annotations 104 | Area_3/storage_1/Annotations 105 | Area_3/storage_2/Annotations 106 | Area_3/WC_1/Annotations 107 | Area_3/WC_2/Annotations 108 | Area_4/conferenceRoom_1/Annotations 109 | Area_4/conferenceRoom_2/Annotations 110 | Area_4/conferenceRoom_3/Annotations 111 | Area_4/hallway_10/Annotations 112 | Area_4/hallway_11/Annotations 113 | Area_4/hallway_12/Annotations 114 | Area_4/hallway_13/Annotations 115 | Area_4/hallway_14/Annotations 116 | Area_4/hallway_1/Annotations 117 | Area_4/hallway_2/Annotations 118 | Area_4/hallway_3/Annotations 119 | Area_4/hallway_4/Annotations 120 | Area_4/hallway_5/Annotations 121 | Area_4/hallway_6/Annotations 122 | Area_4/hallway_7/Annotations 123 | Area_4/hallway_8/Annotations 124 | Area_4/hallway_9/Annotations 125 | Area_4/lobby_1/Annotations 126 | Area_4/lobby_2/Annotations 127 | Area_4/office_10/Annotations 128 | Area_4/office_11/Annotations 129 | Area_4/office_12/Annotations 130 | Area_4/office_13/Annotations 131 | Area_4/office_14/Annotations 132 | Area_4/office_15/Annotations 133 | Area_4/office_16/Annotations 134 | Area_4/office_17/Annotations 135 | Area_4/office_18/Annotations 136 | Area_4/office_19/Annotations 137 | Area_4/office_1/Annotations 138 | Area_4/office_20/Annotations 139 | Area_4/office_21/Annotations 140 | Area_4/office_22/Annotations 141 | Area_4/office_2/Annotations 142 | Area_4/office_3/Annotations 143 | Area_4/office_4/Annotations 144 | Area_4/office_5/Annotations 145 | Area_4/office_6/Annotations 146 | Area_4/office_7/Annotations 147 | Area_4/office_8/Annotations 148 | Area_4/office_9/Annotations 149 | Area_4/storage_1/Annotations 150 | Area_4/storage_2/Annotations 151 | Area_4/storage_3/Annotations 152 | Area_4/storage_4/Annotations 153 | Area_4/WC_1/Annotations 154 | Area_4/WC_2/Annotations 155 | Area_4/WC_3/Annotations 156 | Area_4/WC_4/Annotations 157 | Area_5/conferenceRoom_1/Annotations 158 | Area_5/conferenceRoom_2/Annotations 159 | Area_5/conferenceRoom_3/Annotations 160 | Area_5/hallway_10/Annotations 161 | Area_5/hallway_11/Annotations 162 | Area_5/hallway_12/Annotations 163 | Area_5/hallway_13/Annotations 164 | Area_5/hallway_14/Annotations 165 | Area_5/hallway_15/Annotations 166 | Area_5/hallway_1/Annotations 167 | Area_5/hallway_2/Annotations 168 | Area_5/hallway_3/Annotations 169 | Area_5/hallway_4/Annotations 170 | Area_5/hallway_5/Annotations 171 | Area_5/hallway_6/Annotations 172 | Area_5/hallway_7/Annotations 173 | Area_5/hallway_8/Annotations 174 | Area_5/hallway_9/Annotations 175 | Area_5/lobby_1/Annotations 176 | Area_5/office_10/Annotations 177 | Area_5/office_11/Annotations 178 | Area_5/office_12/Annotations 179 | Area_5/office_13/Annotations 180 | Area_5/office_14/Annotations 181 | Area_5/office_15/Annotations 182 | Area_5/office_16/Annotations 183 | Area_5/office_17/Annotations 184 | Area_5/office_18/Annotations 185 | Area_5/office_19/Annotations 186 | Area_5/office_1/Annotations 187 | Area_5/office_20/Annotations 188 | Area_5/office_21/Annotations 189 | Area_5/office_22/Annotations 190 | Area_5/office_23/Annotations 191 | Area_5/office_24/Annotations 192 | Area_5/office_25/Annotations 193 | Area_5/office_26/Annotations 194 | Area_5/office_27/Annotations 195 | Area_5/office_28/Annotations 196 | Area_5/office_29/Annotations 197 | Area_5/office_2/Annotations 198 | Area_5/office_30/Annotations 199 | Area_5/office_31/Annotations 200 | Area_5/office_32/Annotations 201 | Area_5/office_33/Annotations 202 | Area_5/office_34/Annotations 203 | Area_5/office_35/Annotations 204 | Area_5/office_36/Annotations 205 | Area_5/office_37/Annotations 206 | Area_5/office_38/Annotations 207 | Area_5/office_39/Annotations 208 | Area_5/office_3/Annotations 209 | Area_5/office_40/Annotations 210 | Area_5/office_41/Annotations 211 | Area_5/office_42/Annotations 212 | Area_5/office_4/Annotations 213 | Area_5/office_5/Annotations 214 | Area_5/office_6/Annotations 215 | Area_5/office_7/Annotations 216 | Area_5/office_8/Annotations 217 | Area_5/office_9/Annotations 218 | Area_5/pantry_1/Annotations 219 | Area_5/storage_1/Annotations 220 | Area_5/storage_2/Annotations 221 | Area_5/storage_3/Annotations 222 | Area_5/storage_4/Annotations 223 | Area_5/WC_1/Annotations 224 | Area_5/WC_2/Annotations 225 | Area_6/conferenceRoom_1/Annotations 226 | Area_6/copyRoom_1/Annotations 227 | Area_6/hallway_1/Annotations 228 | Area_6/hallway_2/Annotations 229 | Area_6/hallway_3/Annotations 230 | Area_6/hallway_4/Annotations 231 | Area_6/hallway_5/Annotations 232 | Area_6/hallway_6/Annotations 233 | Area_6/lounge_1/Annotations 234 | Area_6/office_10/Annotations 235 | Area_6/office_11/Annotations 236 | Area_6/office_12/Annotations 237 | Area_6/office_13/Annotations 238 | Area_6/office_14/Annotations 239 | Area_6/office_15/Annotations 240 | Area_6/office_16/Annotations 241 | Area_6/office_17/Annotations 242 | Area_6/office_18/Annotations 243 | Area_6/office_19/Annotations 244 | Area_6/office_1/Annotations 245 | Area_6/office_20/Annotations 246 | Area_6/office_21/Annotations 247 | Area_6/office_22/Annotations 248 | Area_6/office_23/Annotations 249 | Area_6/office_24/Annotations 250 | Area_6/office_25/Annotations 251 | Area_6/office_26/Annotations 252 | Area_6/office_27/Annotations 253 | Area_6/office_28/Annotations 254 | Area_6/office_29/Annotations 255 | Area_6/office_2/Annotations 256 | Area_6/office_30/Annotations 257 | Area_6/office_31/Annotations 258 | Area_6/office_32/Annotations 259 | Area_6/office_33/Annotations 260 | Area_6/office_34/Annotations 261 | Area_6/office_35/Annotations 262 | Area_6/office_36/Annotations 263 | Area_6/office_37/Annotations 264 | Area_6/office_3/Annotations 265 | Area_6/office_4/Annotations 266 | Area_6/office_5/Annotations 267 | Area_6/office_6/Annotations 268 | Area_6/office_7/Annotations 269 | Area_6/office_8/Annotations 270 | Area_6/office_9/Annotations 271 | Area_6/openspace_1/Annotations 272 | Area_6/pantry_1/Annotations 273 | -------------------------------------------------------------------------------- /prepare_data/meta/class_names.txt: -------------------------------------------------------------------------------- 1 | ceiling 2 | floor 3 | wall 4 | beam 5 | column 6 | window 7 | door 8 | table 9 | chair 10 | sofa 11 | bookcase 12 | board 13 | clutter 14 | -------------------------------------------------------------------------------- /prepare_data/meta/partseg_colors.txt: -------------------------------------------------------------------------------- 1 | [ 2 | {"label": "airplane_body","id": 0,"color": [152,223,138]}, 3 | {"label": "airplane_wings","id": 1,"color": [174,199,232]}, 4 | {"label": "airplane_tail","id": 2,"color": [255,105,180]}, 5 | {"label": "airplane_engine","id": 3,"color": [31,119,180]}, 6 | {"label": "bag_handle","id": 4,"color": [255,187,120]}, 7 | {"label": "bag_body","id": 5,"color": [188,189,34]}, 8 | {"label": "cap_peak","id": 6,"color": [140,86,75]}, 9 | {"label": "cap_panel","id": 7,"color": [255,152,150]}, 10 | {"label": "car_roof","id": 8,"color": [214,39,40]}, 11 | {"label": "car_hood","id": 9,"color": [197,176,213]}, 12 | {"label": "car_wheels","id": 10,"color": [148,103,189]}, 13 | {"label": "car_body","id": 11,"color": [196,156,148]}, 14 | {"label": "chair_back","id": 12,"color": [23,190,207]}, 15 | {"label": "chair_seat","id": 13,"color": [186,85,211]}, 16 | {"label": "chair_leg","id": 14,"color": [247,182,210]}, 17 | {"label": "chair_arm","id": 15,"color": [66,188,102]}, 18 | {"label": "earphone_earphone","id": 16,"color": [219,219,141]}, 19 | {"label": "earphone_headband","id": 17,"color": [140,57,197]}, 20 | {"label": "earphone_microphone","id": 18,"color": [202,185,52]}, 21 | {"label": "guitar_head","id": 19,"color": [213,92,176]}, 22 | {"label": "guitar_neck","id": 20,"color": [200,54,131]}, 23 | {"label": "guitar_body","id": 21,"color": [92,193,61]}, 24 | {"label": "knife_blade","id": 22,"color": [78,71,183]}, 25 | {"label": "knife_handle","id": 23,"color": [172,114,82]}, 26 | {"label": "lamp_base","id": 24,"color": [255,127,14]}, 27 | {"label": "lamp_shape","id": 25,"color": [91,163,138]}, 28 | {"label": "lamp_ceiling","id": 26,"color": [153,98,156]}, 29 | {"label": "lamp_tube","id": 27,"color": [140,153,101]}, 30 | {"label": "laptop_keyboard","id": 28,"color": [158,218,229]}, 31 | {"label": "laptop_screen","id": 29,"color": [100,125,154]}, 32 | {"label": "motorbike_gastank","id": 30,"color": [178,127,135]}, 33 | {"label": "motorbike_seat","id": 31,"color": [120,185,128]}, 34 | {"label": "motorbike_wheel","id": 32,"color": [146,111,194]}, 35 | {"label": "motorbike_handle","id": 33,"color": [44,160,44]}, 36 | {"label": "motorbike_light","id": 34,"color": [112,128,144]}, 37 | {"label": "motorbike_body","id": 35,"color": [96,207,209]}, 38 | {"label": "mug_handle","id": 36,"color": [227,119,194]}, 39 | {"label": "mug_body","id": 37,"color": [51,176,203]}, 40 | {"label": "pistol_barrel","id": 38,"color": [94,106,211]}, 41 | {"label": "pistol_handle","id": 39,"color": [82,84,163]}, 42 | {"label": "pistol_trigger","id": 40,"color": [100,85,144]}, 43 | {"label": "rocket_body","id": 41,"color": [255,127,80]}, 44 | {"label": "rocket_fin","id": 42,"color": [0,100,0]}, 45 | {"label": "rocket_nose","id": 43,"color": [173,255,47]}, 46 | {"label": "skateboard_wheel","id": 44,"color": [64,224,208]}, 47 | {"label": "skateboard_deck","id": 45,"color": [0,255,255]}, 48 | {"label": "skateboard_bearing","id": 46,"color": [25,25,112]}, 49 | {"label": "table_top","id": 47,"color": [178,76,76]}, 50 | {"label": "table_leg","id": 48,"color": [255,0,255]}, 51 | {"label": "table_drawer","id": 49,"color": [152,223,138]} 52 | ] -------------------------------------------------------------------------------- /prepare_data/meta/semseg_colors.txt: -------------------------------------------------------------------------------- 1 | [ 2 | {"label": "ceiling","id": 1,"color": [152,223,138]}, 3 | {"label": "floor","id": 2,"color": [174,199,232]}, 4 | {"label": "wall","id": 3,"color": [255,127,14]}, 5 | {"label": "beam","id": 4,"color": [91,163,138]}, 6 | {"label": "column","id": 5,"color": [255,187,120]}, 7 | {"label": "window","id": 6,"color": [188,189,34]}, 8 | {"label": "door","id": 7,"color": [140,86,75]}, 9 | {"label": "table","id": 8,"color": [255,152,150]}, 10 | {"label": "chair","id": 9,"color": [214,39,40]}, 11 | {"label": "sofa","id": 10,"color": [197,176,213]}, 12 | {"label": "bookcase","id": 11,"color": [196,156,148]}, 13 | {"label": "board","id": 12,"color": [23,190,207]}, 14 | {"label": "clutter","id": 13,"color": [112,128,144]} 15 | ] -------------------------------------------------------------------------------- /prepare_data/scannetv2_seg_dataset_rgb21c_pointid.py: -------------------------------------------------------------------------------- 1 | """ 2 | ScanNet v2 data preprocessing. 3 | Extract point clouds data from .ply files to genrate .pickle files for training and testing. 4 | Author: Wenxuan Wu 5 | Date: July 2018 6 | """ 7 | 8 | import os 9 | import sys 10 | import numpy as np 11 | import util 12 | import h5py 13 | import pickle 14 | from plyfile import PlyData, PlyElement 15 | 16 | 17 | def remove_unano(scene_data, scene_label, scene_data_id): 18 | keep_idx = np.where((scene_label > 0) & ( 19 | scene_label < 41)) # 0: unanotated 20 | scene_data_clean = scene_data[keep_idx] 21 | scene_label_clean = scene_label[keep_idx] 22 | scene_data_id_clean = scene_data_id[keep_idx] 23 | return scene_data_clean, scene_label_clean, scene_data_id_clean 24 | 25 | 26 | test_class = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 27 | 10, 11, 12, 14, 16, 24, 28, 33, 34, 36, 39] 28 | 29 | 30 | def gen_label_map(): 31 | label_map = np.zeros(41) 32 | for i in range(41): 33 | if i in test_class: 34 | label_map[i] = test_class.index(i) 35 | else: 36 | label_map[i] = 0 37 | print(label_map) 38 | return label_map 39 | 40 | 41 | def gen_pickle(split="val", keep_unanno=False, root="DataSet/Scannet_v2"): 42 | if split == 'test': 43 | root_new = root + "/scans_test" 44 | else: 45 | root_new = root + "/scans" 46 | file_list = "scannetv2_%s.txt" % (split) 47 | with open(file_list) as fl: 48 | scene_id = fl.read().splitlines() 49 | 50 | scene_data = [] 51 | scene_data_labels = [] 52 | scene_data_id = [] 53 | scene_data_num = [] 54 | label_map = gen_label_map() 55 | for i in range(len(scene_id)): # len(scene_id) 56 | print('process...', i) 57 | scene_namergb = os.path.join( 58 | root_new, scene_id[i], scene_id[i]+'_vh_clean_2.ply') 59 | scene_xyzlabelrgb = PlyData.read(scene_namergb) 60 | scene_vertex_rgb = scene_xyzlabelrgb['vertex'] 61 | scene_data_tmp = np.stack((scene_vertex_rgb['x'], scene_vertex_rgb['y'], 62 | scene_vertex_rgb['z'], scene_vertex_rgb['red'], 63 | scene_vertex_rgb['green'], scene_vertex_rgb['blue']), axis=-1).astype(np.float32) 64 | scene_points_num = scene_data_tmp.shape[0] 65 | scene_point_id = np.array([c for c in range(scene_points_num)]) 66 | if not keep_unanno: 67 | scene_name = os.path.join( 68 | root_new, scene_id[i], scene_id[i]+'_vh_clean_2.labels.ply') 69 | scene_xyzlabel = PlyData.read(scene_name) 70 | scene_vertex = scene_xyzlabel['vertex'] 71 | scene_data_label_tmp = scene_vertex['label'] 72 | scene_data_tmp, scene_data_label_tmp, scene_point_id_tmp = remove_unano( 73 | scene_data_tmp, scene_data_label_tmp, scene_point_id) 74 | scene_data_label_tmp = label_map[scene_data_label_tmp] 75 | elif split != 'test': 76 | scene_name = os.path.join( 77 | root_new, scene_id[i], scene_id[i]+'_vh_clean_2.labels.ply') 78 | scene_xyzlabel = PlyData.read(scene_name) 79 | scene_vertex = scene_xyzlabel['vertex'] 80 | scene_point_id_tmp = scene_point_id 81 | scene_data_label_tmp = scene_vertex['label'] 82 | scene_data_label_tmp[np.where(scene_data_label_tmp > 40)] = 0 83 | scene_data_label_tmp = label_map[scene_data_label_tmp] 84 | else: 85 | scene_data_label_tmp = np.zeros( 86 | (scene_data_tmp.shape[0])).astype(np.int32) 87 | scene_point_id_tmp = scene_point_id 88 | scene_data.append(scene_data_tmp) 89 | scene_data_labels.append(scene_data_label_tmp) 90 | scene_data_id.append(scene_point_id_tmp) 91 | scene_data_num.append(scene_points_num) 92 | 93 | if not keep_unanno: 94 | out_path = os.path.join(root, "scannet_%s_rgb21c_pointid.pickle" % (split)) 95 | else: 96 | out_path = os.path.join(root, "scannet_%s_rgb21c_pointid_keep_unanno.pickle" % (split)) 97 | pickle_out = open(out_path, "wb") 98 | pickle.dump(scene_data, pickle_out, protocol=0) 99 | pickle.dump(scene_data_labels, pickle_out, protocol=0) 100 | pickle.dump(scene_data_id, pickle_out, protocol=0) 101 | pickle.dump(scene_data_num, pickle_out, protocol=0) 102 | pickle_out.close() 103 | 104 | 105 | if __name__ == '__main__': 106 | 107 | # modify this path to your Scannet v2 dataset Path 108 | root = "../data/ScanNet" 109 | gen_pickle(split='train', keep_unanno=False, root=root) 110 | gen_pickle(split='val', keep_unanno=False, root=root) 111 | gen_pickle(split='val', keep_unanno=True, root=root) 112 | gen_pickle(split='test', keep_unanno=True, root=root) 113 | 114 | print('Done!!!') 115 | -------------------------------------------------------------------------------- /prepare_data/scannetv2_test.txt: -------------------------------------------------------------------------------- 1 | scene0707_00 2 | scene0708_00 3 | scene0709_00 4 | scene0710_00 5 | scene0711_00 6 | scene0712_00 7 | scene0713_00 8 | scene0714_00 9 | scene0715_00 10 | scene0716_00 11 | scene0717_00 12 | scene0718_00 13 | scene0719_00 14 | scene0720_00 15 | scene0721_00 16 | scene0722_00 17 | scene0723_00 18 | scene0724_00 19 | scene0725_00 20 | scene0726_00 21 | scene0727_00 22 | scene0728_00 23 | scene0729_00 24 | scene0730_00 25 | scene0731_00 26 | scene0732_00 27 | scene0733_00 28 | scene0734_00 29 | scene0735_00 30 | scene0736_00 31 | scene0737_00 32 | scene0738_00 33 | scene0739_00 34 | scene0740_00 35 | scene0741_00 36 | scene0742_00 37 | scene0743_00 38 | scene0744_00 39 | scene0745_00 40 | scene0746_00 41 | scene0747_00 42 | scene0748_00 43 | scene0749_00 44 | scene0750_00 45 | scene0751_00 46 | scene0752_00 47 | scene0753_00 48 | scene0754_00 49 | scene0755_00 50 | scene0756_00 51 | scene0757_00 52 | scene0758_00 53 | scene0759_00 54 | scene0760_00 55 | scene0761_00 56 | scene0762_00 57 | scene0763_00 58 | scene0764_00 59 | scene0765_00 60 | scene0766_00 61 | scene0767_00 62 | scene0768_00 63 | scene0769_00 64 | scene0770_00 65 | scene0771_00 66 | scene0772_00 67 | scene0773_00 68 | scene0774_00 69 | scene0775_00 70 | scene0776_00 71 | scene0777_00 72 | scene0778_00 73 | scene0779_00 74 | scene0780_00 75 | scene0781_00 76 | scene0782_00 77 | scene0783_00 78 | scene0784_00 79 | scene0785_00 80 | scene0786_00 81 | scene0787_00 82 | scene0788_00 83 | scene0789_00 84 | scene0790_00 85 | scene0791_00 86 | scene0792_00 87 | scene0793_00 88 | scene0794_00 89 | scene0795_00 90 | scene0796_00 91 | scene0797_00 92 | scene0798_00 93 | scene0799_00 94 | scene0800_00 95 | scene0801_00 96 | scene0802_00 97 | scene0803_00 98 | scene0804_00 99 | scene0805_00 100 | scene0806_00 101 | -------------------------------------------------------------------------------- /prepare_data/scannetv2_train.txt: -------------------------------------------------------------------------------- 1 | scene0191_00 2 | scene0191_01 3 | scene0191_02 4 | scene0119_00 5 | scene0230_00 6 | scene0528_00 7 | scene0528_01 8 | scene0705_00 9 | scene0705_01 10 | scene0705_02 11 | scene0415_00 12 | scene0415_01 13 | scene0415_02 14 | scene0007_00 15 | scene0141_00 16 | scene0141_01 17 | scene0141_02 18 | scene0515_00 19 | scene0515_01 20 | scene0515_02 21 | scene0447_00 22 | scene0447_01 23 | scene0447_02 24 | scene0531_00 25 | scene0503_00 26 | scene0285_00 27 | scene0069_00 28 | scene0584_00 29 | scene0584_01 30 | scene0584_02 31 | scene0581_00 32 | scene0581_01 33 | scene0581_02 34 | scene0620_00 35 | scene0620_01 36 | scene0263_00 37 | scene0263_01 38 | scene0481_00 39 | scene0481_01 40 | scene0020_00 41 | scene0020_01 42 | scene0291_00 43 | scene0291_01 44 | scene0291_02 45 | scene0469_00 46 | scene0469_01 47 | scene0469_02 48 | scene0659_00 49 | scene0659_01 50 | scene0024_00 51 | scene0024_01 52 | scene0024_02 53 | scene0564_00 54 | scene0117_00 55 | scene0027_00 56 | scene0027_01 57 | scene0027_02 58 | scene0028_00 59 | scene0330_00 60 | scene0418_00 61 | scene0418_01 62 | scene0418_02 63 | scene0233_00 64 | scene0233_01 65 | scene0673_00 66 | scene0673_01 67 | scene0673_02 68 | scene0673_03 69 | scene0673_04 70 | scene0673_05 71 | scene0585_00 72 | scene0585_01 73 | scene0362_00 74 | scene0362_01 75 | scene0362_02 76 | scene0362_03 77 | scene0035_00 78 | scene0035_01 79 | scene0358_00 80 | scene0358_01 81 | scene0358_02 82 | scene0037_00 83 | scene0194_00 84 | scene0321_00 85 | scene0293_00 86 | scene0293_01 87 | scene0623_00 88 | scene0623_01 89 | scene0592_00 90 | scene0592_01 91 | scene0569_00 92 | scene0569_01 93 | scene0413_00 94 | scene0313_00 95 | scene0313_01 96 | scene0313_02 97 | scene0480_00 98 | scene0480_01 99 | scene0401_00 100 | scene0517_00 101 | scene0517_01 102 | scene0517_02 103 | scene0032_00 104 | scene0032_01 105 | scene0613_00 106 | scene0613_01 107 | scene0613_02 108 | scene0306_00 109 | scene0306_01 110 | scene0052_00 111 | scene0052_01 112 | scene0052_02 113 | scene0053_00 114 | scene0444_00 115 | scene0444_01 116 | scene0055_00 117 | scene0055_01 118 | scene0055_02 119 | scene0560_00 120 | scene0589_00 121 | scene0589_01 122 | scene0589_02 123 | scene0610_00 124 | scene0610_01 125 | scene0610_02 126 | scene0364_00 127 | scene0364_01 128 | scene0383_00 129 | scene0383_01 130 | scene0383_02 131 | scene0006_00 132 | scene0006_01 133 | scene0006_02 134 | scene0275_00 135 | scene0451_00 136 | scene0451_01 137 | scene0451_02 138 | scene0451_03 139 | scene0451_04 140 | scene0451_05 141 | scene0135_00 142 | scene0065_00 143 | scene0065_01 144 | scene0065_02 145 | scene0104_00 146 | scene0674_00 147 | scene0674_01 148 | scene0448_00 149 | scene0448_01 150 | scene0448_02 151 | scene0502_00 152 | scene0502_01 153 | scene0502_02 154 | scene0440_00 155 | scene0440_01 156 | scene0440_02 157 | scene0071_00 158 | scene0072_00 159 | scene0072_01 160 | scene0072_02 161 | scene0509_00 162 | scene0509_01 163 | scene0509_02 164 | scene0649_00 165 | scene0649_01 166 | scene0602_00 167 | scene0694_00 168 | scene0694_01 169 | scene0101_00 170 | scene0101_01 171 | scene0101_02 172 | scene0101_03 173 | scene0101_04 174 | scene0101_05 175 | scene0218_00 176 | scene0218_01 177 | scene0579_00 178 | scene0579_01 179 | scene0579_02 180 | scene0039_00 181 | scene0039_01 182 | scene0493_00 183 | scene0493_01 184 | scene0242_00 185 | scene0242_01 186 | scene0242_02 187 | scene0083_00 188 | scene0083_01 189 | scene0127_00 190 | scene0127_01 191 | scene0662_00 192 | scene0662_01 193 | scene0662_02 194 | scene0018_00 195 | scene0087_00 196 | scene0087_01 197 | scene0087_02 198 | scene0332_00 199 | scene0332_01 200 | scene0332_02 201 | scene0628_00 202 | scene0628_01 203 | scene0628_02 204 | scene0134_00 205 | scene0134_01 206 | scene0134_02 207 | scene0238_00 208 | scene0238_01 209 | scene0092_00 210 | scene0092_01 211 | scene0092_02 212 | scene0092_03 213 | scene0092_04 214 | scene0022_00 215 | scene0022_01 216 | scene0467_00 217 | scene0392_00 218 | scene0392_01 219 | scene0392_02 220 | scene0424_00 221 | scene0424_01 222 | scene0424_02 223 | scene0646_00 224 | scene0646_01 225 | scene0646_02 226 | scene0098_00 227 | scene0098_01 228 | scene0044_00 229 | scene0044_01 230 | scene0044_02 231 | scene0510_00 232 | scene0510_01 233 | scene0510_02 234 | scene0571_00 235 | scene0571_01 236 | scene0166_00 237 | scene0166_01 238 | scene0166_02 239 | scene0563_00 240 | scene0172_00 241 | scene0172_01 242 | scene0388_00 243 | scene0388_01 244 | scene0215_00 245 | scene0215_01 246 | scene0252_00 247 | scene0287_00 248 | scene0668_00 249 | scene0572_00 250 | scene0572_01 251 | scene0572_02 252 | scene0026_00 253 | scene0224_00 254 | scene0113_00 255 | scene0113_01 256 | scene0551_00 257 | scene0381_00 258 | scene0381_01 259 | scene0381_02 260 | scene0371_00 261 | scene0371_01 262 | scene0460_00 263 | scene0118_00 264 | scene0118_01 265 | scene0118_02 266 | scene0417_00 267 | scene0008_00 268 | scene0634_00 269 | scene0521_00 270 | scene0123_00 271 | scene0123_01 272 | scene0123_02 273 | scene0045_00 274 | scene0045_01 275 | scene0511_00 276 | scene0511_01 277 | scene0114_00 278 | scene0114_01 279 | scene0114_02 280 | scene0070_00 281 | scene0029_00 282 | scene0029_01 283 | scene0029_02 284 | scene0129_00 285 | scene0103_00 286 | scene0103_01 287 | scene0002_00 288 | scene0002_01 289 | scene0132_00 290 | scene0132_01 291 | scene0132_02 292 | scene0124_00 293 | scene0124_01 294 | scene0143_00 295 | scene0143_01 296 | scene0143_02 297 | scene0604_00 298 | scene0604_01 299 | scene0604_02 300 | scene0507_00 301 | scene0105_00 302 | scene0105_01 303 | scene0105_02 304 | scene0428_00 305 | scene0428_01 306 | scene0311_00 307 | scene0140_00 308 | scene0140_01 309 | scene0182_00 310 | scene0182_01 311 | scene0182_02 312 | scene0142_00 313 | scene0142_01 314 | scene0399_00 315 | scene0399_01 316 | scene0012_00 317 | scene0012_01 318 | scene0012_02 319 | scene0060_00 320 | scene0060_01 321 | scene0370_00 322 | scene0370_01 323 | scene0370_02 324 | scene0310_00 325 | scene0310_01 326 | scene0310_02 327 | scene0661_00 328 | scene0650_00 329 | scene0152_00 330 | scene0152_01 331 | scene0152_02 332 | scene0158_00 333 | scene0158_01 334 | scene0158_02 335 | scene0482_00 336 | scene0482_01 337 | scene0600_00 338 | scene0600_01 339 | scene0600_02 340 | scene0393_00 341 | scene0393_01 342 | scene0393_02 343 | scene0562_00 344 | scene0174_00 345 | scene0174_01 346 | scene0157_00 347 | scene0157_01 348 | scene0161_00 349 | scene0161_01 350 | scene0161_02 351 | scene0159_00 352 | scene0254_00 353 | scene0254_01 354 | scene0115_00 355 | scene0115_01 356 | scene0115_02 357 | scene0162_00 358 | scene0163_00 359 | scene0163_01 360 | scene0523_00 361 | scene0523_01 362 | scene0523_02 363 | scene0459_00 364 | scene0459_01 365 | scene0175_00 366 | scene0085_00 367 | scene0085_01 368 | scene0279_00 369 | scene0279_01 370 | scene0279_02 371 | scene0201_00 372 | scene0201_01 373 | scene0201_02 374 | scene0283_00 375 | scene0456_00 376 | scene0456_01 377 | scene0429_00 378 | scene0043_00 379 | scene0043_01 380 | scene0419_00 381 | scene0419_01 382 | scene0419_02 383 | scene0368_00 384 | scene0368_01 385 | scene0348_00 386 | scene0348_01 387 | scene0348_02 388 | scene0442_00 389 | scene0178_00 390 | scene0380_00 391 | scene0380_01 392 | scene0380_02 393 | scene0165_00 394 | scene0165_01 395 | scene0165_02 396 | scene0181_00 397 | scene0181_01 398 | scene0181_02 399 | scene0181_03 400 | scene0333_00 401 | scene0614_00 402 | scene0614_01 403 | scene0614_02 404 | scene0404_00 405 | scene0404_01 406 | scene0404_02 407 | scene0185_00 408 | scene0126_00 409 | scene0126_01 410 | scene0126_02 411 | scene0519_00 412 | scene0236_00 413 | scene0236_01 414 | scene0189_00 415 | scene0075_00 416 | scene0267_00 417 | scene0192_00 418 | scene0192_01 419 | scene0192_02 420 | scene0281_00 421 | scene0420_00 422 | scene0420_01 423 | scene0420_02 424 | scene0195_00 425 | scene0195_01 426 | scene0195_02 427 | scene0597_00 428 | scene0597_01 429 | scene0597_02 430 | scene0041_00 431 | scene0041_01 432 | scene0111_00 433 | scene0111_01 434 | scene0111_02 435 | scene0666_00 436 | scene0666_01 437 | scene0666_02 438 | scene0200_00 439 | scene0200_01 440 | scene0200_02 441 | scene0536_00 442 | scene0536_01 443 | scene0536_02 444 | scene0390_00 445 | scene0280_00 446 | scene0280_01 447 | scene0280_02 448 | scene0344_00 449 | scene0344_01 450 | scene0205_00 451 | scene0205_01 452 | scene0205_02 453 | scene0484_00 454 | scene0484_01 455 | scene0009_00 456 | scene0009_01 457 | scene0009_02 458 | scene0302_00 459 | scene0302_01 460 | scene0209_00 461 | scene0209_01 462 | scene0209_02 463 | scene0210_00 464 | scene0210_01 465 | scene0395_00 466 | scene0395_01 467 | scene0395_02 468 | scene0683_00 469 | scene0601_00 470 | scene0601_01 471 | scene0214_00 472 | scene0214_01 473 | scene0214_02 474 | scene0477_00 475 | scene0477_01 476 | scene0439_00 477 | scene0439_01 478 | scene0468_00 479 | scene0468_01 480 | scene0468_02 481 | scene0546_00 482 | scene0466_00 483 | scene0466_01 484 | scene0220_00 485 | scene0220_01 486 | scene0220_02 487 | scene0122_00 488 | scene0122_01 489 | scene0130_00 490 | scene0110_00 491 | scene0110_01 492 | scene0110_02 493 | scene0327_00 494 | scene0156_00 495 | scene0266_00 496 | scene0266_01 497 | scene0001_00 498 | scene0001_01 499 | scene0228_00 500 | scene0199_00 501 | scene0219_00 502 | scene0464_00 503 | scene0232_00 504 | scene0232_01 505 | scene0232_02 506 | scene0299_00 507 | scene0299_01 508 | scene0530_00 509 | scene0363_00 510 | scene0453_00 511 | scene0453_01 512 | scene0570_00 513 | scene0570_01 514 | scene0570_02 515 | scene0183_00 516 | scene0239_00 517 | scene0239_01 518 | scene0239_02 519 | scene0373_00 520 | scene0373_01 521 | scene0241_00 522 | scene0241_01 523 | scene0241_02 524 | scene0188_00 525 | scene0622_00 526 | scene0622_01 527 | scene0244_00 528 | scene0244_01 529 | scene0691_00 530 | scene0691_01 531 | scene0206_00 532 | scene0206_01 533 | scene0206_02 534 | scene0247_00 535 | scene0247_01 536 | scene0061_00 537 | scene0061_01 538 | scene0082_00 539 | scene0250_00 540 | scene0250_01 541 | scene0250_02 542 | scene0501_00 543 | scene0501_01 544 | scene0501_02 545 | scene0320_00 546 | scene0320_01 547 | scene0320_02 548 | scene0320_03 549 | scene0631_00 550 | scene0631_01 551 | scene0631_02 552 | scene0255_00 553 | scene0255_01 554 | scene0255_02 555 | scene0047_00 556 | scene0265_00 557 | scene0265_01 558 | scene0265_02 559 | scene0004_00 560 | scene0336_00 561 | scene0336_01 562 | scene0058_00 563 | scene0058_01 564 | scene0260_00 565 | scene0260_01 566 | scene0260_02 567 | scene0243_00 568 | scene0603_00 569 | scene0603_01 570 | scene0093_00 571 | scene0093_01 572 | scene0093_02 573 | scene0109_00 574 | scene0109_01 575 | scene0434_00 576 | scene0434_01 577 | scene0434_02 578 | scene0290_00 579 | scene0627_00 580 | scene0627_01 581 | scene0470_00 582 | scene0470_01 583 | scene0137_00 584 | scene0137_01 585 | scene0137_02 586 | scene0270_00 587 | scene0270_01 588 | scene0270_02 589 | scene0271_00 590 | scene0271_01 591 | scene0504_00 592 | scene0274_00 593 | scene0274_01 594 | scene0274_02 595 | scene0036_00 596 | scene0036_01 597 | scene0276_00 598 | scene0276_01 599 | scene0272_00 600 | scene0272_01 601 | scene0499_00 602 | scene0698_00 603 | scene0698_01 604 | scene0051_00 605 | scene0051_01 606 | scene0051_02 607 | scene0051_03 608 | scene0108_00 609 | scene0245_00 610 | scene0369_00 611 | scene0369_01 612 | scene0369_02 613 | scene0284_00 614 | scene0289_00 615 | scene0289_01 616 | scene0286_00 617 | scene0286_01 618 | scene0286_02 619 | scene0286_03 620 | scene0031_00 621 | scene0031_01 622 | scene0031_02 623 | scene0545_00 624 | scene0545_01 625 | scene0545_02 626 | scene0557_00 627 | scene0557_01 628 | scene0557_02 629 | scene0533_00 630 | scene0533_01 631 | scene0116_00 632 | scene0116_01 633 | scene0116_02 634 | scene0611_00 635 | scene0611_01 636 | scene0688_00 637 | scene0294_00 638 | scene0294_01 639 | scene0294_02 640 | scene0295_00 641 | scene0295_01 642 | scene0296_00 643 | scene0296_01 644 | scene0596_00 645 | scene0596_01 646 | scene0596_02 647 | scene0532_00 648 | scene0532_01 649 | scene0637_00 650 | scene0638_00 651 | scene0121_00 652 | scene0121_01 653 | scene0121_02 654 | scene0040_00 655 | scene0040_01 656 | scene0197_00 657 | scene0197_01 658 | scene0197_02 659 | scene0410_00 660 | scene0410_01 661 | scene0305_00 662 | scene0305_01 663 | scene0615_00 664 | scene0615_01 665 | scene0703_00 666 | scene0703_01 667 | scene0555_00 668 | scene0297_00 669 | scene0297_01 670 | scene0297_02 671 | scene0582_00 672 | scene0582_01 673 | scene0582_02 674 | scene0023_00 675 | scene0094_00 676 | scene0013_00 677 | scene0013_01 678 | scene0013_02 679 | scene0136_00 680 | scene0136_01 681 | scene0136_02 682 | scene0407_00 683 | scene0407_01 684 | scene0062_00 685 | scene0062_01 686 | scene0062_02 687 | scene0386_00 688 | scene0318_00 689 | scene0554_00 690 | scene0554_01 691 | scene0497_00 692 | scene0213_00 693 | scene0258_00 694 | scene0323_00 695 | scene0323_01 696 | scene0324_00 697 | scene0324_01 698 | scene0016_00 699 | scene0016_01 700 | scene0016_02 701 | scene0681_00 702 | scene0398_00 703 | scene0398_01 704 | scene0227_00 705 | scene0090_00 706 | scene0066_00 707 | scene0262_00 708 | scene0262_01 709 | scene0155_00 710 | scene0155_01 711 | scene0155_02 712 | scene0352_00 713 | scene0352_01 714 | scene0352_02 715 | scene0038_00 716 | scene0038_01 717 | scene0038_02 718 | scene0335_00 719 | scene0335_01 720 | scene0335_02 721 | scene0261_00 722 | scene0261_01 723 | scene0261_02 724 | scene0261_03 725 | scene0640_00 726 | scene0640_01 727 | scene0640_02 728 | scene0080_00 729 | scene0080_01 730 | scene0080_02 731 | scene0403_00 732 | scene0403_01 733 | scene0282_00 734 | scene0282_01 735 | scene0282_02 736 | scene0682_00 737 | scene0173_00 738 | scene0173_01 739 | scene0173_02 740 | scene0522_00 741 | scene0687_00 742 | scene0345_00 743 | scene0345_01 744 | scene0612_00 745 | scene0612_01 746 | scene0411_00 747 | scene0411_01 748 | scene0411_02 749 | scene0625_00 750 | scene0625_01 751 | scene0211_00 752 | scene0211_01 753 | scene0211_02 754 | scene0211_03 755 | scene0676_00 756 | scene0676_01 757 | scene0179_00 758 | scene0498_00 759 | scene0498_01 760 | scene0498_02 761 | scene0547_00 762 | scene0547_01 763 | scene0547_02 764 | scene0269_00 765 | scene0269_01 766 | scene0269_02 767 | scene0366_00 768 | scene0680_00 769 | scene0680_01 770 | scene0588_00 771 | scene0588_01 772 | scene0588_02 773 | scene0588_03 774 | scene0346_00 775 | scene0346_01 776 | scene0359_00 777 | scene0359_01 778 | scene0014_00 779 | scene0120_00 780 | scene0120_01 781 | scene0212_00 782 | scene0212_01 783 | scene0212_02 784 | scene0176_00 785 | scene0049_00 786 | scene0259_00 787 | scene0259_01 788 | scene0586_00 789 | scene0586_01 790 | scene0586_02 791 | scene0309_00 792 | scene0309_01 793 | scene0125_00 794 | scene0455_00 795 | scene0177_00 796 | scene0177_01 797 | scene0177_02 798 | scene0326_00 799 | scene0372_00 800 | scene0171_00 801 | scene0171_01 802 | scene0374_00 803 | scene0654_00 804 | scene0654_01 805 | scene0445_00 806 | scene0445_01 807 | scene0475_00 808 | scene0475_01 809 | scene0475_02 810 | scene0349_00 811 | scene0349_01 812 | scene0234_00 813 | scene0669_00 814 | scene0669_01 815 | scene0375_00 816 | scene0375_01 817 | scene0375_02 818 | scene0387_00 819 | scene0387_01 820 | scene0387_02 821 | scene0312_00 822 | scene0312_01 823 | scene0312_02 824 | scene0384_00 825 | scene0385_00 826 | scene0385_01 827 | scene0385_02 828 | scene0000_00 829 | scene0000_01 830 | scene0000_02 831 | scene0376_00 832 | scene0376_01 833 | scene0376_02 834 | scene0301_00 835 | scene0301_01 836 | scene0301_02 837 | scene0322_00 838 | scene0542_00 839 | scene0079_00 840 | scene0079_01 841 | scene0099_00 842 | scene0099_01 843 | scene0476_00 844 | scene0476_01 845 | scene0476_02 846 | scene0394_00 847 | scene0394_01 848 | scene0147_00 849 | scene0147_01 850 | scene0067_00 851 | scene0067_01 852 | scene0067_02 853 | scene0397_00 854 | scene0397_01 855 | scene0337_00 856 | scene0337_01 857 | scene0337_02 858 | scene0431_00 859 | scene0223_00 860 | scene0223_01 861 | scene0223_02 862 | scene0010_00 863 | scene0010_01 864 | scene0402_00 865 | scene0268_00 866 | scene0268_01 867 | scene0268_02 868 | scene0679_00 869 | scene0679_01 870 | scene0405_00 871 | scene0128_00 872 | scene0408_00 873 | scene0408_01 874 | scene0190_00 875 | scene0107_00 876 | scene0076_00 877 | scene0167_00 878 | scene0361_00 879 | scene0361_01 880 | scene0361_02 881 | scene0216_00 882 | scene0202_00 883 | scene0303_00 884 | scene0303_01 885 | scene0303_02 886 | scene0446_00 887 | scene0446_01 888 | scene0089_00 889 | scene0089_01 890 | scene0089_02 891 | scene0360_00 892 | scene0150_00 893 | scene0150_01 894 | scene0150_02 895 | scene0421_00 896 | scene0421_01 897 | scene0421_02 898 | scene0454_00 899 | scene0626_00 900 | scene0626_01 901 | scene0626_02 902 | scene0186_00 903 | scene0186_01 904 | scene0538_00 905 | scene0479_00 906 | scene0479_01 907 | scene0479_02 908 | scene0656_00 909 | scene0656_01 910 | scene0656_02 911 | scene0656_03 912 | scene0525_00 913 | scene0525_01 914 | scene0525_02 915 | scene0308_00 916 | scene0396_00 917 | scene0396_01 918 | scene0396_02 919 | scene0624_00 920 | scene0292_00 921 | scene0292_01 922 | scene0632_00 923 | scene0253_00 924 | scene0021_00 925 | scene0325_00 926 | scene0325_01 927 | scene0437_00 928 | scene0437_01 929 | scene0438_00 930 | scene0590_00 931 | scene0590_01 932 | scene0400_00 933 | scene0400_01 934 | scene0541_00 935 | scene0541_01 936 | scene0541_02 937 | scene0677_00 938 | scene0677_01 939 | scene0677_02 940 | scene0443_00 941 | scene0315_00 942 | scene0288_00 943 | scene0288_01 944 | scene0288_02 945 | scene0422_00 946 | scene0672_00 947 | scene0672_01 948 | scene0184_00 949 | scene0449_00 950 | scene0449_01 951 | scene0449_02 952 | scene0048_00 953 | scene0048_01 954 | scene0138_00 955 | scene0452_00 956 | scene0452_01 957 | scene0452_02 958 | scene0667_00 959 | scene0667_01 960 | scene0667_02 961 | scene0463_00 962 | scene0463_01 963 | scene0078_00 964 | scene0078_01 965 | scene0078_02 966 | scene0636_00 967 | scene0457_00 968 | scene0457_01 969 | scene0457_02 970 | scene0465_00 971 | scene0465_01 972 | scene0577_00 973 | scene0151_00 974 | scene0151_01 975 | scene0339_00 976 | scene0573_00 977 | scene0573_01 978 | scene0154_00 979 | scene0096_00 980 | scene0096_01 981 | scene0096_02 982 | scene0235_00 983 | scene0168_00 984 | scene0168_01 985 | scene0168_02 986 | scene0594_00 987 | scene0587_00 988 | scene0587_01 989 | scene0587_02 990 | scene0587_03 991 | scene0229_00 992 | scene0229_01 993 | scene0229_02 994 | scene0512_00 995 | scene0106_00 996 | scene0106_01 997 | scene0106_02 998 | scene0472_00 999 | scene0472_01 1000 | scene0472_02 1001 | scene0489_00 1002 | scene0489_01 1003 | scene0489_02 1004 | scene0425_00 1005 | scene0425_01 1006 | scene0641_00 1007 | scene0526_00 1008 | scene0526_01 1009 | scene0317_00 1010 | scene0317_01 1011 | scene0544_00 1012 | scene0017_00 1013 | scene0017_01 1014 | scene0017_02 1015 | scene0042_00 1016 | scene0042_01 1017 | scene0042_02 1018 | scene0576_00 1019 | scene0576_01 1020 | scene0576_02 1021 | scene0347_00 1022 | scene0347_01 1023 | scene0347_02 1024 | scene0436_00 1025 | scene0226_00 1026 | scene0226_01 1027 | scene0485_00 1028 | scene0486_00 1029 | scene0487_00 1030 | scene0487_01 1031 | scene0619_00 1032 | scene0097_00 1033 | scene0367_00 1034 | scene0367_01 1035 | scene0491_00 1036 | scene0492_00 1037 | scene0492_01 1038 | scene0005_00 1039 | scene0005_01 1040 | scene0543_00 1041 | scene0543_01 1042 | scene0543_02 1043 | scene0657_00 1044 | scene0341_00 1045 | scene0341_01 1046 | scene0534_00 1047 | scene0534_01 1048 | scene0319_00 1049 | scene0273_00 1050 | scene0273_01 1051 | scene0225_00 1052 | scene0198_00 1053 | scene0003_00 1054 | scene0003_01 1055 | scene0003_02 1056 | scene0409_00 1057 | scene0409_01 1058 | scene0331_00 1059 | scene0331_01 1060 | scene0505_00 1061 | scene0505_01 1062 | scene0505_02 1063 | scene0505_03 1064 | scene0505_04 1065 | scene0506_00 1066 | scene0057_00 1067 | scene0057_01 1068 | scene0074_00 1069 | scene0074_01 1070 | scene0074_02 1071 | scene0091_00 1072 | scene0112_00 1073 | scene0112_01 1074 | scene0112_02 1075 | scene0240_00 1076 | scene0102_00 1077 | scene0102_01 1078 | scene0513_00 1079 | scene0514_00 1080 | scene0514_01 1081 | scene0537_00 1082 | scene0516_00 1083 | scene0516_01 1084 | scene0495_00 1085 | scene0617_00 1086 | scene0133_00 1087 | scene0520_00 1088 | scene0520_01 1089 | scene0635_00 1090 | scene0635_01 1091 | scene0054_00 1092 | scene0473_00 1093 | scene0473_01 1094 | scene0524_00 1095 | scene0524_01 1096 | scene0379_00 1097 | scene0471_00 1098 | scene0471_01 1099 | scene0471_02 1100 | scene0566_00 1101 | scene0248_00 1102 | scene0248_01 1103 | scene0248_02 1104 | scene0529_00 1105 | scene0529_01 1106 | scene0529_02 1107 | scene0391_00 1108 | scene0264_00 1109 | scene0264_01 1110 | scene0264_02 1111 | scene0675_00 1112 | scene0675_01 1113 | scene0350_00 1114 | scene0350_01 1115 | scene0350_02 1116 | scene0450_00 1117 | scene0068_00 1118 | scene0068_01 1119 | scene0237_00 1120 | scene0237_01 1121 | scene0365_00 1122 | scene0365_01 1123 | scene0365_02 1124 | scene0605_00 1125 | scene0605_01 1126 | scene0539_00 1127 | scene0539_01 1128 | scene0539_02 1129 | scene0540_00 1130 | scene0540_01 1131 | scene0540_02 1132 | scene0170_00 1133 | scene0170_01 1134 | scene0170_02 1135 | scene0433_00 1136 | scene0340_00 1137 | scene0340_01 1138 | scene0340_02 1139 | scene0160_00 1140 | scene0160_01 1141 | scene0160_02 1142 | scene0160_03 1143 | scene0160_04 1144 | scene0059_00 1145 | scene0059_01 1146 | scene0059_02 1147 | scene0056_00 1148 | scene0056_01 1149 | scene0478_00 1150 | scene0478_01 1151 | scene0548_00 1152 | scene0548_01 1153 | scene0548_02 1154 | scene0204_00 1155 | scene0204_01 1156 | scene0204_02 1157 | scene0033_00 1158 | scene0145_00 1159 | scene0483_00 1160 | scene0508_00 1161 | scene0508_01 1162 | scene0508_02 1163 | scene0180_00 1164 | scene0148_00 1165 | scene0556_00 1166 | scene0556_01 1167 | scene0416_00 1168 | scene0416_01 1169 | scene0416_02 1170 | scene0416_03 1171 | scene0416_04 1172 | scene0073_00 1173 | scene0073_01 1174 | scene0073_02 1175 | scene0073_03 1176 | scene0034_00 1177 | scene0034_01 1178 | scene0034_02 1179 | scene0639_00 1180 | scene0561_00 1181 | scene0561_01 1182 | scene0298_00 1183 | scene0692_00 1184 | scene0692_01 1185 | scene0692_02 1186 | scene0692_03 1187 | scene0692_04 1188 | scene0642_00 1189 | scene0642_01 1190 | scene0642_02 1191 | scene0642_03 1192 | scene0630_00 1193 | scene0630_01 1194 | scene0630_02 1195 | scene0630_03 1196 | scene0630_04 1197 | scene0630_05 1198 | scene0630_06 1199 | scene0706_00 1200 | scene0567_00 1201 | scene0567_01 1202 | -------------------------------------------------------------------------------- /prepare_data/scannetv2_val.txt: -------------------------------------------------------------------------------- 1 | scene0568_00 2 | scene0568_01 3 | scene0568_02 4 | scene0304_00 5 | scene0488_00 6 | scene0488_01 7 | scene0412_00 8 | scene0412_01 9 | scene0217_00 10 | scene0019_00 11 | scene0019_01 12 | scene0414_00 13 | scene0575_00 14 | scene0575_01 15 | scene0575_02 16 | scene0426_00 17 | scene0426_01 18 | scene0426_02 19 | scene0426_03 20 | scene0549_00 21 | scene0549_01 22 | scene0578_00 23 | scene0578_01 24 | scene0578_02 25 | scene0665_00 26 | scene0665_01 27 | scene0050_00 28 | scene0050_01 29 | scene0050_02 30 | scene0257_00 31 | scene0025_00 32 | scene0025_01 33 | scene0025_02 34 | scene0583_00 35 | scene0583_01 36 | scene0583_02 37 | scene0701_00 38 | scene0701_01 39 | scene0701_02 40 | scene0580_00 41 | scene0580_01 42 | scene0565_00 43 | scene0169_00 44 | scene0169_01 45 | scene0655_00 46 | scene0655_01 47 | scene0655_02 48 | scene0063_00 49 | scene0221_00 50 | scene0221_01 51 | scene0591_00 52 | scene0591_01 53 | scene0591_02 54 | scene0678_00 55 | scene0678_01 56 | scene0678_02 57 | scene0462_00 58 | scene0427_00 59 | scene0595_00 60 | scene0193_00 61 | scene0193_01 62 | scene0164_00 63 | scene0164_01 64 | scene0164_02 65 | scene0164_03 66 | scene0598_00 67 | scene0598_01 68 | scene0598_02 69 | scene0599_00 70 | scene0599_01 71 | scene0599_02 72 | scene0328_00 73 | scene0300_00 74 | scene0300_01 75 | scene0354_00 76 | scene0458_00 77 | scene0458_01 78 | scene0423_00 79 | scene0423_01 80 | scene0423_02 81 | scene0307_00 82 | scene0307_01 83 | scene0307_02 84 | scene0606_00 85 | scene0606_01 86 | scene0606_02 87 | scene0432_00 88 | scene0432_01 89 | scene0608_00 90 | scene0608_01 91 | scene0608_02 92 | scene0651_00 93 | scene0651_01 94 | scene0651_02 95 | scene0430_00 96 | scene0430_01 97 | scene0689_00 98 | scene0357_00 99 | scene0357_01 100 | scene0574_00 101 | scene0574_01 102 | scene0574_02 103 | scene0329_00 104 | scene0329_01 105 | scene0329_02 106 | scene0153_00 107 | scene0153_01 108 | scene0616_00 109 | scene0616_01 110 | scene0671_00 111 | scene0671_01 112 | scene0618_00 113 | scene0382_00 114 | scene0382_01 115 | scene0490_00 116 | scene0621_00 117 | scene0607_00 118 | scene0607_01 119 | scene0149_00 120 | scene0695_00 121 | scene0695_01 122 | scene0695_02 123 | scene0695_03 124 | scene0389_00 125 | scene0377_00 126 | scene0377_01 127 | scene0377_02 128 | scene0342_00 129 | scene0139_00 130 | scene0629_00 131 | scene0629_01 132 | scene0629_02 133 | scene0496_00 134 | scene0633_00 135 | scene0633_01 136 | scene0518_00 137 | scene0652_00 138 | scene0406_00 139 | scene0406_01 140 | scene0406_02 141 | scene0144_00 142 | scene0144_01 143 | scene0494_00 144 | scene0278_00 145 | scene0278_01 146 | scene0316_00 147 | scene0609_00 148 | scene0609_01 149 | scene0609_02 150 | scene0609_03 151 | scene0084_00 152 | scene0084_01 153 | scene0084_02 154 | scene0696_00 155 | scene0696_01 156 | scene0696_02 157 | scene0351_00 158 | scene0351_01 159 | scene0643_00 160 | scene0644_00 161 | scene0645_00 162 | scene0645_01 163 | scene0645_02 164 | scene0081_00 165 | scene0081_01 166 | scene0081_02 167 | scene0647_00 168 | scene0647_01 169 | scene0535_00 170 | scene0353_00 171 | scene0353_01 172 | scene0353_02 173 | scene0559_00 174 | scene0559_01 175 | scene0559_02 176 | scene0593_00 177 | scene0593_01 178 | scene0246_00 179 | scene0653_00 180 | scene0653_01 181 | scene0064_00 182 | scene0064_01 183 | scene0356_00 184 | scene0356_01 185 | scene0356_02 186 | scene0030_00 187 | scene0030_01 188 | scene0030_02 189 | scene0222_00 190 | scene0222_01 191 | scene0338_00 192 | scene0338_01 193 | scene0338_02 194 | scene0378_00 195 | scene0378_01 196 | scene0378_02 197 | scene0660_00 198 | scene0553_00 199 | scene0553_01 200 | scene0553_02 201 | scene0527_00 202 | scene0663_00 203 | scene0663_01 204 | scene0663_02 205 | scene0664_00 206 | scene0664_01 207 | scene0664_02 208 | scene0334_00 209 | scene0334_01 210 | scene0334_02 211 | scene0046_00 212 | scene0046_01 213 | scene0046_02 214 | scene0203_00 215 | scene0203_01 216 | scene0203_02 217 | scene0088_00 218 | scene0088_01 219 | scene0088_02 220 | scene0088_03 221 | scene0086_00 222 | scene0086_01 223 | scene0086_02 224 | scene0670_00 225 | scene0670_01 226 | scene0256_00 227 | scene0256_01 228 | scene0256_02 229 | scene0249_00 230 | scene0441_00 231 | scene0658_00 232 | scene0704_00 233 | scene0704_01 234 | scene0187_00 235 | scene0187_01 236 | scene0131_00 237 | scene0131_01 238 | scene0131_02 239 | scene0207_00 240 | scene0207_01 241 | scene0207_02 242 | scene0461_00 243 | scene0011_00 244 | scene0011_01 245 | scene0343_00 246 | scene0251_00 247 | scene0077_00 248 | scene0077_01 249 | scene0684_00 250 | scene0684_01 251 | scene0550_00 252 | scene0686_00 253 | scene0686_01 254 | scene0686_02 255 | scene0208_00 256 | scene0500_00 257 | scene0500_01 258 | scene0552_00 259 | scene0552_01 260 | scene0648_00 261 | scene0648_01 262 | scene0435_00 263 | scene0435_01 264 | scene0435_02 265 | scene0435_03 266 | scene0690_00 267 | scene0690_01 268 | scene0693_00 269 | scene0693_01 270 | scene0693_02 271 | scene0700_00 272 | scene0700_01 273 | scene0700_02 274 | scene0699_00 275 | scene0231_00 276 | scene0231_01 277 | scene0231_02 278 | scene0697_00 279 | scene0697_01 280 | scene0697_02 281 | scene0697_03 282 | scene0474_00 283 | scene0474_01 284 | scene0474_02 285 | scene0474_03 286 | scene0474_04 287 | scene0474_05 288 | scene0355_00 289 | scene0355_01 290 | scene0146_00 291 | scene0146_01 292 | scene0146_02 293 | scene0196_00 294 | scene0702_00 295 | scene0702_01 296 | scene0702_02 297 | scene0314_00 298 | scene0277_00 299 | scene0277_01 300 | scene0277_02 301 | scene0095_00 302 | scene0095_01 303 | scene0015_00 304 | scene0100_00 305 | scene0100_01 306 | scene0100_02 307 | scene0558_00 308 | scene0558_01 309 | scene0558_02 310 | scene0685_00 311 | scene0685_01 312 | scene0685_02 313 | -------------------------------------------------------------------------------- /prepare_data/util.py: -------------------------------------------------------------------------------- 1 | import os, sys 2 | import csv 3 | try: 4 | import numpy as np 5 | except: 6 | print("Failed to import numpy package.") 7 | sys.exit(-1) 8 | try: 9 | import imageio 10 | except: 11 | print("Please install the module 'imageio' for image processing, e.g.") 12 | print("pip install imageio") 13 | sys.exit(-1) 14 | 15 | # print an error message and quit 16 | def print_error(message, user_fault=False): 17 | sys.stderr.write('ERROR: ' + str(message) + '\n') 18 | if user_fault: 19 | sys.exit(2) 20 | sys.exit(-1) 21 | 22 | 23 | # if string s represents an int 24 | def represents_int(s): 25 | try: 26 | int(s) 27 | return True 28 | except ValueError: 29 | return False 30 | 31 | 32 | def read_label_mapping(filename, label_from='raw_category', label_to='nyu40id'): 33 | assert os.path.isfile(filename) 34 | mapping = dict() 35 | with open(filename) as csvfile: 36 | reader = csv.DictReader(csvfile, delimiter='\t') 37 | for row in reader: 38 | mapping[row[label_from]] = int(row[label_to]) 39 | # if ints convert 40 | if represents_int(mapping.keys()[0]): 41 | mapping = {int(k):v for k,v in mapping.items()} 42 | return mapping 43 | 44 | 45 | # input: scene_types.txt or scene_types_all.txt 46 | def read_scene_types_mapping(filename, remove_spaces=True): 47 | assert os.path.isfile(filename) 48 | mapping = dict() 49 | lines = open(filename).read().splitlines() 50 | lines = [line.split('\t') for line in lines] 51 | if remove_spaces: 52 | mapping = { x[1].strip():int(x[0]) for x in lines } 53 | else: 54 | mapping = { x[1]:int(x[0]) for x in lines } 55 | return mapping 56 | 57 | 58 | # color by label 59 | def visualize_label_image(filename, image): 60 | height = image.shape[0] 61 | width = image.shape[1] 62 | vis_image = np.zeros([height, width, 3], dtype=np.uint8) 63 | color_palette = create_color_palette() 64 | for idx, color in enumerate(color_palette): 65 | vis_image[image==idx] = color 66 | imageio.imwrite(filename, vis_image) 67 | 68 | 69 | # color by different instances (mod length of color palette) 70 | def visualize_instance_image(filename, image): 71 | height = image.shape[0] 72 | width = image.shape[1] 73 | vis_image = np.zeros([height, width, 3], dtype=np.uint8) 74 | color_palette = create_color_palette() 75 | instances = np.unique(image) 76 | for idx, inst in enumerate(instances): 77 | vis_image[image==inst] = color_palette[inst%len(color_palette)] 78 | imageio.imwrite(filename, vis_image) 79 | 80 | 81 | # color palette for nyu40 labels 82 | def create_color_palette(): 83 | return [ 84 | (0, 0, 0), 85 | (174, 199, 232), # wall 86 | (152, 223, 138), # floor 87 | (31, 119, 180), # cabinet 88 | (255, 187, 120), # bed 89 | (188, 189, 34), # chair 90 | (140, 86, 75), # sofa 91 | (255, 152, 150), # table 92 | (214, 39, 40), # door 93 | (197, 176, 213), # window 94 | (148, 103, 189), # bookshelf 95 | (196, 156, 148), # picture 96 | (23, 190, 207), # counter 97 | (178, 76, 76), 98 | (247, 182, 210), # desk 99 | (66, 188, 102), 100 | (219, 219, 141), # curtain 101 | (140, 57, 197), 102 | (202, 185, 52), 103 | (51, 176, 203), 104 | (200, 54, 131), 105 | (92, 193, 61), 106 | (78, 71, 183), 107 | (172, 114, 82), 108 | (255, 127, 14), # refrigerator 109 | (91, 163, 138), 110 | (153, 98, 156), 111 | (140, 153, 101), 112 | (158, 218, 229), # shower curtain 113 | (100, 125, 154), 114 | (178, 127, 135), 115 | (120, 185, 128), 116 | (146, 111, 194), 117 | (44, 160, 44), # toilet 118 | (112, 128, 144), # sink 119 | (96, 207, 209), 120 | (227, 119, 194), # bathtub 121 | (213, 92, 176), 122 | (94, 106, 211), 123 | (82, 84, 163), # otherfurn 124 | (100, 85, 144) 125 | ] 126 | -------------------------------------------------------------------------------- /pretrained/model.cls.1024.t7: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/antao97/dgcnn.pytorch/07d534c2702905010ec9991619f552d8cacae45b/pretrained/model.cls.1024.t7 -------------------------------------------------------------------------------- /pretrained/model.cls.2048.t7: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/antao97/dgcnn.pytorch/07d534c2702905010ec9991619f552d8cacae45b/pretrained/model.cls.2048.t7 -------------------------------------------------------------------------------- /pretrained/model.partseg.airplane.t7: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/antao97/dgcnn.pytorch/07d534c2702905010ec9991619f552d8cacae45b/pretrained/model.partseg.airplane.t7 -------------------------------------------------------------------------------- /pretrained/model.partseg.car.t7: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/antao97/dgcnn.pytorch/07d534c2702905010ec9991619f552d8cacae45b/pretrained/model.partseg.car.t7 -------------------------------------------------------------------------------- /pretrained/model.partseg.t7: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/antao97/dgcnn.pytorch/07d534c2702905010ec9991619f552d8cacae45b/pretrained/model.partseg.t7 -------------------------------------------------------------------------------- /pretrained/semseg_s3dis/model_1.t7: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/antao97/dgcnn.pytorch/07d534c2702905010ec9991619f552d8cacae45b/pretrained/semseg_s3dis/model_1.t7 -------------------------------------------------------------------------------- /pretrained/semseg_s3dis/model_2.t7: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/antao97/dgcnn.pytorch/07d534c2702905010ec9991619f552d8cacae45b/pretrained/semseg_s3dis/model_2.t7 -------------------------------------------------------------------------------- /pretrained/semseg_s3dis/model_3.t7: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/antao97/dgcnn.pytorch/07d534c2702905010ec9991619f552d8cacae45b/pretrained/semseg_s3dis/model_3.t7 -------------------------------------------------------------------------------- /pretrained/semseg_s3dis/model_4.t7: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/antao97/dgcnn.pytorch/07d534c2702905010ec9991619f552d8cacae45b/pretrained/semseg_s3dis/model_4.t7 -------------------------------------------------------------------------------- /pretrained/semseg_s3dis/model_5.t7: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/antao97/dgcnn.pytorch/07d534c2702905010ec9991619f552d8cacae45b/pretrained/semseg_s3dis/model_5.t7 -------------------------------------------------------------------------------- /pretrained/semseg_s3dis/model_6.t7: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/antao97/dgcnn.pytorch/07d534c2702905010ec9991619f552d8cacae45b/pretrained/semseg_s3dis/model_6.t7 -------------------------------------------------------------------------------- /pretrained/semseg_scannet/models/model_200.pth: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/antao97/dgcnn.pytorch/07d534c2702905010ec9991619f552d8cacae45b/pretrained/semseg_scannet/models/model_200.pth -------------------------------------------------------------------------------- /sync_bn/__init__.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # File : __init__.py 3 | # Author : Jiayuan Mao 4 | # Email : maojiayuan@gmail.com 5 | # Date : 27/01/2018 6 | # 7 | # This file is part of Synchronized-BatchNorm-PyTorch. 8 | # https://github.com/vacancy/Synchronized-BatchNorm-PyTorch 9 | # Distributed under MIT License. 10 | 11 | from .batchnorm import SynchronizedBatchNorm1d, SynchronizedBatchNorm2d, SynchronizedBatchNorm3d 12 | from .replicate import DataParallelWithCallback, patch_replication_callback 13 | -------------------------------------------------------------------------------- /sync_bn/batchnorm.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # File : batchnorm.py 3 | # Author : Jiayuan Mao 4 | # Email : maojiayuan@gmail.com 5 | # Date : 27/01/2018 6 | # 7 | # This file is part of Synchronized-BatchNorm-PyTorch. 8 | # https://github.com/vacancy/Synchronized-BatchNorm-PyTorch 9 | # Distributed under MIT License. 10 | 11 | import collections 12 | 13 | import torch 14 | import torch.nn.functional as F 15 | 16 | from torch.nn.modules.batchnorm import _BatchNorm 17 | from torch.nn.parallel._functions import ReduceAddCoalesced, Broadcast 18 | 19 | from .comm import SyncMaster 20 | 21 | __all__ = ['SynchronizedBatchNorm1d', 'SynchronizedBatchNorm2d', 'SynchronizedBatchNorm3d'] 22 | 23 | 24 | def _sum_ft(tensor): 25 | """sum over the first and last dimention""" 26 | return tensor.sum(dim=0).sum(dim=-1) 27 | 28 | 29 | def _unsqueeze_ft(tensor): 30 | """add new dementions at the front and the tail""" 31 | return tensor.unsqueeze(0).unsqueeze(-1) 32 | 33 | 34 | _ChildMessage = collections.namedtuple('_ChildMessage', ['sum', 'ssum', 'sum_size']) 35 | _MasterMessage = collections.namedtuple('_MasterMessage', ['sum', 'inv_std']) 36 | 37 | 38 | class _SynchronizedBatchNorm(_BatchNorm): 39 | def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True): 40 | super(_SynchronizedBatchNorm, self).__init__(num_features, eps=eps, momentum=momentum, affine=affine) 41 | 42 | self._sync_master = SyncMaster(self._data_parallel_master) 43 | 44 | self._is_parallel = False 45 | self._parallel_id = None 46 | self._slave_pipe = None 47 | 48 | def forward(self, input): 49 | # If it is not parallel computation or is in evaluation mode, use PyTorch's implementation. 50 | if not (self._is_parallel and self.training): 51 | return F.batch_norm( 52 | input, self.running_mean, self.running_var, self.weight, self.bias, 53 | self.training, self.momentum, self.eps) 54 | 55 | # Resize the input to (B, C, -1). 56 | input_shape = input.size() 57 | input = input.view(input.size(0), self.num_features, -1) 58 | 59 | # Compute the sum and square-sum. 60 | sum_size = input.size(0) * input.size(2) 61 | input_sum = _sum_ft(input) 62 | input_ssum = _sum_ft(input ** 2) 63 | 64 | # Reduce-and-broadcast the statistics. 65 | if self._parallel_id == 0: 66 | mean, inv_std = self._sync_master.run_master(_ChildMessage(input_sum, input_ssum, sum_size)) 67 | else: 68 | mean, inv_std = self._slave_pipe.run_slave(_ChildMessage(input_sum, input_ssum, sum_size)) 69 | 70 | # Compute the output. 71 | if self.affine: 72 | # MJY:: Fuse the multiplication for speed. 73 | output = (input - _unsqueeze_ft(mean)) * _unsqueeze_ft(inv_std * self.weight) + _unsqueeze_ft(self.bias) 74 | else: 75 | output = (input - _unsqueeze_ft(mean)) * _unsqueeze_ft(inv_std) 76 | 77 | # Reshape it. 78 | return output.view(input_shape) 79 | 80 | def __data_parallel_replicate__(self, ctx, copy_id): 81 | self._is_parallel = True 82 | self._parallel_id = copy_id 83 | 84 | # parallel_id == 0 means master device. 85 | if self._parallel_id == 0: 86 | ctx.sync_master = self._sync_master 87 | else: 88 | self._slave_pipe = ctx.sync_master.register_slave(copy_id) 89 | 90 | def _data_parallel_master(self, intermediates): 91 | """Reduce the sum and square-sum, compute the statistics, and broadcast it.""" 92 | 93 | # Always using same "device order" makes the ReduceAdd operation faster. 94 | # Thanks to:: Tete Xiao (http://tetexiao.com/) 95 | intermediates = sorted(intermediates, key=lambda i: i[1].sum.get_device()) 96 | 97 | to_reduce = [i[1][:2] for i in intermediates] 98 | to_reduce = [j for i in to_reduce for j in i] # flatten 99 | target_gpus = [i[1].sum.get_device() for i in intermediates] 100 | 101 | sum_size = sum([i[1].sum_size for i in intermediates]) 102 | sum_, ssum = ReduceAddCoalesced.apply(target_gpus[0], 2, *to_reduce) 103 | mean, inv_std = self._compute_mean_std(sum_, ssum, sum_size) 104 | 105 | broadcasted = Broadcast.apply(target_gpus, mean, inv_std) 106 | 107 | outputs = [] 108 | for i, rec in enumerate(intermediates): 109 | outputs.append((rec[0], _MasterMessage(*broadcasted[i*2:i*2+2]))) 110 | 111 | return outputs 112 | 113 | def _compute_mean_std(self, sum_, ssum, size): 114 | """Compute the mean and standard-deviation with sum and square-sum. This method 115 | also maintains the moving average on the master device.""" 116 | assert size > 1, 'BatchNorm computes unbiased standard-deviation, which requires size > 1.' 117 | mean = sum_ / size 118 | sumvar = ssum - sum_ * mean 119 | unbias_var = sumvar / (size - 1) 120 | bias_var = sumvar / size 121 | 122 | self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * mean.data 123 | self.running_var = (1 - self.momentum) * self.running_var + self.momentum * unbias_var.data 124 | 125 | return mean, bias_var.clamp(self.eps) ** -0.5 126 | 127 | 128 | class SynchronizedBatchNorm1d(_SynchronizedBatchNorm): 129 | r"""Applies Synchronized Batch Normalization over a 2d or 3d input that is seen as a 130 | mini-batch. 131 | 132 | .. math:: 133 | 134 | y = \frac{x - mean[x]}{ \sqrt{Var[x] + \epsilon}} * gamma + beta 135 | 136 | This module differs from the built-in PyTorch BatchNorm1d as the mean and 137 | standard-deviation are reduced across all devices during training. 138 | 139 | For example, when one uses `nn.DataParallel` to wrap the network during 140 | training, PyTorch's implementation normalize the tensor on each device using 141 | the statistics only on that device, which accelerated the computation and 142 | is also easy to implement, but the statistics might be inaccurate. 143 | Instead, in this synchronized version, the statistics will be computed 144 | over all training samples distributed on multiple devices. 145 | 146 | Note that, for one-GPU or CPU-only case, this module behaves exactly same 147 | as the built-in PyTorch implementation. 148 | 149 | The mean and standard-deviation are calculated per-dimension over 150 | the mini-batches and gamma and beta are learnable parameter vectors 151 | of size C (where C is the input size). 152 | 153 | During training, this layer keeps a running estimate of its computed mean 154 | and variance. The running sum is kept with a default momentum of 0.1. 155 | 156 | During evaluation, this running mean/variance is used for normalization. 157 | 158 | Because the BatchNorm is done over the `C` dimension, computing statistics 159 | on `(N, L)` slices, it's common terminology to call this Temporal BatchNorm 160 | 161 | Args: 162 | num_features: num_features from an expected input of size 163 | `batch_size x num_features [x width]` 164 | eps: a value added to the denominator for numerical stability. 165 | Default: 1e-5 166 | momentum: the value used for the running_mean and running_var 167 | computation. Default: 0.1 168 | affine: a boolean value that when set to ``True``, gives the layer learnable 169 | affine parameters. Default: ``True`` 170 | 171 | Shape: 172 | - Input: :math:`(N, C)` or :math:`(N, C, L)` 173 | - Output: :math:`(N, C)` or :math:`(N, C, L)` (same shape as input) 174 | 175 | Examples: 176 | >>> # With Learnable Parameters 177 | >>> m = SynchronizedBatchNorm1d(100) 178 | >>> # Without Learnable Parameters 179 | >>> m = SynchronizedBatchNorm1d(100, affine=False) 180 | >>> input = torch.autograd.Variable(torch.randn(20, 100)) 181 | >>> output = m(input) 182 | """ 183 | 184 | def _check_input_dim(self, input): 185 | if input.dim() != 2 and input.dim() != 3: 186 | raise ValueError('expected 2D or 3D input (got {}D input)' 187 | .format(input.dim())) 188 | super(SynchronizedBatchNorm1d, self)._check_input_dim(input) 189 | 190 | 191 | class SynchronizedBatchNorm2d(_SynchronizedBatchNorm): 192 | r"""Applies Batch Normalization over a 4d input that is seen as a mini-batch 193 | of 3d inputs 194 | 195 | .. math:: 196 | 197 | y = \frac{x - mean[x]}{ \sqrt{Var[x] + \epsilon}} * gamma + beta 198 | 199 | This module differs from the built-in PyTorch BatchNorm2d as the mean and 200 | standard-deviation are reduced across all devices during training. 201 | 202 | For example, when one uses `nn.DataParallel` to wrap the network during 203 | training, PyTorch's implementation normalize the tensor on each device using 204 | the statistics only on that device, which accelerated the computation and 205 | is also easy to implement, but the statistics might be inaccurate. 206 | Instead, in this synchronized version, the statistics will be computed 207 | over all training samples distributed on multiple devices. 208 | 209 | Note that, for one-GPU or CPU-only case, this module behaves exactly same 210 | as the built-in PyTorch implementation. 211 | 212 | The mean and standard-deviation are calculated per-dimension over 213 | the mini-batches and gamma and beta are learnable parameter vectors 214 | of size C (where C is the input size). 215 | 216 | During training, this layer keeps a running estimate of its computed mean 217 | and variance. The running sum is kept with a default momentum of 0.1. 218 | 219 | During evaluation, this running mean/variance is used for normalization. 220 | 221 | Because the BatchNorm is done over the `C` dimension, computing statistics 222 | on `(N, H, W)` slices, it's common terminology to call this Spatial BatchNorm 223 | 224 | Args: 225 | num_features: num_features from an expected input of 226 | size batch_size x num_features x height x width 227 | eps: a value added to the denominator for numerical stability. 228 | Default: 1e-5 229 | momentum: the value used for the running_mean and running_var 230 | computation. Default: 0.1 231 | affine: a boolean value that when set to ``True``, gives the layer learnable 232 | affine parameters. Default: ``True`` 233 | 234 | Shape: 235 | - Input: :math:`(N, C, H, W)` 236 | - Output: :math:`(N, C, H, W)` (same shape as input) 237 | 238 | Examples: 239 | >>> # With Learnable Parameters 240 | >>> m = SynchronizedBatchNorm2d(100) 241 | >>> # Without Learnable Parameters 242 | >>> m = SynchronizedBatchNorm2d(100, affine=False) 243 | >>> input = torch.autograd.Variable(torch.randn(20, 100, 35, 45)) 244 | >>> output = m(input) 245 | """ 246 | 247 | def _check_input_dim(self, input): 248 | if input.dim() != 4: 249 | raise ValueError('expected 4D input (got {}D input)' 250 | .format(input.dim())) 251 | super(SynchronizedBatchNorm2d, self)._check_input_dim(input) 252 | 253 | 254 | class SynchronizedBatchNorm3d(_SynchronizedBatchNorm): 255 | r"""Applies Batch Normalization over a 5d input that is seen as a mini-batch 256 | of 4d inputs 257 | 258 | .. math:: 259 | 260 | y = \frac{x - mean[x]}{ \sqrt{Var[x] + \epsilon}} * gamma + beta 261 | 262 | This module differs from the built-in PyTorch BatchNorm3d as the mean and 263 | standard-deviation are reduced across all devices during training. 264 | 265 | For example, when one uses `nn.DataParallel` to wrap the network during 266 | training, PyTorch's implementation normalize the tensor on each device using 267 | the statistics only on that device, which accelerated the computation and 268 | is also easy to implement, but the statistics might be inaccurate. 269 | Instead, in this synchronized version, the statistics will be computed 270 | over all training samples distributed on multiple devices. 271 | 272 | Note that, for one-GPU or CPU-only case, this module behaves exactly same 273 | as the built-in PyTorch implementation. 274 | 275 | The mean and standard-deviation are calculated per-dimension over 276 | the mini-batches and gamma and beta are learnable parameter vectors 277 | of size C (where C is the input size). 278 | 279 | During training, this layer keeps a running estimate of its computed mean 280 | and variance. The running sum is kept with a default momentum of 0.1. 281 | 282 | During evaluation, this running mean/variance is used for normalization. 283 | 284 | Because the BatchNorm is done over the `C` dimension, computing statistics 285 | on `(N, D, H, W)` slices, it's common terminology to call this Volumetric BatchNorm 286 | or Spatio-temporal BatchNorm 287 | 288 | Args: 289 | num_features: num_features from an expected input of 290 | size batch_size x num_features x depth x height x width 291 | eps: a value added to the denominator for numerical stability. 292 | Default: 1e-5 293 | momentum: the value used for the running_mean and running_var 294 | computation. Default: 0.1 295 | affine: a boolean value that when set to ``True``, gives the layer learnable 296 | affine parameters. Default: ``True`` 297 | 298 | Shape: 299 | - Input: :math:`(N, C, D, H, W)` 300 | - Output: :math:`(N, C, D, H, W)` (same shape as input) 301 | 302 | Examples: 303 | >>> # With Learnable Parameters 304 | >>> m = SynchronizedBatchNorm3d(100) 305 | >>> # Without Learnable Parameters 306 | >>> m = SynchronizedBatchNorm3d(100, affine=False) 307 | >>> input = torch.autograd.Variable(torch.randn(20, 100, 35, 45, 10)) 308 | >>> output = m(input) 309 | """ 310 | 311 | def _check_input_dim(self, input): 312 | if input.dim() != 5: 313 | raise ValueError('expected 5D input (got {}D input)' 314 | .format(input.dim())) 315 | super(SynchronizedBatchNorm3d, self)._check_input_dim(input) 316 | -------------------------------------------------------------------------------- /sync_bn/comm.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # File : comm.py 3 | # Author : Jiayuan Mao 4 | # Email : maojiayuan@gmail.com 5 | # Date : 27/01/2018 6 | # 7 | # This file is part of Synchronized-BatchNorm-PyTorch. 8 | # https://github.com/vacancy/Synchronized-BatchNorm-PyTorch 9 | # Distributed under MIT License. 10 | 11 | import queue 12 | import collections 13 | import threading 14 | 15 | __all__ = ['FutureResult', 'SlavePipe', 'SyncMaster'] 16 | 17 | 18 | class FutureResult(object): 19 | """A thread-safe future implementation. Used only as one-to-one pipe.""" 20 | 21 | def __init__(self): 22 | self._result = None 23 | self._lock = threading.Lock() 24 | self._cond = threading.Condition(self._lock) 25 | 26 | def put(self, result): 27 | with self._lock: 28 | assert self._result is None, 'Previous result has\'t been fetched.' 29 | self._result = result 30 | self._cond.notify() 31 | 32 | def get(self): 33 | with self._lock: 34 | if self._result is None: 35 | self._cond.wait() 36 | 37 | res = self._result 38 | self._result = None 39 | return res 40 | 41 | 42 | _MasterRegistry = collections.namedtuple('MasterRegistry', ['result']) 43 | _SlavePipeBase = collections.namedtuple('_SlavePipeBase', ['identifier', 'queue', 'result']) 44 | 45 | 46 | class SlavePipe(_SlavePipeBase): 47 | """Pipe for master-slave communication.""" 48 | 49 | def run_slave(self, msg): 50 | self.queue.put((self.identifier, msg)) 51 | ret = self.result.get() 52 | self.queue.put(True) 53 | return ret 54 | 55 | 56 | class SyncMaster(object): 57 | """An abstract `SyncMaster` object. 58 | 59 | - During the replication, as the data parallel will trigger an callback of each module, all slave devices should 60 | call `register(id)` and obtain an `SlavePipe` to communicate with the master. 61 | - During the forward pass, master device invokes `run_master`, all messages from slave devices will be collected, 62 | and passed to a registered callback. 63 | - After receiving the messages, the master device should gather the information and determine to message passed 64 | back to each slave devices. 65 | """ 66 | 67 | def __init__(self, master_callback): 68 | """ 69 | 70 | Args: 71 | master_callback: a callback to be invoked after having collected messages from slave devices. 72 | """ 73 | self._master_callback = master_callback 74 | self._queue = queue.Queue() 75 | self._registry = collections.OrderedDict() 76 | self._activated = False 77 | 78 | def __getstate__(self): 79 | return {'master_callback': self._master_callback} 80 | 81 | def __setstate__(self, state): 82 | self.__init__(state['master_callback']) 83 | 84 | def register_slave(self, identifier): 85 | """ 86 | Register an slave device. 87 | 88 | Args: 89 | identifier: an identifier, usually is the device id. 90 | 91 | Returns: a `SlavePipe` object which can be used to communicate with the master device. 92 | 93 | """ 94 | if self._activated: 95 | assert self._queue.empty(), 'Queue is not clean before next initialization.' 96 | self._activated = False 97 | self._registry.clear() 98 | future = FutureResult() 99 | self._registry[identifier] = _MasterRegistry(future) 100 | return SlavePipe(identifier, self._queue, future) 101 | 102 | def run_master(self, master_msg): 103 | """ 104 | Main entry for the master device in each forward pass. 105 | The messages were first collected from each devices (including the master device), and then 106 | an callback will be invoked to compute the message to be sent back to each devices 107 | (including the master device). 108 | 109 | Args: 110 | master_msg: the message that the master want to send to itself. This will be placed as the first 111 | message when calling `master_callback`. For detailed usage, see `_SynchronizedBatchNorm` for an example. 112 | 113 | Returns: the message to be sent back to the master device. 114 | 115 | """ 116 | self._activated = True 117 | 118 | intermediates = [(0, master_msg)] 119 | for i in range(self.nr_slaves): 120 | intermediates.append(self._queue.get()) 121 | 122 | results = self._master_callback(intermediates) 123 | assert results[0][0] == 0, 'The first result should belongs to the master.' 124 | 125 | for i, res in results: 126 | if i == 0: 127 | continue 128 | self._registry[i].result.put(res) 129 | 130 | for i in range(self.nr_slaves): 131 | assert self._queue.get() is True 132 | 133 | return results[0][1] 134 | 135 | @property 136 | def nr_slaves(self): 137 | return len(self._registry) 138 | -------------------------------------------------------------------------------- /sync_bn/replicate.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # File : replicate.py 3 | # Author : Jiayuan Mao 4 | # Email : maojiayuan@gmail.com 5 | # Date : 27/01/2018 6 | # 7 | # This file is part of Synchronized-BatchNorm-PyTorch. 8 | # https://github.com/vacancy/Synchronized-BatchNorm-PyTorch 9 | # Distributed under MIT License. 10 | 11 | import functools 12 | 13 | from torch.nn.parallel.data_parallel import DataParallel 14 | 15 | __all__ = [ 16 | 'CallbackContext', 17 | 'execute_replication_callbacks', 18 | 'DataParallelWithCallback', 19 | 'patch_replication_callback' 20 | ] 21 | 22 | 23 | class CallbackContext(object): 24 | pass 25 | 26 | 27 | def execute_replication_callbacks(modules): 28 | """ 29 | Execute an replication callback `__data_parallel_replicate__` on each module created by original replication. 30 | 31 | The callback will be invoked with arguments `__data_parallel_replicate__(ctx, copy_id)` 32 | 33 | Note that, as all modules are isomorphism, we assign each sub-module with a context 34 | (shared among multiple copies of this module on different devices). 35 | Through this context, different copies can share some information. 36 | 37 | We guarantee that the callback on the master copy (the first copy) will be called ahead of calling the callback 38 | of any slave copies. 39 | """ 40 | master_copy = modules[0] 41 | nr_modules = len(list(master_copy.modules())) 42 | ctxs = [CallbackContext() for _ in range(nr_modules)] 43 | 44 | for i, module in enumerate(modules): 45 | for j, m in enumerate(module.modules()): 46 | if hasattr(m, '__data_parallel_replicate__'): 47 | m.__data_parallel_replicate__(ctxs[j], i) 48 | 49 | 50 | class DataParallelWithCallback(DataParallel): 51 | """ 52 | Data Parallel with a replication callback. 53 | 54 | An replication callback `__data_parallel_replicate__` of each module will be invoked after being created by 55 | original `replicate` functions. 56 | The callback will be invoked with arguments `__data_parallel_replicate__(ctx, copy_id)` 57 | 58 | Examples: 59 | > sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False) 60 | > sync_bn = DataParallelWithCallback(sync_bn, device_ids=[0, 1]) 61 | # sync_bn.__data_parallel_replicate__ will be invoked. 62 | """ 63 | 64 | def replicate(self, module, device_ids): 65 | modules = super(DataParallelWithCallback, self).replicate(module, device_ids) 66 | execute_replication_callbacks(modules) 67 | return modules 68 | 69 | 70 | def patch_replication_callback(data_parallel): 71 | """ 72 | Monkey-patch an existing `DataParallel` object. Add the replication callback. 73 | Useful when you have customized `DataParallel` implementation. 74 | 75 | Examples: 76 | > sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False) 77 | > sync_bn = DataParallel(sync_bn, device_ids=[0, 1]) 78 | > patch_replication_callback(sync_bn) 79 | # this is equivalent to 80 | > sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False) 81 | > sync_bn = DataParallelWithCallback(sync_bn, device_ids=[0, 1]) 82 | """ 83 | 84 | assert isinstance(data_parallel, DataParallel) 85 | 86 | old_replicate = data_parallel.replicate 87 | 88 | @functools.wraps(old_replicate) 89 | def new_replicate(module, device_ids): 90 | modules = old_replicate(module, device_ids) 91 | execute_replication_callbacks(modules) 92 | return modules 93 | 94 | data_parallel.replicate = new_replicate 95 | -------------------------------------------------------------------------------- /sync_bn/unittest.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # File : unittest.py 3 | # Author : Jiayuan Mao 4 | # Email : maojiayuan@gmail.com 5 | # Date : 27/01/2018 6 | # 7 | # This file is part of Synchronized-BatchNorm-PyTorch. 8 | # https://github.com/vacancy/Synchronized-BatchNorm-PyTorch 9 | # Distributed under MIT License. 10 | 11 | import unittest 12 | 13 | import numpy as np 14 | from torch.autograd import Variable 15 | 16 | 17 | def as_numpy(v): 18 | if isinstance(v, Variable): 19 | v = v.data 20 | return v.cpu().numpy() 21 | 22 | 23 | class TorchTestCase(unittest.TestCase): 24 | def assertTensorClose(self, a, b, atol=1e-3, rtol=1e-3): 25 | npa, npb = as_numpy(a), as_numpy(b) 26 | self.assertTrue( 27 | np.allclose(npa, npb, atol=atol), 28 | 'Tensor close check failed\n{}\n{}\nadiff={}, rdiff={}'.format(a, b, np.abs(npa - npb).max(), np.abs((npa - npb) / np.fmax(npa, 1e-5)).max()) 29 | ) 30 | -------------------------------------------------------------------------------- /util.py: -------------------------------------------------------------------------------- 1 | import os 2 | import glob 3 | import copy 4 | import random 5 | import pickle 6 | import numpy as np 7 | from plyfile import PlyData 8 | 9 | import torch 10 | from torch import nn 11 | from torch.nn.modules.conv import _ConvNd 12 | from torch.nn.modules.batchnorm import _BatchNorm 13 | import torch.nn.init as initer 14 | import torch.nn.functional as F 15 | 16 | 17 | class AverageMeter(object): 18 | """Computes and stores the average and current value""" 19 | 20 | def __init__(self): 21 | self.reset() 22 | 23 | def reset(self): 24 | self.val = 0 25 | self.avg = 0 26 | self.sum = 0 27 | self.count = 0 28 | 29 | def update(self, val, n=1): 30 | self.val = val 31 | self.sum += val * n 32 | self.count += n 33 | self.avg = self.sum / self.count 34 | 35 | 36 | def set_seed(seed=1): 37 | print('Using random seed', seed) 38 | random.seed(seed) 39 | np.random.seed(seed) 40 | torch.manual_seed(seed) 41 | torch.cuda.manual_seed_all(seed) 42 | 43 | 44 | def str2bool(v): 45 | return v.lower() in ("yes", "true", "t", "1") 46 | 47 | 48 | def get_lr(optimizer): 49 | return optimizer.param_groups[0]['lr'] 50 | 51 | 52 | def adjust_lr(optimizer, new_lr): 53 | for param_group in optimizer.param_groups: 54 | param_group['lr'] = new_lr 55 | 56 | 57 | def weights_init(m): 58 | classname = m.__class__.__name__ 59 | if classname.find('Conv2d') != -1: 60 | nn.init.xavier_normal_(m.weight.data) 61 | try: 62 | nn.init.constant_(m.bias.data, 0.0) 63 | except AttributeError: 64 | pass 65 | elif classname.find('Linear') != -1: 66 | nn.init.xavier_normal_(m.weight.data) 67 | try: 68 | nn.init.constant_(m.bias.data, 0.0) 69 | except AttributeError: 70 | pass 71 | 72 | 73 | def bn_momentum_adjust(m, momentum): 74 | if isinstance(m, nn.BatchNorm2d) or \ 75 | isinstance(m, nn.BatchNorm1d): 76 | m.momentum = momentum 77 | 78 | 79 | def intersectionAndUnion(output, target, K, ignore_index=255): 80 | # 'K' classes, output and target sizes are N or N * L or N * H * W, each value in range 0 to K - 1. 81 | assert (output.ndim in [1, 2, 3]) 82 | assert output.shape == target.shape 83 | output = output.reshape(output.size).copy() 84 | target = target.reshape(target.size) 85 | output[np.where(target == ignore_index)[0]] = 255 86 | target[np.where(target == ignore_index)[0]] = 255 87 | intersection = output[np.where(output == target)[0]] 88 | area_intersection, _ = np.histogram(intersection, bins=np.arange(K+1)) 89 | area_output, _ = np.histogram(output, bins=np.arange(K+1)) 90 | area_target, _ = np.histogram(target, bins=np.arange(K+1)) 91 | area_union = area_output + area_target - area_intersection 92 | return area_intersection, area_union, area_target 93 | 94 | 95 | def calc_victim_value(class_value, label, victim_class): 96 | values = [] 97 | for lbl in victim_class: 98 | if label is None or (label == lbl).any(): 99 | values.append(class_value[lbl]) 100 | return np.mean(values) 101 | 102 | 103 | def check_makedirs(dir_name): 104 | if not os.path.exists(dir_name): 105 | os.makedirs(dir_name) 106 | 107 | 108 | def init_weights(model, conv='kaiming', batchnorm='normal', linear='kaiming', lstm='kaiming'): 109 | """ 110 | :param model: Pytorch Model which is nn.Module 111 | :param conv: 'kaiming' or 'xavier' 112 | :param batchnorm: 'normal' or 'constant' 113 | :param linear: 'kaiming' or 'xavier' 114 | :param lstm: 'kaiming' or 'xavier' 115 | """ 116 | for m in model.modules(): 117 | if isinstance(m, (_ConvNd)): 118 | if conv == 'kaiming': 119 | initer.kaiming_normal_(m.weight) 120 | elif conv == 'xavier': 121 | initer.xavier_normal_(m.weight) 122 | else: 123 | raise ValueError("init type of conv error.\n") 124 | if m.bias is not None: 125 | initer.constant_(m.bias, 0) 126 | 127 | elif isinstance(m, _BatchNorm): 128 | if batchnorm == 'normal': 129 | initer.normal_(m.weight, 1.0, 0.02) 130 | elif batchnorm == 'constant': 131 | initer.constant_(m.weight, 1.0) 132 | else: 133 | raise ValueError("init type of batchnorm error.\n") 134 | initer.constant_(m.bias, 0.0) 135 | 136 | elif isinstance(m, nn.Linear): 137 | if linear == 'kaiming': 138 | initer.kaiming_normal_(m.weight) 139 | elif linear == 'xavier': 140 | initer.xavier_normal_(m.weight) 141 | else: 142 | raise ValueError("init type of linear error.\n") 143 | if m.bias is not None: 144 | initer.constant_(m.bias, 0) 145 | 146 | elif isinstance(m, nn.LSTM): 147 | for name, param in m.named_parameters(): 148 | if 'weight' in name: 149 | if lstm == 'kaiming': 150 | initer.kaiming_normal_(param) 151 | elif lstm == 'xavier': 152 | initer.xavier_normal_(param) 153 | else: 154 | raise ValueError("init type of lstm error.\n") 155 | elif 'bias' in name: 156 | initer.constant_(param, 0) 157 | 158 | 159 | def convert_to_syncbn(model): 160 | def recursive_set(cur_module, name, module): 161 | if len(name.split('.')) > 1: 162 | recursive_set( 163 | getattr(cur_module, name[:name.find('.')]), name[name.find('.')+1:], module) 164 | else: 165 | setattr(cur_module, name, module) 166 | from sync_bn import SynchronizedBatchNorm1d, SynchronizedBatchNorm2d, \ 167 | SynchronizedBatchNorm3d 168 | for name, m in model.named_modules(): 169 | if isinstance(m, nn.BatchNorm1d): 170 | recursive_set(model, name, SynchronizedBatchNorm1d( 171 | m.num_features, m.eps, m.momentum, m.affine)) 172 | elif isinstance(m, nn.BatchNorm2d): 173 | recursive_set(model, name, SynchronizedBatchNorm2d( 174 | m.num_features, m.eps, m.momentum, m.affine)) 175 | elif isinstance(m, nn.BatchNorm3d): 176 | recursive_set(model, name, SynchronizedBatchNorm3d( 177 | m.num_features, m.eps, m.momentum, m.affine)) 178 | 179 | 180 | def lbl2rgb(label, names): 181 | """Convert label to rgb colors. 182 | label: [N] 183 | """ 184 | from config import NAME2COLOR 185 | if len(names) == 13: 186 | colors = NAME2COLOR['S3DIS'] 187 | else: 188 | colors = NAME2COLOR['ScanNet'] 189 | rgb = np.zeros((label.shape[0], 3)) 190 | uni_lbl = np.unique(label).astype(np.uint8) 191 | for lbl in uni_lbl: 192 | mask = (label == lbl) 193 | rgb[mask] = np.tile(np.array( 194 | colors[names[lbl]])[None, :], (mask.sum(), 1)) 195 | return rgb 196 | 197 | 198 | def convert2vis(xyz, label, names): 199 | """Assign color to each point according to label.""" 200 | rgb = lbl2rgb(label, names) * 255. 201 | data = np.concatenate([xyz, rgb], axis=1) 202 | return data 203 | 204 | 205 | def proc_pert(points, gt, pred, folder, 206 | names, part=False, ignore_label=255): 207 | """Process and save files for visulization in perturbation attack.""" 208 | check_makedirs(folder) 209 | lbl2cls = {i: names[i] for i in range(len(names))} 210 | 211 | np.savetxt(os.path.join(folder, 'all_points.txt'), points, delimiter=';') 212 | gt_seg = convert2vis(points[gt != ignore_label, :3], 213 | gt[gt != ignore_label], names) 214 | pred_seg = convert2vis(points[gt != ignore_label, :3], 215 | pred[gt != ignore_label], names) 216 | np.savetxt(os.path.join(folder, 'gt.txt'), 217 | gt_seg, delimiter=';') 218 | np.savetxt(os.path.join(folder, 'pred.txt'), 219 | pred_seg, delimiter=';') 220 | if part: 221 | uni_lbl = np.unique(gt[gt != ignore_label]).astype(np.uint8) 222 | for lbl in uni_lbl: 223 | lbl = int(lbl) 224 | mask = (gt == lbl) 225 | sel_points = points[mask] 226 | mask = (gt[gt != ignore_label] == lbl) 227 | sel_seg = pred_seg[mask] 228 | np.savetxt( 229 | os.path.join(folder, '{}_{}_points.txt'.format( 230 | lbl, lbl2cls[lbl])), 231 | sel_points, delimiter=';') 232 | np.savetxt( 233 | os.path.join(folder, '{}_{}_pred.txt'.format( 234 | lbl, lbl2cls[lbl])), 235 | sel_seg, delimiter=';') 236 | 237 | 238 | def proc_add(points, noise, gt, pred, noise_pred, folder, 239 | names, part=False, ignore_label=255): 240 | """Process and save files for visulization in adding attack.""" 241 | check_makedirs(folder) 242 | lbl2cls = {i: names[i] for i in range(len(names))} 243 | 244 | np.savetxt(os.path.join(folder, 'all_points.txt'), points, delimiter=';') 245 | np.savetxt(os.path.join(folder, 'noise_points.txt'), noise, delimiter=';') 246 | gt_seg = convert2vis(points[gt != ignore_label, :3], 247 | gt[gt != ignore_label], names) 248 | pred_seg = convert2vis(points[gt != ignore_label, :3], 249 | pred[gt != ignore_label], names) 250 | noise_seg = convert2vis(noise[:, :3], noise_pred, names) 251 | np.savetxt(os.path.join(folder, 'gt.txt'), 252 | gt_seg, delimiter=';') 253 | np.savetxt(os.path.join(folder, 'pred.txt'), 254 | pred_seg, delimiter=';') 255 | np.savetxt(os.path.join(folder, 'noise_pred.txt'), 256 | noise_seg, delimiter=';') 257 | if part: 258 | uni_lbl = np.unique(gt[gt != ignore_label]).astype(np.uint8) 259 | for lbl in uni_lbl: 260 | lbl = int(lbl) 261 | mask = (gt == lbl) 262 | sel_points = points[mask] 263 | mask = (gt[gt != ignore_label] == lbl) 264 | sel_seg = pred_seg[mask] 265 | np.savetxt( 266 | os.path.join(folder, '{}_{}_points.txt'.format( 267 | lbl, lbl2cls[lbl])), 268 | sel_points, delimiter=';') 269 | np.savetxt( 270 | os.path.join(folder, '{}_{}_pred.txt'.format( 271 | lbl, lbl2cls[lbl])), 272 | sel_seg, delimiter=';') 273 | 274 | 275 | def save_vis(pred_root, save_root, data_root): 276 | from config import CLASS_NAMES 277 | if 'S3DIS' in data_root: # save Area5 data 278 | names = CLASS_NAMES['S3DIS']['other'] 279 | gt_save = load_pickle( 280 | os.path.join(pred_root, 'gt_5.pickle'))['gt'] 281 | pred_save = load_pickle( 282 | os.path.join(pred_root, 'pred_5.pickle'))['pred'] 283 | assert len(gt_save) == len(pred_save) 284 | all_rooms = sorted(os.listdir(data_root)) 285 | all_rooms = [ 286 | room for room in all_rooms if 'Area_5' in room 287 | ] 288 | assert len(gt_save) == len(all_rooms) 289 | check_makedirs(save_root) 290 | for i, room in enumerate(all_rooms): 291 | points = np.load(os.path.join(data_root, room))[:, :6] 292 | folder = os.path.join(save_root, room[:-4]) 293 | check_makedirs(folder) 294 | proc_pert(points, gt_save[i], pred_save[i], 295 | folder, names, part=True) 296 | elif 'ScanNet' in data_root: # save val set data 297 | names = CLASS_NAMES['ScanNet']['other'] 298 | gt_save = load_pickle( 299 | os.path.join(pred_root, 'gt_val.pickle'))['gt'] 300 | pred_save = load_pickle( 301 | os.path.join(pred_root, 'pred_val.pickle'))['pred'] 302 | assert len(gt_save) == len(pred_save) 303 | data_file = os.path.join( 304 | data_root, 'scannet_val_rgb21c_pointid.pickle') 305 | file_pickle = open(data_file, 'rb') 306 | xyz_all = pickle.load(file_pickle) 307 | file_pickle.close() 308 | assert len(xyz_all) == len(gt_save) 309 | with open(os.path.join( 310 | data_root, 'meta_data/scannetv2_val.txt')) as fl: 311 | scene_id = fl.read().splitlines() 312 | assert len(scene_id) == len(gt_save) 313 | check_makedirs(save_root) 314 | for i in range(len(gt_save)): 315 | points = xyz_all[i][:, :6] 316 | folder = os.path.join(save_root, scene_id[i]) 317 | check_makedirs(folder) 318 | proc_pert(points, gt_save[i], pred_save[i], 319 | folder, names, part=True) 320 | 321 | 322 | def save_vis_mink(pred_root, save_root, data_root): 323 | from config import CLASS_NAMES 324 | 325 | def load_data(file_name): 326 | plydata = PlyData.read(file_name) 327 | data = plydata.elements[0].data 328 | coords = np.array([data['x'], data['y'], data['z']], 329 | dtype=np.float32).T 330 | colors = np.array([data['red'], data['green'], 331 | data['blue']], dtype=np.float32).T 332 | return np.concatenate([coords, colors], axis=1) 333 | 334 | if 'S3DIS' in data_root: # save Area5 data 335 | names = CLASS_NAMES['S3DIS']['mink'] 336 | gt_save = load_pickle( 337 | os.path.join(pred_root, 'gt_5.pickle'))['gt'] 338 | pred_save = load_pickle( 339 | os.path.join(pred_root, 'pred_5.pickle'))['pred'] 340 | assert len(gt_save) == len(pred_save) 341 | data_root = os.path.join(data_root, 'Area_5') 342 | all_rooms = sorted(os.listdir(data_root)) 343 | assert len(all_rooms) == len(gt_save) 344 | check_makedirs(save_root) 345 | 346 | for i, room in enumerate(all_rooms): 347 | data = os.path.join(data_root, room) 348 | points = load_data(data) 349 | folder = os.path.join( 350 | save_root, 'Area_5_{}'.format(room[:-4])) 351 | check_makedirs(folder) 352 | proc_pert(points, gt_save[i], pred_save[i], 353 | folder, names, part=True) 354 | elif 'ScanNet' in data_root: # save val set 355 | names = CLASS_NAMES['ScanNet']['mink'] 356 | gt_save = load_pickle( 357 | os.path.join(pred_root, 'gt_val.pickle'))['gt'] 358 | pred_save = load_pickle( 359 | os.path.join(pred_root, 'pred_val.pickle'))['pred'] 360 | assert len(gt_save) == len(pred_save) 361 | data_root = os.path.join(data_root, 'train') 362 | with open(os.path.join( 363 | data_root, 'scannetv2_val.txt'), 'r') as f: 364 | all_rooms = f.readlines() 365 | all_rooms = [room[:-1] for room in all_rooms] 366 | assert len(all_rooms) == len(gt_save) 367 | check_makedirs(save_root) 368 | 369 | for i, room in enumerate(all_rooms): 370 | data = os.path.join(data_root, room) 371 | points = load_data(data) 372 | folder = os.path.join(save_root, room[:-4]) 373 | check_makedirs(folder) 374 | proc_pert(points, gt_save[i], pred_save[i], 375 | folder, names, part=True) 376 | 377 | 378 | def save_vis_from_pickle(pkl_root, save_root=None, room_idx=52, 379 | room_name='scene0354_00'): 380 | names = [ 381 | 'wall', 'floor', 'cabinet', 'bed', 'chair', 'sofa', 'table', 382 | 'door', 'window', 'bookshelf', 'picture', 'counter', 'desk', 383 | 'curtain', 'refrigerator', 'showercurtain', 'toilet', 'sink', 384 | 'bathtub', 'otherfurniture' 385 | ] 386 | data = load_pickle(pkl_root) 387 | points = data['data'][room_idx] 388 | pred = data['pred'][room_idx] 389 | gt = data['gt'][room_idx] 390 | if save_root is None: 391 | save_root = os.path.dirname(pkl_root) 392 | save_folder = os.path.join(save_root, room_name) 393 | proc_pert(points, gt, pred, save_folder, names, part=True) 394 | 395 | 396 | def save_pickle(filename, dict_data): 397 | with open(filename, 'wb') as handle: 398 | pickle.dump(dict_data, handle, 399 | protocol=pickle.HIGHEST_PROTOCOL) 400 | 401 | 402 | def load_pickle(filename): 403 | with open(filename, 'rb') as f: 404 | data = pickle.load(f) 405 | return data 406 | 407 | 408 | def load_s3dis_instance(folder, name2cls, load_name=['chair']): 409 | """Load S3DIS room in a Inst Seg format. 410 | Get each instance separately. 411 | 412 | If load_name is None or [], return all instances. 413 | Returns a list of [np.array of [N, 6], label] 414 | """ 415 | cls2name = {name2cls[name]: name for name in name2cls.keys()} 416 | anno_path = os.path.join(folder, 'Annotations') 417 | points_list = [] 418 | labels_list = [] 419 | idx = 0 420 | files = glob.glob(os.path.join(anno_path, '*.txt')) 421 | files.sort() 422 | 423 | for f in files: 424 | cls = os.path.basename(f).split('_')[0] 425 | if cls not in name2cls.keys(): 426 | cls = 'clutter' 427 | points = np.loadtxt(f) # [N, 6] 428 | num = points.shape[0] 429 | points_list.append(points) 430 | labels_list.append((idx, idx + num, name2cls[cls])) 431 | idx += num 432 | 433 | # normalize points coords by minus min 434 | data = np.concatenate(points_list, 0) 435 | xyz_min = np.amin(data, axis=0)[0:3] 436 | data[:, 0:3] -= xyz_min 437 | 438 | # rearrange to separate instances 439 | if load_name is None or not load_name: 440 | load_name = list(name2cls.keys()) 441 | instances = [ 442 | [data[pair[0]:pair[1]], pair[2]] for pair in labels_list if 443 | cls2name[pair[2]] in load_name 444 | ] 445 | return instances 446 | 447 | 448 | def cal_loss(pred, gold, smoothing=False, ignore_index=255): 449 | ''' Calculate cross entropy loss, apply label smoothing if needed. ''' 450 | 451 | gold = gold.contiguous().view(-1) 452 | 453 | if smoothing: 454 | eps = 0.2 455 | n_class = pred.size(1) 456 | 457 | one_hot = torch.zeros_like(pred).scatter(1, gold.view(-1, 1), 1) 458 | one_hot = one_hot * (1 - eps) + (1 - one_hot) * eps / (n_class - 1) 459 | log_prb = F.log_softmax(pred, dim=1) 460 | 461 | loss = -(one_hot * log_prb).sum(dim=1).mean() 462 | else: 463 | loss = F.cross_entropy( 464 | pred, gold, reduction='mean', 465 | ignore_index=ignore_index) 466 | 467 | return loss 468 | 469 | 470 | class IOStream(): 471 | def __init__(self, path): 472 | self.f = open(path, 'a') 473 | 474 | def cprint(self, text): 475 | print(text) 476 | self.f.write(text+'\n') 477 | self.f.flush() 478 | 479 | def close(self): 480 | self.f.close() --------------------------------------------------------------------------------