├── LICENSE ├── README.md ├── data_conversions ├── classes_ObjClassification-ShapeNetCore55.txt ├── download_datasets.py ├── download_semantic3d.sh ├── extract_scannet_objs.py ├── prepare_partseg_data.py ├── prepare_s3dis_data.py ├── prepare_s3dis_filelists.py ├── prepare_s3dis_label.py ├── prepare_scannet_cls_data.py ├── prepare_scannet_seg_data.py ├── prepare_scannet_seg_filelists.py ├── prepare_semantic3d_data.py ├── prepare_semantic3d_filelists.py ├── preprocess_semantic3d_data.py ├── scannet-labels.combined.tsv └── un7z_semantic3d.sh ├── evaluation ├── eval_s3dis.py ├── eval_scannet.py ├── eval_shapenet_seg.py ├── s3dis_merge.py └── semantic3d_merge.py ├── log ├── cls │ └── .gitignore └── seg │ └── .gitignore ├── setting ├── cls_modelnet40.py ├── seg_s3dis.py ├── seg_scannet.py ├── seg_semantic3d.py └── seg_shapenet.py ├── shellconv.py ├── test_cls_modelnet40.py ├── test_seg_3dis.py ├── test_seg_scannet.py ├── test_seg_semantic3d.py ├── test_seg_shapenet.py ├── tf_ops └── sampling │ ├── LICENSE │ ├── __init__.py │ ├── tf_sampling.cpp │ ├── tf_sampling.py │ ├── tf_sampling_compile.sh │ └── tf_sampling_g.cu ├── train_val_cls.py ├── train_val_seg.py └── utils ├── __init__.py ├── data └── .gitignore ├── data_utils.py ├── plyfile.py ├── pointfly.py └── provider.py /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | ShellNet 4 | Copyright (c) 2019 Zhiyuan Zhang 5 | 6 | Permission is hereby granted, free of charge, to any person obtaining a copy 7 | of this software and associated documentation files (the "Software"), to deal 8 | in the Software without restriction, including without limitation the rights 9 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 10 | copies of the Software, and to permit persons to whom the Software is 11 | furnished to do so, subject to the following conditions: 12 | 13 | The above copyright notice and this permission notice shall be included in all 14 | copies or substantial portions of the Software. 15 | 16 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 17 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 18 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 19 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 20 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 21 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 22 | SOFTWARE. -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # ShellNet: Efficient Point Cloud Convolutional Neural Networks using Concentric Shells Statistics 2 | 3 | International Conference on Computer Vision (ICCV) 2019 (Oral) 4 | 5 | Zhiyuan Zhang, Binh-Son Hua, Sai-Kit Yeung. 6 | 7 | ## Introduction 8 | This is the code release of our paper about building a convolution and neural network for point cloud learning such that the training is fast and accurate. We address this problem by learning point features in regions called 'shells', which resolves point orders and produces local features altogether. Please find the details of the technique in our [project page](https://hkust-vgd.github.io/shellnet/). 9 | 10 | If you found this paper useful in your research, please cite: 11 | ``` 12 | @inproceedings{zhang-shellnet-iccv19, 13 | title = {ShellNet: Efficient Point Cloud Convolutional Neural Networks using Concentric Shells Statistics}, 14 | author = {Zhiyuan Zhang and Binh-Son Hua and Sai-Kit Yeung}, 15 | booktitle = {International Conference on Computer Vision (ICCV)}, 16 | year = {2019} 17 | } 18 | ``` 19 | 20 | ## Installation 21 | The code is based on [PointCNN](https://github.com/yangyanli/PointCNN). Please install [TensorFlow](https://www.tensorflow.org/install/), and follow the instruction in [PointNet++](https://github.com/charlesq34/pointnet2) to compile the customized TF operators in the `tf_ops` folder. 22 | 23 | The code has been tested with Python 3.6, TensorFlow 1.13.2, CUDA 10.0 and cuDNN 7.3 on Ubuntu 14.04. 24 | 25 | ## Code Explanation 26 | The core convolution, ShellConv, and the neural network, ShellNet, are defined in [shellconv.py](shellconv.py). 27 | 28 | ### Convolution Parameters 29 | Let us take the `sconv_params` from [s3dis.py](setting/seg_s3dis.py) as an example: 30 | ``` 31 | ss = 8 32 | sconv_param_name = ('K', 'D', 'P', 'C') 33 | sconv_params = [dict(zip(sconv_param_name, sconv_param)) for sconv_param in 34 | [ 35 | (ss*4, 4, 512, 128), 36 | (ss*2, 2, 128, 256), 37 | (ss*1, 1, 32, 512)]] 38 | ``` 39 | `ss` indicates the shell size which is defined as the number of points contained in each shell. Each element in `sconv_params` is a tuple of `(K, D, P, C)`, where `K` is the neighborhood size, `D` is number of shells, `P` is the representative point number in the output, and `C` is the output channel number. Each tuple specifies the parameters of one `ShellConv` layer, and they are stacked to create a deep network. 40 | 41 | ### Deconvolution Parameters 42 | Similarly, for deconvolution, let us look at `sdconv_params` from [s3dis.py](setting/seg_s3dis.py): 43 | ``` 44 | sdconv_param_name = ('K', 'D', 'pts_layer_idx', 'qrs_layer_idx') 45 | sdconv_params = [dict(zip(sdconv_param_name, sdconv_param)) for sdconv_param in 46 | [ 47 | (ss*1, 1, 2, 1), 48 | (ss*2, 2, 1, 0), 49 | (ss*4, 4, 0, -1)]] 50 | ``` 51 | Each element in `sdconv_params` is a tuple of `(K, D, pts_layer_idx, qrs_layer_idx)`, where `K` and `D` have the same meaning as that in `sconv_params`, `pts_layer_idx` specifies the output of which `ShellConv` layer (from the `sconv_params`) will be the input of this `ShellDeConv` layer, and `qrs_layer_idx` specifies the output of which `ShellConv` layer (from the `sconv_params`) will be forwarded and fused with the output of this `ShellDeConv` layer. The `P` and `C` parameters of this `ShellDeConv` layer is also determined by `qrs_layer_idx`. Similarly, each tuple specifies the parameters of one `ShellDeConv` layer, and they are stacked to create a deep network. 52 | 53 | 54 | 55 | ## Usage 56 | ### Classification 57 | 58 | To train a ShellNet model to classify shapes in the ModelNet40 dataset: 59 | ``` 60 | cd data_conversions 61 | python3 ./download_datasets.py -d modelnet 62 | cd .. 63 | python3 train_val_cls.py 64 | python3 test_cls_modelnet40.py -l log/cls/xxxx 65 | ``` 66 | Our pretrained model can be downloaded [here](https://gohkust-my.sharepoint.com/:u:/g/personal/saikit_ust_hk/EQsJFdSR6YJMsco8Dqss_3kBgTKD8rWfHQcZt7yR9rigTQ?e=3Bj4mK). Please put it to `log/cls/modelnet_pretrained` folder to test. 67 | 68 | 69 | ### Segmentation 70 | We perform segmentation with various datasets, as follows. 71 | 72 | #### ShapeNet 73 | ``` 74 | cd data_conversions 75 | python3 ./download_datasets.py -d shapenet_partseg 76 | python3 ./prepare_partseg_data.py -f ../../data/shapenet_partseg 77 | cd .. 78 | python3 train_val_seg.py -x seg_shapenet 79 | python3 test_seg_shapenet.py -l log/seg/shellconv_seg_shapenet_xxxx/ckpts/epoch-xxx 80 | cd evaluation 81 | python3 eval_shapenet_seg.py -g ../../data/shapenet_partseg/test_label -p ../../data/shapenet_partseg/test_pred_shellnet_1 -a 82 | ``` 83 | 84 | #### ScanNet 85 | Please refer to ScanNet [homepage](http://www.scan-net.org) and PointNet++ preprocessed [data](https://github.com/charlesq34/pointnet2/tree/master/scannet) to download ScanNet. After that, the following script can be used for training and testing: 86 | ``` 87 | cd data_conversions 88 | python3 prepare_scannet_seg_data.py 89 | python3 prepare_scannet_seg_filelists.py 90 | cd .. 91 | python3 train_val_seg.py -x seg_scannet 92 | python3 test_seg_scannet.py -l log/seg/shellconv_seg_scannet_xxxx/ckpts/epoch-xxx 93 | cd evaluation 94 | python3 eval_scannet.py -d -p 95 | ``` 96 | 97 | #### S3DIS 98 | Please download the [S3DIS dataset](http://buildingparser.stanford.edu/dataset.html#Download). The following script performs training and testing: 99 | ``` 100 | cd data_conversions 101 | python3 prepare_s3dis_label.py 102 | python3 prepare_s3dis_data.py 103 | python3 prepare_s3dis_filelists.py 104 | cd .. 105 | python3 train_val_seg.py -x seg_s3dis 106 | python3 test_seg_s3dis.py -l log/seg/shellconv_seg_s3dis_xxxx/ckpts/epoch-xxx 107 | cd evaluation 108 | python3 s3dis_merge.py -d 109 | python3 eval_s3dis.py 110 | ``` 111 | Please notice that these command just for `Area 1` validation. Results on other Areas can be computed by modifying the `filelist` and `filelist_val` in [s3dis.py](setting/seg_s3dis.py). 112 | 113 | #### Semantic3D 114 | You can download our preprocessed hdf5 files and labels [here](https://gohkust-my.sharepoint.com/:u:/g/personal/saikit_ust_hk/Ea_S7Yb-n7JLp1wK5xFihfYBI6vAzHWaA548ytu5k84kdQ?e=YXRcaK). Then: 115 | ``` 116 | python3 train_val_seg.py -x seg_semantic3d 117 | python3 test_seg_semantic3d.py -l log/seg/shellconv_seg_semantic3d_xxxx/ckpts/epoch-xxx 118 | cd evaluation 119 | python3 semantic3d_merge.py -d -v 120 | ``` 121 | 122 | If you prefer to process the data by yourself, here are the steps we used. In general, this data preprocessing of this dataset is more involved. First, please download the original [Semantic3D dataset](http://www.semantic3d.net/view_dbase.php). We then downsample the data using this [script](https://github.com/intel-isl/Open3D-PointNet2-Semantic3D). Finally, we follow PointCNN's [script](https://github.com/yangyanli/PointCNN/tree/master/data_conversions) to split the data into training and validation set, and prepare the .h5 files. 123 | ## License 124 | This repository is released under MIT License (see LICENSE file for details). -------------------------------------------------------------------------------- /data_conversions/classes_ObjClassification-ShapeNetCore55.txt: -------------------------------------------------------------------------------- 1 | 1 trash 2 | 3 basket 3 | 4 bathtub 4 | 5 bed 5 | 9 shelf 6 | 13 cabinet 7 | 18 chair 8 | 20 keyboard 9 | 22 tv 10 | 30 lamp 11 | 31 laptop 12 | 35 microwave 13 | 39 pillow 14 | 42 printer 15 | 47 sofa 16 | 48 stove 17 | 49 table 18 | -------------------------------------------------------------------------------- /data_conversions/download_datasets.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/python3 2 | '''Download datasets for this project.''' 3 | 4 | from __future__ import absolute_import 5 | from __future__ import division 6 | from __future__ import print_function 7 | 8 | import os 9 | import sys 10 | import gzip 11 | import html 12 | import shutil 13 | import tarfile 14 | import zipfile 15 | import requests 16 | import argparse 17 | from tqdm import tqdm 18 | 19 | 20 | # from https://gist.github.com/hrouault/1358474 21 | def query_yes_no(question, default="yes"): 22 | """Ask a yes/no question via raw_input() and return their answer. 23 | "question" is a string that is presented to the user. 24 | "default" is the presumed answer if the user just hits . 25 | It must be "yes" (the default), "no" or None (meaning 26 | an answer is required of the user). 27 | The "answer" return value is one of "yes" or "no". 28 | """ 29 | valid = {"yes": True, "y": True, "ye": True, 30 | "no": False, "n": False} 31 | if default == None: 32 | prompt = " [y/n] " 33 | elif default == "yes": 34 | prompt = " [Y/n] " 35 | elif default == "no": 36 | prompt = " [y/N] " 37 | else: 38 | raise ValueError("invalid default answer: '%s'" % default) 39 | 40 | while True: 41 | sys.stdout.write(question + prompt) 42 | choice = input().lower() 43 | if default is not None and choice == '': 44 | return valid[default] 45 | elif choice in valid: 46 | return valid[choice] 47 | else: 48 | sys.stdout.write("Please respond with 'yes' or 'no' (or 'y' or 'n').\n") 49 | 50 | 51 | def download_from_url(url, dst): 52 | download = True 53 | if os.path.exists(dst): 54 | download = query_yes_no('Seems you have downloaded %s to %s, overwrite?' % (url, dst), default='no') 55 | if download: 56 | os.remove(dst) 57 | 58 | if download: 59 | response = requests.get(url, stream=True) 60 | total_size = int(response.headers.get('content-length', 0)) 61 | chunk_size = 1024 * 1024 62 | bars = total_size // chunk_size 63 | with open(dst, "wb") as handle: 64 | for data in tqdm(response.iter_content(chunk_size=chunk_size), total=bars, desc=url.split('/')[-1], 65 | unit='M'): 66 | handle.write(data) 67 | 68 | 69 | def download_and_unzip(url, root, dataset): 70 | folder = os.path.join(root, dataset) 71 | folder_zips = os.path.join(folder, 'zips') 72 | if not os.path.exists(folder_zips): 73 | os.makedirs(folder_zips) 74 | filename_zip = os.path.join(folder_zips, url.split('/')[-1]) 75 | 76 | download_from_url(url, filename_zip) 77 | 78 | if filename_zip.endswith('.zip'): 79 | zip_ref = zipfile.ZipFile(filename_zip, 'r') 80 | zip_ref.extractall(folder) 81 | zip_ref.close() 82 | elif filename_zip.endswith(('.tar.gz', '.tgz')): 83 | tarfile.open(name=filename_zip, mode="r:gz").extractall(folder) 84 | elif filename_zip.endswith('.gz'): 85 | filename_no_gz = filename_zip[:-3] 86 | with gzip.open(filename_zip, 'rb') as f_in, open(filename_no_gz, 'wb') as f_out: 87 | shutil.copyfileobj(f_in, f_out) 88 | 89 | 90 | def main(): 91 | parser = argparse.ArgumentParser() 92 | parser.add_argument('--folder', '-f', help='Path to data folder.') 93 | parser.add_argument('--dataset', '-d', help='Dataset to download.') 94 | args = parser.parse_args() 95 | print(args) 96 | 97 | root = args.folder if args.folder else '../../data' 98 | if args.dataset == 'modelnet': 99 | download_and_unzip('https://shapenet.cs.stanford.edu/media/modelnet40_ply_hdf5_2048.zip', root, args.dataset) 100 | folder = os.path.join(root, args.dataset) 101 | folder_h5 = os.path.join(folder, 'modelnet40_ply_hdf5_2048') 102 | for filename in os.listdir(folder_h5): 103 | shutil.move(os.path.join(folder_h5, filename), os.path.join(folder, filename)) 104 | shutil.rmtree(folder_h5) 105 | elif args.dataset == 'shapenet_partseg': 106 | download_and_unzip('https://shapenet.cs.stanford.edu/iccv17/partseg/train_data.zip', root, args.dataset) 107 | download_and_unzip('https://shapenet.cs.stanford.edu/iccv17/partseg/train_label.zip', root, args.dataset) 108 | download_and_unzip('https://shapenet.cs.stanford.edu/iccv17/partseg/val_data.zip', root, args.dataset) 109 | download_and_unzip('https://shapenet.cs.stanford.edu/iccv17/partseg/val_label.zip', root, args.dataset) 110 | download_and_unzip('https://shapenet.cs.stanford.edu/iccv17/partseg/test_data.zip', root, args.dataset) 111 | download_and_unzip('https://shapenet.cs.stanford.edu/iccv17/partseg/test_label.zip', root, args.dataset) 112 | 113 | 114 | if __name__ == '__main__': 115 | main() 116 | -------------------------------------------------------------------------------- /data_conversions/download_semantic3d.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | BASE_DIR=${1-../../data/semantic3d} 3 | 4 | # Training data 5 | wget -c -N http://semantic3d.net/data/point-clouds/training1/bildstein_station1_xyz_intensity_rgb.7z -P $BASE_DIR/train/ 6 | wget -c -N http://semantic3d.net/data/point-clouds/training1/bildstein_station5_xyz_intensity_rgb.7z -P $BASE_DIR/train/ 7 | wget -c -N http://semantic3d.net/data/point-clouds/training1/domfountain_station1_xyz_intensity_rgb.7z -P $BASE_DIR/train/ 8 | wget -c -N http://semantic3d.net/data/point-clouds/training1/domfountain_station3_xyz_intensity_rgb.7z -P $BASE_DIR/train/ 9 | wget -c -N http://semantic3d.net/data/point-clouds/training1/neugasse_station1_xyz_intensity_rgb.7z -P $BASE_DIR/train/ 10 | wget -c -N http://semantic3d.net/data/point-clouds/training1/sg27_station1_intensity_rgb.7z -P $BASE_DIR/train/ 11 | wget -c -N http://semantic3d.net/data/point-clouds/training1/sg27_station2_intensity_rgb.7z -P $BASE_DIR/train/ 12 | wget -c -N http://semantic3d.net/data/point-clouds/training1/sg27_station5_intensity_rgb.7z -P $BASE_DIR/train/ 13 | wget -c -N http://semantic3d.net/data/point-clouds/training1/sg27_station9_intensity_rgb.7z -P $BASE_DIR/train/ 14 | wget -c -N http://semantic3d.net/data/point-clouds/training1/sg28_station4_intensity_rgb.7z -P $BASE_DIR/train/ 15 | wget -c -N http://semantic3d.net/data/point-clouds/training1/untermaederbrunnen_station1_xyz_intensity_rgb.7z -P $BASE_DIR/train/ 16 | wget -c -N http://semantic3d.net/data/sem8_labels_training.7z -P $BASE_DIR/train/ 17 | 18 | # Validation data 19 | wget -c -N http://semantic3d.net/data/point-clouds/training1/bildstein_station3_xyz_intensity_rgb.7z -P $BASE_DIR/val/ 20 | wget -c -N http://semantic3d.net/data/point-clouds/training1/domfountain_station2_xyz_intensity_rgb.7z -P $BASE_DIR/val/ 21 | wget -c -N http://semantic3d.net/data/point-clouds/training1/sg27_station4_intensity_rgb.7z -P $BASE_DIR/val/ 22 | wget -c -N http://semantic3d.net/data/point-clouds/training1/untermaederbrunnen_station3_xyz_intensity_rgb.7z -P $BASE_DIR/val/ 23 | mv $BASE_DIR/train/bildstein_station3_xyz_intensity_rgb.labels $BASE_DIR/val 24 | mv $BASE_DIR/train/domfountain_station2_xyz_intensity_rgb.labels $BASE_DIR/val 25 | mv $BASE_DIR/train/sg27_station4_intensity_rgb.labels $BASE_DIR/val 26 | mv $BASE_DIR/train/untermaederbrunnen_station3_xyz_intensity_rgb.labels $BASE_DIR/val 27 | 28 | # Test data 29 | wget -c -N http://semantic3d.net/data/point-clouds/testing2/MarketplaceFeldkirch_Station4_rgb_intensity-reduced.txt.7z -P $BASE_DIR/test_reduced/ 30 | wget -c -N http://semantic3d.net/data/point-clouds/testing2/StGallenCathedral_station6_rgb_intensity-reduced.txt.7z -P $BASE_DIR/test_reduced/ 31 | wget -c -N http://semantic3d.net/data/point-clouds/testing2/sg27_station10_rgb_intensity-reduced.txt.7z -P $BASE_DIR/test_reduced/ 32 | wget -c -N http://semantic3d.net/data/point-clouds/testing2/sg28_Station2_rgb_intensity-reduced.txt.7z -P $BASE_DIR/test_reduced/ -------------------------------------------------------------------------------- /data_conversions/extract_scannet_objs.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/python3 2 | """Convert original Scannet data to PointCNN Classification dataset""" 3 | 4 | from __future__ import absolute_import 5 | from __future__ import division 6 | from __future__ import print_function 7 | 8 | import os 9 | import json 10 | import plyfile 11 | from plyfile import PlyData 12 | import argparse 13 | import numpy as np 14 | 15 | 16 | def dir(root, type='f', addroot=True): 17 | dirList = [] 18 | fileList = [] 19 | root = root + "/" 20 | files = os.listdir(root) 21 | 22 | for f in files: 23 | if (os.path.isdir(root + f)): 24 | if addroot == True: 25 | dirList.append(root + f) 26 | else: 27 | dirList.append(f) 28 | 29 | if (os.path.isfile(root + f)): 30 | if addroot == True: 31 | fileList.append(root + f) 32 | else: 33 | fileList.append(f) 34 | 35 | if type == "f": 36 | return fileList 37 | elif type == "d": 38 | return dirList 39 | else: 40 | print("ERROR: TMC.dir(root,type) type must be [f] for file or [d] for dir") 41 | return 0 42 | 43 | 44 | def save_ply(points, colors, filename): 45 | vertex = np.array([tuple(p) for p in points], dtype=[('x', 'f4'), ('y', 'f4'), ('z', 'f4')]) 46 | 47 | vertex_color = np.array([tuple(c) for c in colors], dtype=[('red', 'u1'), ('green', 'u1'), ('blue', 'u1')]) 48 | 49 | n = len(vertex) 50 | assert len(vertex_color) == n 51 | 52 | vertex_all = np.empty(n, dtype=vertex.dtype.descr + vertex_color.dtype.descr) 53 | 54 | for prop in vertex.dtype.names: 55 | vertex_all[prop] = vertex[prop] 56 | 57 | for prop in vertex_color.dtype.names: 58 | vertex_all[prop] = vertex_color[prop] 59 | 60 | ply = plyfile.PlyData([plyfile.PlyElement.describe(vertex_all, 'vertex')], text=False) 61 | ply.write(filename) 62 | 63 | print("save ply to", filename) 64 | 65 | 66 | def pc_getbbox(pc): 67 | x = [] 68 | y = [] 69 | z = [] 70 | 71 | for pts in pc: 72 | x.append(pts[0]) 73 | y.append(pts[1]) 74 | z.append(pts[2]) 75 | 76 | boundary = [min(x), max(x), min(y), max(y), min(z), max(z)] 77 | 78 | return boundary 79 | 80 | 81 | def scene2instances(scene_path, out_root, all_label, label_map, issave_ply): 82 | print("Process Scene:", scene_path) 83 | 84 | sceneid = scene_path.strip().split("scene")[1] 85 | spaceid = sceneid.split("_")[0] 86 | scanid = sceneid.split("_")[1] 87 | 88 | label_list = all_label[0] 89 | label_info = all_label[1] 90 | 91 | pts_dir = out_root + "/pts/" 92 | 93 | # check the path 94 | if not os.path.exists(pts_dir): 95 | print(pts_dir, "Not Exists! Create", pts_dir) 96 | os.makedirs(pts_dir) 97 | 98 | if save_ply: 99 | 100 | ply_dir = out_root + "/ply/" + spaceid + scanid + "/" 101 | 102 | if not os.path.exists(ply_dir): 103 | print(ply_dir, "Not Exists! Create", ply_dir) 104 | os.makedirs(ply_dir) 105 | 106 | ply_file = scene_path + "/scene" + sceneid + "_vh_clean_2.ply" 107 | jsonflie = scene_path + "/scene" + sceneid + "_vh_clean_2.0.010000.segs.json" 108 | aggjsonfile = scene_path + "/scene" + sceneid + ".aggregation.json" 109 | 110 | # Read ply file 111 | print("\nRead ply file:", ply_file) 112 | plydata = PlyData.read(ply_file).elements[0].data 113 | pts_num = len(plydata) 114 | print("points num:", pts_num) 115 | 116 | # Read json file 117 | print("Read json file:", jsonflie) 118 | json_data = json.load(open(jsonflie)) 119 | 120 | # check json file 121 | if json_data['sceneId'].strip() == ('scene' + sceneid): 122 | 123 | segIndices = json_data['segIndices'] 124 | seg_num = len(segIndices) 125 | 126 | # check num 127 | if seg_num != pts_num: 128 | print("seg num != pts num!") 129 | os.exit(0) 130 | 131 | else: 132 | 133 | print("Read Wrong Json File!") 134 | os.exit(0) 135 | 136 | # Read aggregation json file 137 | print("Read aggregation json file:", aggjsonfile) 138 | aggjson_data = json.load(open(aggjsonfile)) 139 | 140 | # check json file 141 | if aggjson_data['sceneId'].strip() == ('scene' + sceneid): 142 | 143 | segGroups = aggjson_data['segGroups'] 144 | 145 | else: 146 | 147 | print("Read Wrong Aggregation Json File!") 148 | os.exit(0) 149 | 150 | # split pts 151 | obj_dict = {} 152 | 153 | for k, pts in enumerate(plydata): 154 | 155 | seg_indice = segIndices[k] 156 | 157 | # find obj 158 | label = "unannotated" 159 | objid = -1 160 | 161 | for seg in segGroups: 162 | 163 | segments = seg['segments'] 164 | 165 | if seg_indice in segments: 166 | label = seg['label'].strip() 167 | objid = seg['objectId'] 168 | 169 | break 170 | 171 | obj_key = str(label + "_" + str(objid)).strip() 172 | 173 | if obj_key not in obj_dict.keys(): 174 | obj_dict[obj_key] = [] 175 | 176 | obj_dict[obj_key].append(pts) 177 | 178 | # save data file by obj 179 | for objkey in obj_dict.keys(): 180 | 181 | obj_label = objkey.split("_")[0] 182 | obj_id = objkey.split("_")[1] 183 | 184 | if obj_label in label_list: 185 | 186 | label_id = label_list.index(obj_label) 187 | 188 | else: 189 | 190 | label_id = 0 191 | 192 | label_full = label_info[label_id] 193 | 194 | label_id = label_full[0] 195 | label_s55 = label_full[3] 196 | label_s55_id = 0 197 | 198 | for l in label_map: 199 | 200 | if label_s55 == l[1]: 201 | label_s55_id = l[0] 202 | 203 | if label_s55_id != 0: 204 | 205 | pts_out_file = pts_dir + spaceid + scanid + "%04d" % int(obj_id) + "_" + str(label_s55_id) + ".pts" 206 | f_pts = open(pts_out_file, "w") 207 | 208 | pts = [] 209 | rgb = [] 210 | 211 | for p in obj_dict[objkey]: 212 | pts.append([p[0], p[1], p[2]]) 213 | rgb.append([p[3], p[4], p[5]]) 214 | 215 | bbox = pc_getbbox(pts) 216 | dimxy = [bbox[1] - bbox[0], bbox[3] - bbox[2]] 217 | 218 | # retrans 219 | for i in range(len(pts)): 220 | pts[i] = [pts[i][0] - bbox[0] - dimxy[0] / 2, pts[i][2] - bbox[4], pts[i][1] - bbox[2] - dimxy[1] / 2] 221 | 222 | if issave_ply: 223 | 224 | ply_out_file = ply_dir + spaceid + "_" + scanid + "_" + "%04d" % int( 225 | obj_id) + "_" + obj_label + "_" + str(label_id) + "_" + str(label_s55) + "_" + str( 226 | label_s55_id) + ".ply" 227 | 228 | seg_pts = [] 229 | seg_ply = [] 230 | 231 | for k, p in enumerate(pts): 232 | seg_pts.append((p[0], p[1], p[2])) 233 | seg_ply.append((rgb[k][0], rgb[k][1], rgb[k][2])) 234 | 235 | save_ply(seg_pts, seg_ply, ply_out_file) 236 | 237 | # write pts xyzrgb 238 | for k, p in enumerate(pts): 239 | f_pts.writelines(str(p[0]) + " " + str(p[1]) + " " + str(p[2]) + " " + str(rgb[k][0]) + " " + str( 240 | rgb[k][1]) + " " + str(rgb[k][2]) + "\n") 241 | 242 | f_pts.close() 243 | 244 | 245 | def main(): 246 | parser = argparse.ArgumentParser() 247 | parser.add_argument('--folder', '-f', help='Path to data folder') 248 | parser.add_argument('--benchmark', '-b', help='Path to benchmark folder') 249 | parser.add_argument('--outpath', '-o', help='Path to output folder') 250 | parser.add_argument('--saveply', '-s', action='store_true', help='Save color ply or not') 251 | args = parser.parse_args() 252 | print(args) 253 | 254 | label_tsv = args.benchmark + "/scannet-labels.combined.tsv" 255 | trainval_list_file = args.benchmark + "/scannet_trainval.txt" 256 | test_list_file = args.benchmark + "/scannet_test.txt" 257 | label_shapenetcore55 = args.benchmark + "/classes_ObjClassification-ShapeNetCore55.txt" 258 | 259 | ##########################################################Read Source########################################################## 260 | 261 | print("read scene dir:", args.folder) 262 | scenedir = dir(args.folder, 'd') 263 | 264 | print("read trainval list:", trainval_list_file) 265 | train_scene_list = [] 266 | with open(trainval_list_file, 'r') as train_f: 267 | for line in train_f.readlines(): 268 | sceneid = line.strip().split("scene")[1] 269 | spaceid = sceneid.split("_")[0] 270 | scanid = sceneid.split("_")[1] 271 | train_scene_list.append(spaceid + scanid) 272 | 273 | print("read test list:", test_list_file) 274 | test_scene_list = [] 275 | with open(test_list_file, 'r') as train_f: 276 | for line in train_f.readlines(): 277 | sceneid = line.strip().split("scene")[1] 278 | spaceid = sceneid.split("_")[0] 279 | scanid = sceneid.split("_")[1] 280 | test_scene_list.append(spaceid + scanid) 281 | 282 | print("read label tsv file:", label_tsv) 283 | label_map = [] 284 | label_info = [] 285 | 286 | with open(label_tsv, 'r') as tsv_f: 287 | 288 | for k, line in enumerate(tsv_f.readlines()): 289 | 290 | if k > 0: 291 | line_s = line.strip().split('\t') 292 | 293 | label_id = int(line_s[0]) 294 | category = line_s[1] 295 | count = int(line_s[2]) 296 | ShapeNetCore55 = line_s[11] 297 | 298 | label_map.append(category) 299 | label_info.append([label_id, category, count, ShapeNetCore55]) 300 | 301 | print("read shapenetcore55 label file:", label_shapenetcore55) 302 | label_shapenetcore55_map = [] 303 | with open(label_shapenetcore55, 'r') as label_shapenetcore55_f: 304 | 305 | for k, line in enumerate(label_shapenetcore55_f.readlines()): 306 | line_s = line.strip().split('\t') 307 | label_id = line_s[0] 308 | category = line_s[1] 309 | label_shapenetcore55_map.append([label_id, category]) 310 | 311 | # split scene to train and test 312 | process_train_list = [] 313 | process_test_list = [] 314 | 315 | for scene in scenedir: 316 | 317 | sceneid = scene.strip().split("scene")[1] 318 | spaceid = sceneid.split("_")[0] 319 | scanid = sceneid.split("_")[1] 320 | scenename = spaceid + scanid 321 | 322 | if scenename in train_scene_list: 323 | 324 | process_train_list.append(scene) 325 | 326 | elif scenename in test_scene_list: 327 | 328 | process_test_list.append(scene) 329 | 330 | print("Train all:", len(train_scene_list), "Test all:", len(test_scene_list), "Dir all:", len(scenedir)) 331 | print("Process Train:", len(process_train_list), "Process Test:", len(process_test_list)) 332 | 333 | ##########################################################Process Data########################################################## 334 | print("Process Train Scene:") 335 | 336 | for scene in process_train_list: 337 | scene2instances(scene, args.outpath + "/train/", [label_map, label_info], label_shapenetcore55_map, 338 | args.saveply) 339 | 340 | print("Process Test Scene:") 341 | 342 | for scene in process_test_list: 343 | scene2instances(scene, args.outpath + "/test/", [label_map, label_info], label_shapenetcore55_map, args.saveply) 344 | 345 | 346 | if __name__ == '__main__': 347 | main() 348 | -------------------------------------------------------------------------------- /data_conversions/prepare_partseg_data.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/python3 2 | '''Prepare Data for ShapeNet Segmentation Task.''' 3 | 4 | from __future__ import absolute_import 5 | from __future__ import division 6 | from __future__ import print_function 7 | 8 | import os 9 | import sys 10 | import h5py 11 | import argparse 12 | import numpy as np 13 | from datetime import datetime 14 | 15 | sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) 16 | from utils import data_utils 17 | 18 | def main(): 19 | parser = argparse.ArgumentParser() 20 | parser.add_argument('--folder', '-f', help='Path to data folder') 21 | parser.add_argument('--save_ply', '-s', help='Convert .pts to .ply', action='store_true') 22 | args = parser.parse_args() 23 | print(args) 24 | 25 | root = args.folder if args.folder else '../../data/shapenet_partseg' 26 | folders = [(os.path.join(root,'train_data'), os.path.join(root,'train_label')), 27 | (os.path.join(root,'val_data'), os.path.join(root,'val_label')), 28 | (os.path.join(root,'test_data'), os.path.join(root,'test_label'))] 29 | category_label_seg_max_dict = dict() 30 | max_point_num = 0 31 | label_seg_min = sys.maxsize 32 | for data_folder, label_folder in folders: 33 | if not os.path.exists(data_folder): 34 | continue 35 | for category in sorted(os.listdir(data_folder)): 36 | if category not in category_label_seg_max_dict: 37 | category_label_seg_max_dict[category] = 0 38 | data_category_folder = os.path.join(data_folder, category) 39 | category_label_seg_max = 0 40 | for filename in sorted(os.listdir(data_category_folder)): 41 | data_filepath = os.path.join(data_category_folder, filename) 42 | coordinates = [xyz for xyz in open(data_filepath, 'r') if len(xyz.split(' ')) == 3] 43 | max_point_num = max(max_point_num, len(coordinates)) 44 | 45 | if label_folder is not None: 46 | label_filepath = os.path.join(label_folder, category, filename[0:-3] + 'seg') 47 | label_seg_this = np.loadtxt(label_filepath).astype(np.int32) 48 | assert (len(coordinates) == len(label_seg_this)) 49 | category_label_seg_max = max(category_label_seg_max, max(label_seg_this)) 50 | label_seg_min = min(label_seg_min, min(label_seg_this)) 51 | category_label_seg_max_dict[category] = max(category_label_seg_max_dict[category], category_label_seg_max) 52 | category_label_seg_max_list = [(key, category_label_seg_max_dict[key]) for key in 53 | sorted(category_label_seg_max_dict.keys())] 54 | 55 | category_label = dict() 56 | offset = 0 57 | category_offset = dict() 58 | label_seg_max = max([category_label_seg_max for _, category_label_seg_max in category_label_seg_max_list]) 59 | with open(os.path.join(root, 'categories.txt'), 'w') as file_categories: 60 | for idx, (category, category_label_seg_max) in enumerate(category_label_seg_max_list): 61 | file_categories.write('%s %d\n' % (category, category_label_seg_max - label_seg_min + 1)) 62 | category_label[category] = idx 63 | category_offset[category] = offset 64 | offset = offset + category_label_seg_max - label_seg_min + 1 65 | 66 | print('part_num:', offset) 67 | print('max_point_num:', max_point_num) 68 | print(category_label_seg_max_list) 69 | 70 | batch_size = 2048 71 | data = np.zeros((batch_size, max_point_num, 3)) 72 | data_num = np.zeros((batch_size), dtype=np.int32) 73 | label = np.zeros((batch_size), dtype=np.int32) 74 | label_seg = np.zeros((batch_size, max_point_num), dtype=np.int32) 75 | for data_folder, label_folder in folders: 76 | if not os.path.exists(data_folder): 77 | continue 78 | data_folder_ply = data_folder + '_ply' 79 | file_num = 0 80 | for category in sorted(os.listdir(data_folder)): 81 | data_category_folder = os.path.join(data_folder, category) 82 | file_num = file_num + len(os.listdir(data_category_folder)) 83 | idx_h5 = 0 84 | idx = 0 85 | 86 | save_path = '%s/%s' % (os.path.dirname(data_folder), os.path.basename(data_folder)[0:-5]) 87 | filename_txt = '%s_files.txt' % (save_path) 88 | ply_filepath_list = [] 89 | with open(filename_txt, 'w') as filelist: 90 | for category in sorted(os.listdir(data_folder)): 91 | data_category_folder = os.path.join(data_folder, category) 92 | for filename in sorted(os.listdir(data_category_folder)): 93 | data_filepath = os.path.join(data_category_folder, filename) 94 | coordinates = [[float(value) for value in xyz.split(' ')] 95 | for xyz in open(data_filepath, 'r') if len(xyz.split(' ')) == 3] 96 | idx_in_batch = idx % batch_size 97 | data[idx_in_batch, 0:len(coordinates), ...] = np.array(coordinates) 98 | data_num[idx_in_batch] = len(coordinates) 99 | label[idx_in_batch] = category_label[category] 100 | 101 | if label_folder is not None: 102 | label_filepath = os.path.join(label_folder, category, filename[0:-3] + 'seg') 103 | label_seg_this = np.loadtxt(label_filepath).astype(np.int32) - label_seg_min 104 | assert (len(coordinates) == label_seg_this.shape[0]) 105 | label_seg[idx_in_batch, 0:len(coordinates)] = label_seg_this + category_offset[category] 106 | 107 | data_ply_filepath = os.path.join(data_folder_ply, category, filename[:-3] + 'ply') 108 | ply_filepath_list.append(data_ply_filepath) 109 | 110 | if ((idx + 1) % batch_size == 0) or idx == file_num - 1: 111 | item_num = idx_in_batch + 1 112 | filename_h5 = '%s_%d.h5' % (save_path, idx_h5) 113 | print('{}-Saving {}...'.format(datetime.now(), filename_h5)) 114 | filelist.write('./%s_%d.h5\n' % (os.path.basename(data_folder)[0:-5], idx_h5)) 115 | 116 | file = h5py.File(filename_h5, 'w') 117 | file.create_dataset('data', data=data[0:item_num, ...]) 118 | file.create_dataset('data_num', data=data_num[0:item_num, ...]) 119 | file.create_dataset('label', data=label[0:item_num, ...]) 120 | file.create_dataset('label_seg', data=label_seg[0:item_num, ...]) 121 | file.close() 122 | 123 | if args.save_ply: 124 | data_utils.save_ply_property_batch(data[0:item_num, ...], label_seg[0:item_num, ...], 125 | ply_filepath_list, data_num[0:item_num, ...], 126 | label_seg_max - label_seg_min) 127 | ply_filepath_list = [] 128 | idx_h5 = idx_h5 + 1 129 | idx = idx + 1 130 | 131 | train_val_txt = os.path.join(root, "train_val_files.txt") 132 | with open(train_val_txt, "w") as train_val: 133 | for part in ("train", "val"): 134 | part_txt = os.path.join(root, "%s_files.txt" % part) 135 | train_val.write(open(part_txt, "r").read()) 136 | 137 | if __name__ == '__main__': 138 | main() 139 | -------------------------------------------------------------------------------- /data_conversions/prepare_s3dis_data.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/python3 2 | '''Prepare Data for S3DIS Segmentation Task.''' 3 | 4 | from __future__ import absolute_import 5 | from __future__ import division 6 | from __future__ import print_function 7 | 8 | import os 9 | import sys 10 | import math 11 | import h5py 12 | import argparse 13 | import numpy as np 14 | from datetime import datetime 15 | 16 | BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) 17 | sys.path.append(BASE_DIR) 18 | sys.path.append(os.path.join(BASE_DIR, 'utils')) 19 | 20 | import data_utils 21 | 22 | 23 | def main(): 24 | parser = argparse.ArgumentParser() 25 | parser.add_argument('--folder', '-f', default = '../data/3DIS/prepare_label_rgb', help='Path to data folder') 26 | parser.add_argument('--max_point_num', '-m', help='Max point number of each sample', type=int, default=8192) 27 | parser.add_argument('--block_size', '-b', help='Block size', type=float, default=1.5) 28 | parser.add_argument('--grid_size', '-g', help='Grid size', type=float, default=0.03) 29 | parser.add_argument('--save_ply', '-s', help='Convert .pts to .ply', action='store_true', default=True) 30 | 31 | args = parser.parse_args() 32 | print(args) 33 | 34 | root = args.folder if args.folder else '../../data/s3dis/prepare_label_rgb' 35 | max_point_num = args.max_point_num 36 | 37 | batch_size = 2048 38 | data = np.zeros((batch_size, max_point_num, 6)) 39 | data_num = np.zeros((batch_size), dtype=np.int32) 40 | label = np.zeros((batch_size), dtype=np.int32) 41 | label_seg = np.zeros((batch_size, max_point_num), dtype=np.int32) 42 | indices_split_to_full = np.zeros((batch_size, max_point_num), dtype=np.int32) 43 | 44 | for area_idx in range(1, 7): 45 | folder = os.path.join(root, 'Area_%d' % area_idx) 46 | datasets = [dataset for dataset in os.listdir(folder)] 47 | for dataset_idx, dataset in enumerate(datasets): 48 | dataset_marker = os.path.join(folder, dataset, ".dataset") 49 | if os.path.exists(dataset_marker): 50 | print('{}-{}/{} already processed, skipping'.format(datetime.now(), folder, dataset)) 51 | continue 52 | filename_data = os.path.join(folder, dataset, 'xyzrgb.npy') 53 | print('{}-Loading {}...'.format(datetime.now(), filename_data)) 54 | xyzrgb = np.load(filename_data) 55 | 56 | filename_labels = os.path.join(folder, dataset, 'label.npy') 57 | print('{}-Loading {}...'.format(datetime.now(), filename_labels)) 58 | labels = np.load(filename_labels).astype(int).flatten() 59 | 60 | xyz, rgb = np.split(xyzrgb, [3], axis=-1) 61 | xyz_min = np.amin(xyz, axis=0, keepdims=True) 62 | xyz_max = np.amax(xyz, axis=0, keepdims=True) 63 | xyz_center = (xyz_min + xyz_max) / 2 64 | xyz_center[0][-1] = xyz_min[0][-1] 65 | xyz = xyz - xyz_center # align to room bottom center 66 | rgb = rgb / 255 - 0.5 67 | 68 | offsets = [('zero', 0.0), ('half', args.block_size / 2)] 69 | for offset_name, offset in offsets: 70 | idx_h5 = 0 71 | idx = 0 72 | 73 | print('{}-Computing block id of {} points...'.format(datetime.now(), xyzrgb.shape[0])) 74 | xyz_min = np.amin(xyz, axis=0, keepdims=True) - offset 75 | xyz_max = np.amax(xyz, axis=0, keepdims=True) 76 | block_size = (args.block_size, args.block_size, 2 * (xyz_max[0, -1] - xyz_min[0, -1])) 77 | xyz_blocks = np.floor((xyz - xyz_min) / block_size).astype(np.int) 78 | 79 | print('{}-Collecting points belong to each block...'.format(datetime.now(), xyzrgb.shape[0])) 80 | blocks, point_block_indices, block_point_counts = np.unique(xyz_blocks, return_inverse=True, 81 | return_counts=True, axis=0) 82 | block_point_indices = np.split(np.argsort(point_block_indices), np.cumsum(block_point_counts[:-1])) 83 | print('{}-{} is split into {} blocks.'.format(datetime.now(), dataset, blocks.shape[0])) 84 | 85 | block_to_block_idx_map = dict() 86 | for block_idx in range(blocks.shape[0]): 87 | block = (blocks[block_idx][0], blocks[block_idx][1]) 88 | block_to_block_idx_map[(block[0], block[1])] = block_idx 89 | 90 | # merge small blocks into one of their big neighbors 91 | block_point_count_threshold = max_point_num/10 92 | nbr_block_offsets = [(0, 1), (1, 0), (0, -1), (-1, 0), (-1, 1), (1, 1), (1, -1), (-1, -1)] 93 | block_merge_count = 0 94 | for block_idx in range(blocks.shape[0]): 95 | if block_point_counts[block_idx] >= block_point_count_threshold: 96 | continue 97 | 98 | block = (blocks[block_idx][0], blocks[block_idx][1]) 99 | for x, y in nbr_block_offsets: 100 | nbr_block = (block[0] + x, block[1] + y) 101 | if nbr_block not in block_to_block_idx_map: 102 | continue 103 | 104 | nbr_block_idx = block_to_block_idx_map[nbr_block] 105 | if block_point_counts[nbr_block_idx] < block_point_count_threshold: 106 | continue 107 | 108 | block_point_indices[nbr_block_idx] = np.concatenate( 109 | [block_point_indices[nbr_block_idx], block_point_indices[block_idx]], axis=-1) 110 | block_point_indices[block_idx] = np.array([], dtype=np.int) 111 | block_merge_count = block_merge_count + 1 112 | break 113 | print('{}-{} of {} blocks are merged.'.format(datetime.now(), block_merge_count, blocks.shape[0])) 114 | 115 | idx_last_non_empty_block = 0 116 | for block_idx in reversed(range(blocks.shape[0])): 117 | if block_point_indices[block_idx].shape[0] != 0: 118 | idx_last_non_empty_block = block_idx 119 | break 120 | 121 | # uniformly sample each block 122 | for block_idx in range(idx_last_non_empty_block + 1): 123 | point_indices = block_point_indices[block_idx] 124 | if point_indices.shape[0] == 0: 125 | continue 126 | block_points = xyz[point_indices] 127 | block_min = np.amin(block_points, axis=0, keepdims=True) 128 | xyz_grids = np.floor((block_points - block_min) / args.grid_size).astype(np.int) 129 | grids, point_grid_indices, grid_point_counts = np.unique(xyz_grids, return_inverse=True, 130 | return_counts=True, axis=0) 131 | grid_point_indices = np.split(np.argsort(point_grid_indices), np.cumsum(grid_point_counts[:-1])) 132 | grid_point_count_avg = int(np.average(grid_point_counts)) 133 | point_indices_repeated = [] 134 | for grid_idx in range(grids.shape[0]): 135 | point_indices_in_block = grid_point_indices[grid_idx] 136 | repeat_num = math.ceil(grid_point_count_avg / point_indices_in_block.shape[0]) 137 | if repeat_num > 1: 138 | point_indices_in_block = np.repeat(point_indices_in_block, repeat_num) 139 | np.random.shuffle(point_indices_in_block) 140 | point_indices_in_block = point_indices_in_block[:grid_point_count_avg] 141 | point_indices_repeated.extend(list(point_indices[point_indices_in_block])) 142 | block_point_indices[block_idx] = np.array(point_indices_repeated) 143 | block_point_counts[block_idx] = len(point_indices_repeated) 144 | 145 | for block_idx in range(idx_last_non_empty_block + 1): 146 | point_indices = block_point_indices[block_idx] 147 | if point_indices.shape[0] == 0: 148 | continue 149 | 150 | block_point_num = point_indices.shape[0] 151 | block_split_num = int(math.ceil(block_point_num * 1.0 / max_point_num)) 152 | point_num_avg = int(math.ceil(block_point_num * 1.0 / block_split_num)) 153 | point_nums = [point_num_avg] * block_split_num 154 | point_nums[-1] = block_point_num - (point_num_avg * (block_split_num - 1)) 155 | starts = [0] + list(np.cumsum(point_nums)) 156 | 157 | np.random.shuffle(point_indices) 158 | block_points = xyz[point_indices] 159 | block_rgb = rgb[point_indices] 160 | block_labels = labels[point_indices] 161 | x, y, z = np.split(block_points, (1, 2), axis=-1) 162 | block_xzyrgb = np.concatenate([x, z, y, block_rgb], axis=-1) 163 | 164 | for block_split_idx in range(block_split_num): 165 | start = starts[block_split_idx] 166 | point_num = point_nums[block_split_idx] 167 | end = start + point_num 168 | idx_in_batch = idx % batch_size 169 | data[idx_in_batch, 0:point_num, ...] = block_xzyrgb[start:end, :] 170 | data_num[idx_in_batch] = point_num 171 | label[idx_in_batch] = dataset_idx # won't be used... 172 | label_seg[idx_in_batch, 0:point_num] = block_labels[start:end] 173 | indices_split_to_full[idx_in_batch, 0:point_num] = point_indices[start:end] 174 | 175 | if ((idx + 1) % batch_size == 0) or \ 176 | (block_idx == idx_last_non_empty_block and block_split_idx == block_split_num - 1): 177 | item_num = idx_in_batch + 1 178 | filename_h5 = os.path.join(folder, dataset, '%s_%d.h5' % (offset_name, idx_h5)) 179 | print('{}-Saving {}...'.format(datetime.now(), filename_h5)) 180 | 181 | file = h5py.File(filename_h5, 'w') 182 | file.create_dataset('data', data=data[0:item_num, ...]) 183 | file.create_dataset('data_num', data=data_num[0:item_num, ...]) 184 | file.create_dataset('label', data=label[0:item_num, ...]) 185 | file.create_dataset('label_seg', data=label_seg[0:item_num, ...]) 186 | file.create_dataset('indices_split_to_full', data=indices_split_to_full[0:item_num, ...]) 187 | file.close() 188 | 189 | if args.save_ply: 190 | print('{}-Saving ply of {}...'.format(datetime.now(), filename_h5)) 191 | filepath_label_ply = os.path.join(folder, dataset, 'ply_label', 192 | 'label_%s_%d' % (offset_name, idx_h5)) 193 | data_utils.save_ply_property_batch(data[0:item_num, :, 0:3], 194 | label_seg[0:item_num, ...], 195 | filepath_label_ply, data_num[0:item_num, ...], 14) 196 | 197 | filepath_rgb_ply = os.path.join(folder, dataset, 'ply_rgb', 198 | 'rgb_%s_%d' % (offset_name, idx_h5)) 199 | data_utils.save_ply_color_batch(data[0:item_num, :, 0:3], 200 | (data[0:item_num, :, 3:] + 0.5) * 255, 201 | filepath_rgb_ply, data_num[0:item_num, ...]) 202 | 203 | idx_h5 = idx_h5 + 1 204 | idx = idx + 1 205 | 206 | # Marker indicating we've processed this dataset 207 | open(dataset_marker, "w").close() 208 | 209 | if __name__ == '__main__': 210 | main() 211 | print('{}-Done.'.format(datetime.now())) 212 | -------------------------------------------------------------------------------- /data_conversions/prepare_s3dis_filelists.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/python3 2 | '''Prepare Filelists for S3DIS Segmentation Task.''' 3 | 4 | from __future__ import absolute_import 5 | from __future__ import division 6 | from __future__ import print_function 7 | 8 | import os 9 | import math 10 | import random 11 | import argparse 12 | from datetime import datetime 13 | 14 | 15 | def main(): 16 | parser = argparse.ArgumentParser() 17 | parser.add_argument('--folder', '-f', help='Path to data folder') 18 | parser.add_argument('--h5_num', '-d', help='Number of h5 files to be loaded each time', type=int, default=8) 19 | parser.add_argument('--repeat_num', '-r', help='Number of repeatly using each loaded h5 list', type=int, default=1) 20 | 21 | args = parser.parse_args() 22 | print(args) 23 | 24 | root = args.folder if args.folder else '../../data/3DIS/prepare_label_rgb/' 25 | 26 | area_h5s = [[] for _ in range(6)] 27 | for area_idx in range(1, 7): 28 | folder = os.path.join(root, 'Area_%d' % area_idx) 29 | datasets = [dataset for dataset in os.listdir(folder)] 30 | for dataset in datasets: 31 | folder_dataset = os.path.join(folder, dataset) 32 | filename_h5s = ['./Area_%d/%s/%s\n' % (area_idx, dataset, filename) for filename in 33 | os.listdir(folder_dataset) 34 | if filename.endswith('.h5')] 35 | area_h5s[area_idx - 1].extend(filename_h5s) 36 | 37 | for area_idx in range(1, 7): 38 | train_h5 = [filename for idx in range(6) if idx + 1 != area_idx for filename in area_h5s[idx]] 39 | random.shuffle(train_h5) 40 | train_list = os.path.join(root, 'train_files_for_val_on_Area_%d.txt' % area_idx) 41 | print('{}-Saving {}...'.format(datetime.now(), train_list)) 42 | with open(train_list, 'w') as filelist: 43 | list_num = math.ceil(len(train_h5) / args.h5_num) 44 | for list_idx in range(list_num): 45 | train_val_list_i = os.path.join(root, 'filelists', 46 | 'train_files_for_val_on_Area_%d_g_%d.txt' % (area_idx, list_idx)) 47 | os.makedirs(os.path.dirname(train_val_list_i), exist_ok=True) 48 | with open(train_val_list_i, 'w') as filelist_i: 49 | for h5_idx in range(args.h5_num): 50 | filename_idx = list_idx * args.h5_num + h5_idx 51 | if filename_idx > len(train_h5) - 1: 52 | break 53 | filename_h5 = train_h5[filename_idx] 54 | filelist_i.write('../' + filename_h5) 55 | for repeat_idx in range(args.repeat_num): 56 | filelist.write('./filelists/train_files_for_val_on_Area_%d_g_%d.txt\n' % (area_idx, list_idx)) 57 | 58 | val_h5 = area_h5s[area_idx - 1] 59 | val_list = os.path.join(root, 'val_files_Area_%d.txt' % area_idx) 60 | print('{}-Saving {}...'.format(datetime.now(), val_list)) 61 | with open(val_list, 'w') as filelist: 62 | for filename_h5 in val_h5: 63 | filelist.write(filename_h5) 64 | 65 | 66 | if __name__ == '__main__': 67 | main() 68 | -------------------------------------------------------------------------------- /data_conversions/prepare_s3dis_label.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/python() 2 | from __future__ import absolute_import 3 | from __future__ import division 4 | from __future__ import print_function 5 | 6 | import argparse 7 | import os 8 | import numpy as np 9 | 10 | DEFAULT_DATA_DIR = '../../data/3DIS/Stanford3dDataset_v1.2_Aligned_Version' 11 | DEFAULT_OUTPUT_DIR = '../../data/3DIS/prepare_label_rgb' 12 | 13 | p = argparse.ArgumentParser() 14 | p.add_argument( 15 | "-d", "--data", dest='data_dir', 16 | default=DEFAULT_DATA_DIR, 17 | help="Path to S3DIS data (default is %s)" % DEFAULT_DATA_DIR) 18 | p.add_argument( 19 | "-f", "--folder", dest='output_dir', 20 | default=DEFAULT_OUTPUT_DIR, 21 | help="Folder to write labels (default is %s)" % DEFAULT_OUTPUT_DIR) 22 | 23 | args = p.parse_args() 24 | 25 | object_dict = { 26 | 'clutter': 0, 27 | 'ceiling': 1, 28 | 'floor': 2, 29 | 'wall': 3, 30 | 'beam': 4, 31 | 'column': 5, 32 | 'door': 6, 33 | 'window': 7, 34 | 'table': 8, 35 | 'chair': 9, 36 | 'sofa': 10, 37 | 'bookcase': 11, 38 | 'board': 12} 39 | 40 | path_dir_areas = os.listdir(args.data_dir) 41 | 42 | for area in path_dir_areas: 43 | path_area = os.path.join(args.data_dir, area) 44 | if not os.path.isdir(path_area): 45 | continue 46 | path_dir_rooms = os.listdir(path_area) 47 | for room in path_dir_rooms: 48 | path_annotations = os.path.join(args.data_dir, area, room, "Annotations") 49 | if not os.path.isdir(path_annotations): 50 | continue 51 | print(path_annotations) 52 | path_prepare_label = os.path.join(args.output_dir, area, room) 53 | if os.path.exists(os.path.join(path_prepare_label, ".labels")): 54 | print("%s already processed, skipping" % path_prepare_label) 55 | continue 56 | xyz_room = np.zeros((1,6)) 57 | label_room = np.zeros((1,1)) 58 | # make store directories 59 | if not os.path.exists(path_prepare_label): 60 | os.makedirs(path_prepare_label) 61 | ############################# 62 | path_objects = os.listdir(path_annotations) 63 | for obj in path_objects: 64 | object_key = obj.split("_", 1)[0] 65 | try: 66 | val = object_dict[object_key] 67 | except KeyError: 68 | continue 69 | print("%s/%s" % (room, obj[:-4])) 70 | xyz_object_path = os.path.join(path_annotations, obj) 71 | try: 72 | xyz_object = np.loadtxt(xyz_object_path)[:,:] # (N,6) 73 | except ValueError as e: 74 | print("ERROR: cannot load %s: %s" % (xyz_object_path, e)) 75 | continue 76 | label_object = np.tile(val, (xyz_object.shape[0], 1)) # (N,1) 77 | xyz_room = np.vstack((xyz_room, xyz_object)) 78 | label_room = np.vstack((label_room, label_object)) 79 | 80 | xyz_room = np.delete(xyz_room, [0], 0) 81 | label_room = np.delete(label_room, [0], 0) 82 | 83 | np.save(path_prepare_label+"/xyzrgb.npy", xyz_room) 84 | np.save(path_prepare_label+"/label.npy", label_room) 85 | 86 | # Marker indicating we've processed this room 87 | open(os.path.join(path_prepare_label, ".labels"), "w").close() 88 | -------------------------------------------------------------------------------- /data_conversions/prepare_scannet_cls_data.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/python3 2 | '''Convert ScanNet pts to h5.''' 3 | 4 | from __future__ import absolute_import 5 | from __future__ import division 6 | from __future__ import print_function 7 | 8 | import os 9 | import h5py 10 | import argparse 11 | import numpy as np 12 | from datetime import datetime 13 | 14 | 15 | def main(): 16 | parser = argparse.ArgumentParser() 17 | parser.add_argument('--folder', '-f', help='Path to data folder') 18 | args = parser.parse_args() 19 | print(args) 20 | 21 | batch_size = 2048 22 | sample_num = 2048 23 | 24 | folder_scanenet = args.folder if args.folder else '../../data/scannet/cls' 25 | train_test_folders = ['train', 'test'] 26 | 27 | label_list = [] 28 | for folder in train_test_folders: 29 | folder_pts = os.path.join(folder_scanenet, folder, 'pts') 30 | for filename in os.listdir(folder_pts): 31 | label_list.append(int(filename[:-4].split('_')[-1])) 32 | label_list = sorted(set(label_list)) 33 | print('label_num:', len(label_list)) 34 | label_dict = dict() 35 | for idx, label in enumerate(label_list): 36 | label_dict[label] = idx 37 | 38 | data = np.zeros((batch_size, sample_num, 6)) 39 | label = np.zeros((batch_size), dtype=np.int32) 40 | for folder in train_test_folders: 41 | folder_pts = os.path.join(folder_scanenet, folder, 'pts') 42 | 43 | idx_h5 = 0 44 | filename_filelist_h5 = os.path.join(folder_scanenet, '%s_files.txt' % folder) 45 | with open(filename_filelist_h5, 'w') as filelist_h5: 46 | filelist = os.listdir(folder_pts) 47 | for idx_pts, filename in enumerate(filelist): 48 | label_object = label_dict[int(filename[:-4].split('_')[-1])] 49 | filename_pts = os.path.join(folder_pts, filename) 50 | xyzrgbs = np.array([[float(value) for value in xyzrgb.split(' ')] 51 | for xyzrgb in open(filename_pts, 'r') if len(xyzrgb.split(' ')) == 6]) 52 | np.random.shuffle(xyzrgbs) 53 | pt_num = xyzrgbs.shape[0] 54 | indices = np.random.choice(pt_num, sample_num, replace=(pt_num < sample_num)) 55 | points_array = xyzrgbs[indices] 56 | points_array[..., 3:] = points_array[..., 3:]/255 - 0.5 # normalize colors 57 | 58 | idx_in_batch = idx_pts % batch_size 59 | data[idx_in_batch, ...] = points_array 60 | label[idx_in_batch] = label_object 61 | if ((idx_pts + 1) % batch_size == 0) or idx_pts == len(filelist) - 1: 62 | item_num = idx_in_batch + 1 63 | filename_h5 = os.path.join(folder_scanenet, '%s_%d.h5' % (folder, idx_h5)) 64 | print('{}-Saving {}...'.format(datetime.now(), filename_h5)) 65 | filelist_h5.write('./%s_%d.h5\n' % (folder, idx_h5)) 66 | 67 | file = h5py.File(filename_h5, 'w') 68 | file.create_dataset('data', data=data[0:item_num, ...]) 69 | file.create_dataset('label', data=label[0:item_num, ...]) 70 | file.close() 71 | 72 | idx_h5 = idx_h5 + 1 73 | 74 | if __name__ == '__main__': 75 | main() 76 | -------------------------------------------------------------------------------- /data_conversions/prepare_scannet_seg_data.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/python3 2 | '''Prepare Data for ScanNet Segmentation Task.''' 3 | 4 | from __future__ import absolute_import 5 | from __future__ import division 6 | from __future__ import print_function 7 | 8 | import os 9 | import sys 10 | import math 11 | import h5py 12 | import pickle 13 | import argparse 14 | import numpy as np 15 | from datetime import datetime 16 | 17 | BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) 18 | sys.path.append(BASE_DIR) 19 | sys.path.append(os.path.join(BASE_DIR, 'utils')) 20 | 21 | import data_utils 22 | 23 | 24 | 25 | def main(): 26 | parser = argparse.ArgumentParser() 27 | parser.add_argument('--folder', '-f', help='Path to data folder') 28 | parser.add_argument('--max_point_num', '-m', help='Max point number of each sample', type=int, default=8192) 29 | parser.add_argument('--block_size', '-b', help='Block size', type=float, default=1.5) 30 | parser.add_argument('--grid_size', '-g', help='Grid size', type=float, default=0.03) 31 | parser.add_argument('--save_ply', '-s', help='Convert .pts to .ply', action='store_true') 32 | 33 | args = parser.parse_args() 34 | print(args) 35 | 36 | root = args.folder if args.folder else '../../data/scannet/seg' 37 | max_point_num = args.max_point_num 38 | 39 | batch_size = 2048 40 | data = np.zeros((batch_size, max_point_num, 3)) 41 | data_num = np.zeros((batch_size), dtype=np.int32) 42 | label = np.zeros((batch_size), dtype=np.int32) 43 | label_seg = np.zeros((batch_size, max_point_num), dtype=np.int32) 44 | indices_split_to_full = np.zeros((batch_size, max_point_num, 2), dtype=np.int32) 45 | 46 | datasets = ['train', 'test'] 47 | for dataset_idx, dataset in enumerate(datasets): 48 | if dataset_idx == 0: 49 | continue 50 | filename = os.path.abspath(os.path.join(root, 'scannet_%s.pickle' % dataset)) 51 | 52 | print('{}-Loading {}...'.format(datetime.now(), filename)) 53 | file_pickle = open(filename, 'rb') 54 | xyz_all = pickle.load(file_pickle, encoding='latin1') 55 | labels_all = pickle.load(file_pickle, encoding='latin1') 56 | file_pickle.close() 57 | 58 | offsets = [('zero', 0.0), ('half', args.block_size / 2)] 59 | for offset_name, offset in offsets: 60 | idx_h5 = 0 61 | idx = 0 62 | 63 | for room_idx, xyz in enumerate(xyz_all): 64 | # align to room bottom center 65 | xyz_min = np.amin(xyz, axis=0, keepdims=True) 66 | xyz_max = np.amax(xyz, axis=0, keepdims=True) 67 | xyz_center = (xyz_min + xyz_max) / 2 68 | xyz_center[0][-1] = xyz_min[0][-1] 69 | xyz = xyz - xyz_center 70 | 71 | labels = labels_all[room_idx] 72 | print('{}-Computing block id of {} points...'.format(datetime.now(), xyz.shape[0])) 73 | xyz_min = np.amin(xyz, axis=0, keepdims=True) - offset 74 | xyz_max = np.amax(xyz, axis=0, keepdims=True) 75 | block_size = (args.block_size, args.block_size, 2 * (xyz_max[0, -1] - xyz_min[0, -1])) 76 | xyz_blocks = np.floor((xyz - xyz_min) / block_size).astype(np.int) 77 | 78 | print('{}-Collecting points belong to each block...'.format(datetime.now(), xyz.shape[0])) 79 | blocks, point_block_indices, block_point_counts = np.unique(xyz_blocks, return_inverse=True, 80 | return_counts=True, axis=0) 81 | block_point_indices = np.split(np.argsort(point_block_indices), np.cumsum(block_point_counts[:-1])) 82 | print('{}-{} is split into {} blocks.'.format(datetime.now(), dataset, blocks.shape[0])) 83 | 84 | block_to_block_idx_map = dict() 85 | for block_idx in range(blocks.shape[0]): 86 | block = (blocks[block_idx][0], blocks[block_idx][1]) 87 | block_to_block_idx_map[(block[0], block[1])] = block_idx 88 | 89 | # merge small blocks into one of their big neighbors 90 | block_point_count_threshold = max_point_num / 10 91 | nbr_block_offsets = [(0, 1), (1, 0), (0, -1), (-1, 0), (-1, 1), (1, 1), (1, -1), (-1, -1)] 92 | block_merge_count = 0 93 | for block_idx in range(blocks.shape[0]): 94 | if block_point_counts[block_idx] >= block_point_count_threshold: 95 | continue 96 | 97 | block = (blocks[block_idx][0], blocks[block_idx][1]) 98 | for x, y in nbr_block_offsets: 99 | nbr_block = (block[0] + x, block[1] + y) 100 | if nbr_block not in block_to_block_idx_map: 101 | continue 102 | 103 | nbr_block_idx = block_to_block_idx_map[nbr_block] 104 | if block_point_counts[nbr_block_idx] < block_point_count_threshold: 105 | continue 106 | 107 | block_point_indices[nbr_block_idx] = np.concatenate( 108 | [block_point_indices[nbr_block_idx], block_point_indices[block_idx]], axis=-1) 109 | block_point_indices[block_idx] = np.array([], dtype=np.int) 110 | block_merge_count = block_merge_count + 1 111 | break 112 | print('{}-{} of {} blocks are merged.'.format(datetime.now(), block_merge_count, blocks.shape[0])) 113 | 114 | idx_last_non_empty_block = 0 115 | for block_idx in reversed(range(blocks.shape[0])): 116 | if block_point_indices[block_idx].shape[0] != 0: 117 | idx_last_non_empty_block = block_idx 118 | break 119 | 120 | # uniformly sample each block 121 | for block_idx in range(idx_last_non_empty_block + 1): 122 | point_indices = block_point_indices[block_idx] 123 | if point_indices.shape[0] == 0: 124 | continue 125 | block_points = xyz[point_indices] 126 | block_min = np.amin(block_points, axis=0, keepdims=True) 127 | xyz_grids = np.floor((block_points - block_min) / args.grid_size).astype(np.int) 128 | grids, point_grid_indices, grid_point_counts = np.unique(xyz_grids, return_inverse=True, 129 | return_counts=True, axis=0) 130 | grid_point_indices = np.split(np.argsort(point_grid_indices), np.cumsum(grid_point_counts[:-1])) 131 | grid_point_count_avg = int(np.average(grid_point_counts)) 132 | point_indices_repeated = [] 133 | for grid_idx in range(grids.shape[0]): 134 | point_indices_in_block = grid_point_indices[grid_idx] 135 | repeat_num = math.ceil(grid_point_count_avg / point_indices_in_block.shape[0]) 136 | if repeat_num > 1: 137 | point_indices_in_block = np.repeat(point_indices_in_block, repeat_num) 138 | np.random.shuffle(point_indices_in_block) 139 | point_indices_in_block = point_indices_in_block[:grid_point_count_avg] 140 | point_indices_repeated.extend(list(point_indices[point_indices_in_block])) 141 | block_point_indices[block_idx] = np.array(point_indices_repeated) 142 | block_point_counts[block_idx] = len(point_indices_repeated) 143 | 144 | for block_idx in range(idx_last_non_empty_block + 1): 145 | point_indices = block_point_indices[block_idx] 146 | if point_indices.shape[0] == 0: 147 | continue 148 | 149 | block_point_num = point_indices.shape[0] 150 | block_split_num = int(math.ceil(block_point_num * 1.0 / max_point_num)) 151 | point_num_avg = int(math.ceil(block_point_num * 1.0 / block_split_num)) 152 | point_nums = [point_num_avg] * block_split_num 153 | point_nums[-1] = block_point_num - (point_num_avg * (block_split_num - 1)) 154 | starts = [0] + list(np.cumsum(point_nums)) 155 | 156 | np.random.shuffle(point_indices) 157 | block_points = xyz[point_indices] 158 | block_labels = labels[point_indices] 159 | x, y, z = np.split(block_points, (1, 2), axis=-1) 160 | block_xzy = np.concatenate([x, z, y], axis=-1) 161 | 162 | for block_split_idx in range(block_split_num): 163 | start = starts[block_split_idx] 164 | point_num = point_nums[block_split_idx] 165 | end = start + point_num 166 | idx_in_batch = idx % batch_size 167 | data[idx_in_batch, 0:point_num, ...] = block_xzy[start:end, :] 168 | data_num[idx_in_batch] = point_num 169 | label[idx_in_batch] = dataset_idx # won't be used... 170 | label_seg[idx_in_batch, 0:point_num] = block_labels[start:end] 171 | 172 | ind_in_room = point_indices[start:end] 173 | indices_split_to_full[idx_in_batch, 0:point_num] = np.stack([np.zeros_like(ind_in_room) + room_idx ,ind_in_room], -1) 174 | 175 | 176 | if ((idx + 1) % batch_size == 0) \ 177 | or (room_idx == len(xyz_all) - 1 178 | and block_idx == idx_last_non_empty_block 179 | and block_split_idx == block_split_num - 1): 180 | item_num = idx_in_batch + 1 181 | filename_h5 = os.path.join(root, dataset, '%s_%d.h5' % (offset_name, idx_h5)) 182 | os.makedirs(os.path.dirname(filename_h5), exist_ok=True) 183 | print('{}-Saving {}...'.format(datetime.now(), filename_h5)) 184 | 185 | file = h5py.File(filename_h5, 'w') 186 | file.create_dataset('data', data=data[0:item_num, ...]) 187 | file.create_dataset('data_num', data=data_num[0:item_num, ...]) 188 | file.create_dataset('label', data=label[0:item_num, ...]) 189 | file.create_dataset('label_seg', data=label_seg[0:item_num, ...]) 190 | file.create_dataset('indices_split_to_full', data=indices_split_to_full[0:item_num, ...]) 191 | file.close() 192 | 193 | if args.save_ply: 194 | print('{}-Saving ply of {}...'.format(datetime.now(), filename_h5)) 195 | filepath_label_ply = os.path.join(root, dataset, 'ply_label', 196 | 'label_%s_%d' % (offset_name, idx_h5)) 197 | data_utils.save_ply_property_batch(data[0:item_num, :, 0:3], 198 | label_seg[0:item_num, ...], 199 | filepath_label_ply, data_num[0:item_num, ...], 22) 200 | 201 | idx_h5 = idx_h5 + 1 202 | idx = idx + 1 203 | 204 | 205 | if __name__ == '__main__': 206 | main() 207 | print('{}-Done.'.format(datetime.now())) 208 | -------------------------------------------------------------------------------- /data_conversions/prepare_scannet_seg_filelists.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/python3 2 | '''Prepare Filelists for ScanNet Segmentation Task.''' 3 | 4 | from __future__ import absolute_import 5 | from __future__ import division 6 | from __future__ import print_function 7 | 8 | import os 9 | import math 10 | import random 11 | import argparse 12 | from datetime import datetime 13 | 14 | 15 | def main(): 16 | parser = argparse.ArgumentParser() 17 | parser.add_argument('--folder', '-f', help='Path to data folder') 18 | parser.add_argument('--h5_num', '-d', help='Number of h5 files to be loaded each time', type=int, default=4) 19 | parser.add_argument('--repeat_num', '-r', help='Number of repeatly using each loaded h5 list', type=int, default=1) 20 | 21 | args = parser.parse_args() 22 | print(args) 23 | 24 | root = args.folder if args.folder else '../../data/scannet/seg/' 25 | 26 | train_h5 = ['./train/%s\n' % (filename) for filename in os.listdir(os.path.join(root, 'train')) 27 | if filename.endswith('.h5')] 28 | random.shuffle(train_h5) 29 | train_list = os.path.join(root, 'train_files.txt') 30 | print('{}-Saving {}...'.format(datetime.now(), train_list)) 31 | with open(train_list, 'w') as filelist: 32 | list_num = math.ceil(len(train_h5) / args.h5_num) 33 | for list_idx in range(list_num): 34 | train_list_i = os.path.join(root, 'filelists', 'train_files_g_%d.txt' % (list_idx)) 35 | os.makedirs(os.path.dirname(train_list_i), exist_ok=True) 36 | with open(train_list_i, 'w') as filelist_i: 37 | for h5_idx in range(args.h5_num): 38 | filename_idx = list_idx * args.h5_num + h5_idx 39 | if filename_idx > len(train_h5) - 1: 40 | break 41 | filename_h5 = train_h5[filename_idx] 42 | filelist_i.write('../' + filename_h5) 43 | for repeat_idx in range(args.repeat_num): 44 | filelist.write('./filelists/train_files_g_%d.txt\n' % (list_idx)) 45 | 46 | test_h5 = ['./test/%s\n' % (filename) for filename in os.listdir(os.path.join(root, 'test')) 47 | if filename.endswith('.h5')] 48 | test_list = os.path.join(root, 'test_files.txt') 49 | print('{}-Saving {}...'.format(datetime.now(), test_list)) 50 | with open(test_list, 'w') as filelist: 51 | for filename_h5 in test_h5: 52 | filelist.write(filename_h5) 53 | 54 | 55 | if __name__ == '__main__': 56 | main() 57 | -------------------------------------------------------------------------------- /data_conversions/prepare_semantic3d_data.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/python3 2 | '''Prepare Data for Semantic3D Segmentation Task.''' 3 | 4 | from __future__ import absolute_import 5 | from __future__ import division 6 | from __future__ import print_function 7 | 8 | import os 9 | import sys 10 | import math 11 | import h5py 12 | import argparse 13 | import numpy as np 14 | import pandas as pd 15 | from datetime import datetime 16 | 17 | def read_points(f): 18 | # reads Semantic3D .txt file f into a pandas dataframe 19 | col_names = ['x', 'y', 'z', 'i', 'r', 'g', 'b'] 20 | col_dtype = {'x': np.float32, 'y': np.float32, 'z': np.float32, 'i': np.int32, 21 | 'r': np.uint8, 'g': np.uint8, 'b': np.uint8} 22 | xyzirgb = pd.read_csv(f, names=col_names, dtype=col_dtype, delim_whitespace=True) 23 | xyzirgb = xyzirgb.to_numpy() 24 | return xyzirgb[:,0:3], xyzirgb[:,4:7] # xyz, rgb 25 | 26 | def load_labels(label_path): 27 | # Assuming each line is a valid int 28 | with open(label_path, "r") as f: 29 | labels = [int(line) for line in f] 30 | return np.array(labels, dtype=np.int32) 31 | 32 | def main(): 33 | parser = argparse.ArgumentParser() 34 | parser.add_argument('--folder', '-f', help='Path to data folder') 35 | parser.add_argument('--max_point_num', '-m', help='Max point number of each sample', type=int, default=8192) 36 | parser.add_argument('--block_size', '-b', help='Block size', type=float, default=5.0) 37 | parser.add_argument('--grid_size', '-g', help='Grid size', type=float, default=0.1) 38 | 39 | args = parser.parse_args() 40 | print(args) 41 | 42 | root = args.folder if args.folder else '../../data/semantic3d/' 43 | max_point_num = args.max_point_num 44 | 45 | batch_size = 2048 46 | data = np.zeros((batch_size, max_point_num, 6)) 47 | data_num = np.zeros((batch_size), dtype=np.int32) 48 | label = np.zeros((batch_size), dtype=np.int32) 49 | label_seg = np.zeros((batch_size, max_point_num), dtype=np.int32) 50 | indices_split_to_full = np.zeros((batch_size, max_point_num), dtype=np.int32) 51 | 52 | 53 | folders = [os.path.join(root, folder) for folder in ['train', 'val', 'test_reduced']] 54 | for folder in folders: 55 | datasets = [filename[:-4] for filename in os.listdir(folder) if filename.endswith('.txt')] 56 | for dataset_idx, dataset in enumerate(datasets): 57 | h5_file_name = os.path.join(folder,dataset + '_half_0.h5') 58 | if os.path.isfile(h5_file_name): 59 | continue 60 | 61 | filename_txt = os.path.join(folder, dataset + '.txt') 62 | # print('{}-Loading {}...'.format(datetime.now(), filename_txt)) 63 | # xyzirgb = np.loadtxt(filename_txt) 64 | # print('{}-Loading finished{}...'.format(datetime.now(), filename_txt)) 65 | 66 | print('{}-Loading {}...'.format(datetime.now(), filename_txt)) 67 | [xyz, rgb] = read_points(filename_txt) 68 | print('{}-Loading finished{}...'.format(datetime.now(), filename_txt)) 69 | 70 | filename_labels = os.path.join(folder, dataset + '.labels') 71 | has_labels = os.path.exists(filename_labels) 72 | if has_labels: 73 | print('{}-Loading {}...'.format(datetime.now(), filename_labels)) 74 | labels = load_labels(filename_labels) 75 | indices = (labels != 0) 76 | labels = labels[indices] - 1 # since labels == 0 have been removed 77 | xyz = xyz[indices, :] 78 | rgb = rgb[indices, :] 79 | else: 80 | labels = np.zeros((xyz.shape[0])) 81 | 82 | offsets = [('zero', 0.0), ('half', args.block_size / 2)] 83 | for offset_name, offset in offsets: 84 | idx_h5 = 0 85 | idx = 0 86 | 87 | print('{}-Computing block id of {} points...'.format(datetime.now(), xyz.shape[0])) 88 | xyz_min = np.amin(xyz, axis=0, keepdims=True) - offset 89 | xyz_max = np.amax(xyz, axis=0, keepdims=True) 90 | block_size = (args.block_size, args.block_size, 2 * (xyz_max[0, -1] - xyz_min[0, -1])) 91 | xyz_blocks = np.floor((xyz - xyz_min) / block_size).astype(np.int) 92 | 93 | print('{}-Collecting points belong to each block...'.format(datetime.now(), xyz.shape[0])) 94 | blocks, point_block_indices, block_point_counts = np.unique(xyz_blocks, return_inverse=True, 95 | return_counts=True, axis=0) 96 | block_point_indices = np.split(np.argsort(point_block_indices), np.cumsum(block_point_counts[:-1])) 97 | print('{}-{} is split into {} blocks.'.format(datetime.now(), dataset, blocks.shape[0])) 98 | 99 | block_to_block_idx_map = dict() 100 | for block_idx in range(blocks.shape[0]): 101 | block = (blocks[block_idx][0], blocks[block_idx][1]) 102 | block_to_block_idx_map[(block[0], block[1])] = block_idx 103 | 104 | # merge small blocks into one of their big neighbors 105 | block_point_count_threshold = max_point_num / 10 106 | nbr_block_offsets = [(0, 1), (1, 0), (0, -1), (-1, 0), (-1, 1), (1, 1), (1, -1), (-1, -1)] 107 | block_merge_count = 0 108 | for block_idx in range(blocks.shape[0]): 109 | if block_point_counts[block_idx] >= block_point_count_threshold: 110 | continue 111 | 112 | block = (blocks[block_idx][0], blocks[block_idx][1]) 113 | for x, y in nbr_block_offsets: 114 | nbr_block = (block[0] + x, block[1] + y) 115 | if nbr_block not in block_to_block_idx_map: 116 | continue 117 | 118 | nbr_block_idx = block_to_block_idx_map[nbr_block] 119 | if block_point_counts[nbr_block_idx] < block_point_count_threshold: 120 | continue 121 | 122 | block_point_indices[nbr_block_idx] = np.concatenate( 123 | [block_point_indices[nbr_block_idx], block_point_indices[block_idx]], axis=-1) 124 | block_point_indices[block_idx] = np.array([], dtype=np.int) 125 | block_merge_count = block_merge_count + 1 126 | break 127 | print('{}-{} of {} blocks are merged.'.format(datetime.now(), block_merge_count, blocks.shape[0])) 128 | 129 | idx_last_non_empty_block = 0 130 | for block_idx in reversed(range(blocks.shape[0])): 131 | if block_point_indices[block_idx].shape[0] != 0: 132 | idx_last_non_empty_block = block_idx 133 | break 134 | 135 | # uniformly sample each block 136 | for block_idx in range(idx_last_non_empty_block + 1): 137 | point_indices = block_point_indices[block_idx] 138 | if point_indices.shape[0] == 0: 139 | continue 140 | block_points = xyz[point_indices] 141 | block_min = np.amin(block_points, axis=0, keepdims=True) 142 | xyz_grids = np.floor((block_points - block_min) / args.grid_size).astype(np.int) 143 | grids, point_grid_indices, grid_point_counts = np.unique(xyz_grids, return_inverse=True, 144 | return_counts=True, axis=0) 145 | grid_point_indices = np.split(np.argsort(point_grid_indices), np.cumsum(grid_point_counts[:-1])) 146 | grid_point_count_avg = int(np.average(grid_point_counts)) 147 | point_indices_repeated = [] 148 | for grid_idx in range(grids.shape[0]): 149 | point_indices_in_block = grid_point_indices[grid_idx] 150 | repeat_num = math.ceil(grid_point_count_avg / point_indices_in_block.shape[0]) 151 | if repeat_num > 1: 152 | point_indices_in_block = np.repeat(point_indices_in_block, repeat_num) 153 | np.random.shuffle(point_indices_in_block) 154 | point_indices_in_block = point_indices_in_block[:grid_point_count_avg] 155 | point_indices_repeated.extend(list(point_indices[point_indices_in_block])) 156 | block_point_indices[block_idx] = np.array(point_indices_repeated) 157 | block_point_counts[block_idx] = len(point_indices_repeated) 158 | 159 | for block_idx in range(idx_last_non_empty_block + 1): 160 | point_indices = block_point_indices[block_idx] 161 | if point_indices.shape[0] == 0: 162 | continue 163 | 164 | block_point_num = point_indices.shape[0] 165 | block_split_num = int(math.ceil(block_point_num * 1.0 / max_point_num)) 166 | point_num_avg = int(math.ceil(block_point_num * 1.0 / block_split_num)) 167 | point_nums = [point_num_avg] * block_split_num 168 | point_nums[-1] = block_point_num - (point_num_avg * (block_split_num - 1)) 169 | starts = [0] + list(np.cumsum(point_nums)) 170 | 171 | np.random.shuffle(point_indices) 172 | block_points = xyz[point_indices] 173 | block_min = np.amin(block_points, axis=0, keepdims=True) 174 | block_max = np.amax(block_points, axis=0, keepdims=True) 175 | block_center = (block_min + block_max) / 2 176 | block_center[0][-1] = block_min[0][-1] 177 | block_points = block_points - block_center # align to block bottom center 178 | x, y, z = np.split(block_points, (1, 2), axis=-1) 179 | block_xzyrgbi = np.concatenate([x, z, y, rgb[point_indices]], axis=-1) 180 | block_labels = labels[point_indices] 181 | 182 | for block_split_idx in range(block_split_num): 183 | start = starts[block_split_idx] 184 | point_num = point_nums[block_split_idx] 185 | end = start + point_num 186 | idx_in_batch = idx % batch_size 187 | data[idx_in_batch, 0:point_num, ...] = block_xzyrgbi[start:end, :] 188 | data_num[idx_in_batch] = point_num 189 | label[idx_in_batch] = dataset_idx # won't be used... 190 | label_seg[idx_in_batch, 0:point_num] = block_labels[start:end] 191 | indices_split_to_full[idx_in_batch, 0:point_num] = point_indices[start:end] 192 | 193 | 194 | if ((idx + 1) % batch_size == 0) or \ 195 | (block_idx == idx_last_non_empty_block and block_split_idx == block_split_num - 1): 196 | item_num = idx_in_batch + 1 197 | filename_h5 = os.path.join(folder, dataset + '_%s_%d.h5' % (offset_name, idx_h5)) 198 | print('{}-Saving {}...'.format(datetime.now(), filename_h5)) 199 | 200 | file = h5py.File(filename_h5, 'w') 201 | file.create_dataset('data', data=data[0:item_num, ...]) 202 | file.create_dataset('data_num', data=data_num[0:item_num, ...]) 203 | file.create_dataset('label', data=label[0:item_num, ...]) 204 | file.create_dataset('label_seg', data=label_seg[0:item_num, ...]) 205 | file.create_dataset('indices_split_to_full', data=indices_split_to_full[0:item_num, ...]) 206 | file.close() 207 | 208 | idx_h5 = idx_h5 + 1 209 | idx = idx + 1 210 | 211 | 212 | if __name__ == '__main__': 213 | main() 214 | print('{}-Done.'.format(datetime.now())) -------------------------------------------------------------------------------- /data_conversions/prepare_semantic3d_filelists.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/python3 2 | '''Prepare Filelists for Semantic3D Segmentation Task.''' 3 | 4 | from __future__ import absolute_import 5 | from __future__ import division 6 | from __future__ import print_function 7 | 8 | import os 9 | import math 10 | import random 11 | import argparse 12 | from datetime import datetime 13 | 14 | 15 | def main(): 16 | parser = argparse.ArgumentParser() 17 | parser.add_argument('--folder', '-f', help='Path to data folder') 18 | parser.add_argument('--h5_num', '-d', help='Number of h5 files to be loaded each time', type=int, default=4) 19 | parser.add_argument('--repeat_num', '-r', help='Number of repeatly using each loaded h5 list', type=int, default=1) 20 | 21 | args = parser.parse_args() 22 | print(args) 23 | 24 | root = args.folder if args.folder else '../../data/semantic3d/downsampled' 25 | 26 | splits = ['train', 'val', 'test', 'test_reduced'] 27 | split_filelists = dict() 28 | for split in splits: 29 | split_filelists[split] = ['./%s/%s\n' % (split, filename) for filename in os.listdir(os.path.join(root, split)) 30 | if filename.endswith('.h5')] 31 | 32 | train_h5 = split_filelists['train'] 33 | random.shuffle(train_h5) 34 | train_list = os.path.join(root, 'train_data_files.txt') 35 | print('{}-Saving {}...'.format(datetime.now(), train_list)) 36 | with open(train_list, 'w') as filelist: 37 | list_num = math.ceil(len(train_h5) / args.h5_num) 38 | for list_idx in range(list_num): 39 | train_list_i = os.path.join(root, 'filelists', 'train_files_g_%d.txt' % list_idx) 40 | with open(train_list_i, 'w') as filelist_i: 41 | for h5_idx in range(args.h5_num): 42 | filename_idx = list_idx * args.h5_num + h5_idx 43 | if filename_idx > len(train_h5) - 1: 44 | break 45 | filename_h5 = train_h5[filename_idx] 46 | filelist_i.write('../' + filename_h5) 47 | for repeat_idx in range(args.repeat_num): 48 | filelist.write('./filelists/train_files_g_%d.txt\n' % list_idx) 49 | 50 | val_h5 = split_filelists['val'] 51 | val_list = os.path.join(root, 'val_data_files.txt') 52 | print('{}-Saving {}...'.format(datetime.now(), val_list)) 53 | with open(val_list, 'w') as filelist: 54 | for filename_h5 in val_h5: 55 | filelist.write(filename_h5) 56 | 57 | test_h5 = split_filelists['test'] 58 | test_list = os.path.join(root, 'test_files.txt') 59 | print('{}-Saving {}...'.format(datetime.now(), test_list)) 60 | with open(test_list, 'w') as filelist: 61 | for filename_h5 in test_h5: 62 | filelist.write(filename_h5) 63 | 64 | test_h5 = split_filelists['test_reduced'] 65 | test_list = os.path.join(root, 'test_reduced_files.txt') 66 | print('{}-Saving {}...'.format(datetime.now(), test_list)) 67 | with open(test_list, 'w') as filelist: 68 | for filename_h5 in test_h5: 69 | filelist.write(filename_h5) 70 | 71 | 72 | if __name__ == '__main__': 73 | main() 74 | -------------------------------------------------------------------------------- /data_conversions/preprocess_semantic3d_data.py: -------------------------------------------------------------------------------- 1 | import os 2 | import subprocess 3 | import shutil 4 | import open3d 5 | 6 | from dataset.semantic_dataset import all_file_prefixes 7 | 8 | 9 | def wc(file_name): 10 | out = subprocess.Popen( 11 | ["wc", "-l", file_name], stdout=subprocess.PIPE, stderr=subprocess.STDOUT 12 | ).communicate()[0] 13 | return int(out.partition(b" ")[0]) 14 | 15 | 16 | def prepend_line(file_name, line): 17 | with open(file_name, "r+") as f: 18 | content = f.read() 19 | f.seek(0, 0) 20 | f.write(line.rstrip("\r\n") + "\n" + content) 21 | 22 | 23 | def point_cloud_txt_to_pcd(raw_dir, file_prefix): 24 | # File names 25 | txt_file = os.path.join(raw_dir, file_prefix + ".txt") 26 | pts_file = os.path.join(raw_dir, file_prefix + ".pts") 27 | pcd_file = os.path.join(raw_dir, file_prefix + ".pcd") 28 | 29 | # Skip if already done 30 | if os.path.isfile(pcd_file): 31 | print("pcd {} exists, skipped".format(pcd_file)) 32 | return 33 | 34 | # .txt to .pts 35 | # We could just prepend the line count, however, there are some intensity value 36 | # which are non-integers. 37 | print("[txt->pts]") 38 | print("txt: {}".format(txt_file)) 39 | print("pts: {}".format(pts_file)) 40 | with open(txt_file, "r") as txt_f, open(pts_file, "w") as pts_f: 41 | for line in txt_f: 42 | # x, y, z, i, r, g, b 43 | tokens = line.split() 44 | tokens[3] = str(int(float(tokens[3]))) 45 | line = " ".join(tokens) 46 | pts_f.write(line + "\n") 47 | prepend_line(pts_file, str(wc(txt_file))) 48 | 49 | # .pts -> .pcd 50 | print("[pts->pcd]") 51 | print("pts: {}".format(pts_file)) 52 | print("pcd: {}".format(pcd_file)) 53 | point_cloud = open3d.read_point_cloud(pts_file) 54 | open3d.write_point_cloud(pcd_file, point_cloud) 55 | os.remove(pts_file) 56 | 57 | 58 | if __name__ == "__main__": 59 | # By default 60 | # raw data: "dataset/semantic_raw" 61 | current_dir = os.path.dirname(os.path.realpath(__file__)) 62 | dataset_dir = os.path.join(current_dir, "dataset") 63 | raw_dir = os.path.join(dataset_dir, "semantic_raw") 64 | 65 | for file_prefix in all_file_prefixes: 66 | point_cloud_txt_to_pcd(raw_dir, file_prefix) 67 | -------------------------------------------------------------------------------- /data_conversions/un7z_semantic3d.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash -e 2 | BASE_DIR=${1-../../data/semantic3d} 3 | 4 | # Helper function to skip unpacking if already unpacked. Uses markers 5 | # to indicate when a file is successfully unpacked. 6 | unpack() { 7 | local path=${1} 8 | local marker=$path.unpacked 9 | if [ -e $marker ]; then 10 | echo "$path already unpacked, skipping" 11 | return 12 | fi 13 | 7z x $path -o$(dirname $path) -y 14 | touch $marker 15 | } 16 | 17 | # Training data 18 | unpack $BASE_DIR/train/bildstein_station1_xyz_intensity_rgb.7z 19 | unpack $BASE_DIR/train/bildstein_station5_xyz_intensity_rgb.7z 20 | unpack $BASE_DIR/train/domfountain_station1_xyz_intensity_rgb.7z 21 | unpack $BASE_DIR/train/domfountain_station3_xyz_intensity_rgb.7z 22 | unpack $BASE_DIR/train/neugasse_station1_xyz_intensity_rgb.7z 23 | unpack $BASE_DIR/train/sg27_station1_intensity_rgb.7z 24 | unpack $BASE_DIR/train/sg27_station2_intensity_rgb.7z 25 | unpack $BASE_DIR/train/sg27_station5_intensity_rgb.7z 26 | unpack $BASE_DIR/train/sg27_station9_intensity_rgb.7z 27 | unpack $BASE_DIR/train/sg28_station4_intensity_rgb.7z 28 | unpack $BASE_DIR/train/untermaederbrunnen_station1_xyz_intensity_rgb.7z 29 | unpack $BASE_DIR/train/sem8_labels_training.7z 30 | 31 | # Validation data 32 | unpack $BASE_DIR/val/bildstein_station3_xyz_intensity_rgb.7z 33 | unpack $BASE_DIR/val/domfountain_station2_xyz_intensity_rgb.7z 34 | unpack $BASE_DIR/val/sg27_station4_intensity_rgb.7z 35 | unpack $BASE_DIR/val/untermaederbrunnen_station3_xyz_intensity_rgb.7z 36 | 37 | # Testing data 38 | unpack $BASE_DIR/test_reduced/MarketplaceFeldkirch_Station4_rgb_intensity-reduced.txt.7z 39 | unpack $BASE_DIR/test_reduced/StGallenCathedral_station6_rgb_intensity-reduced.txt.7z 40 | unpack $BASE_DIR/test_reduced/sg27_station10_rgb_intensity-reduced.txt.7z 41 | unpack $BASE_DIR/test_reduced/sg28_Station2_rgb_intensity-reduced.txt.7z 42 | -------------------------------------------------------------------------------- /evaluation/eval_s3dis.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/python3 2 | 3 | from __future__ import absolute_import 4 | from __future__ import division 5 | from __future__ import print_function 6 | 7 | import argparse 8 | import os 9 | 10 | import numpy as np 11 | 12 | 13 | gt_label_filenames = [] 14 | pred_label_filenames = [] 15 | 16 | DEFAULT_DATA_DIR = '../../data/3DIS/prepare_label_rgb/' 17 | NUM_CLASSES = 13 18 | 19 | p = argparse.ArgumentParser() 20 | p.add_argument( 21 | "-d", "--data", dest='data_dir', 22 | default=DEFAULT_DATA_DIR, 23 | help="Path to S3DIS data (default is %s)" % DEFAULT_DATA_DIR) 24 | 25 | args = p.parse_args() 26 | 27 | for area in os.listdir(args.data_dir): 28 | path_area = os.path.join(args.data_dir, area) 29 | if not os.path.isdir(path_area): 30 | continue 31 | Rooms = os.listdir(path_area) 32 | for room in Rooms: 33 | path_room = os.path.join(path_area, room) 34 | if not os.path.isdir(path_room): 35 | continue 36 | path_gt_label = os.path.join(path_room, 'label.npy') 37 | if not os.path.exists(path_gt_label): 38 | print("%s does not exist, skipping" % path_gt_label) 39 | continue 40 | path_pred_label = os.path.join(path_room, 'pred.npy') 41 | if not os.path.exists(path_pred_label): 42 | print("%s does not exist, skipping" % path_pred_label) 43 | continue 44 | pred_label_filenames.append(path_pred_label) 45 | gt_label_filenames.append(path_gt_label) 46 | 47 | num_room = len(gt_label_filenames) 48 | num_preds = len(pred_label_filenames) 49 | assert num_room == num_preds 50 | 51 | print("Found {} predictions".format(num_room)) 52 | 53 | gt_classes = [0] * NUM_CLASSES 54 | positive_classes = [0] * NUM_CLASSES 55 | true_positive_classes = [0] * NUM_CLASSES 56 | 57 | print("Evaluating predictions:") 58 | for i in range(num_room): 59 | print(" {} ({}/{})".format(pred_label_filenames[i], i + 1, num_room)) 60 | pred_label = np.loadtxt(pred_label_filenames[i]) 61 | gt_label = np.load(gt_label_filenames[i]) 62 | for j in range(gt_label.shape[0]): 63 | gt_l = int(gt_label[j]) 64 | pred_l = int(pred_label[j]) 65 | gt_classes[gt_l] += 1 66 | positive_classes[pred_l] += 1 67 | true_positive_classes[gt_l] += int(gt_l==pred_l) 68 | 69 | print("Classes:\t{}".format("\t".join(map(str, gt_classes)))) 70 | print("Positive:\t{}".format("\t".join(map(str, positive_classes)))) 71 | print("True positive:\t{}".format("\t".join(map(str, true_positive_classes)))) 72 | print("Overall accuracy: {0}".format(sum(true_positive_classes)/float(sum(positive_classes)))) 73 | 74 | print("Class IoU:") 75 | iou_list = [] 76 | for i in range(13): 77 | iou = true_positive_classes[i]/float(gt_classes[i]+positive_classes[i]-true_positive_classes[i]) 78 | print(" {}: {}".format(i, iou)) 79 | iou_list.append(iou) 80 | 81 | print("Average IoU: {}".format(sum(iou_list)/13.0)) 82 | -------------------------------------------------------------------------------- /evaluation/eval_scannet.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/python3 2 | """Merge blocks and evaluate scannet""" 3 | 4 | from __future__ import absolute_import 5 | from __future__ import division 6 | from __future__ import print_function 7 | 8 | import os,sys 9 | import plyfile 10 | import numpy as np 11 | import argparse 12 | import h5py 13 | import pickle 14 | 15 | def main(): 16 | parser = argparse.ArgumentParser() 17 | parser.add_argument('--datafolder', '-d', help='Path to input *_pred.h5', required=True) 18 | parser.add_argument('--picklefile', '-p', help='Path to scannet_test.pickle', required=True) 19 | args = parser.parse_args() 20 | print(args) 21 | 22 | file_list = os.listdir(args.datafolder) 23 | pred_list = [pred for pred in file_list if pred.split(".")[-1] == "h5" and "pred" in pred] 24 | 25 | pts_acc_list = [] 26 | vox_acc_list = [] 27 | 28 | #load scannet_test.pickle file 29 | file_pickle = open(args.picklefile, 'rb') 30 | xyz_all = pickle.load(file_pickle, encoding='latin1') # encoding keyword for python3 31 | labels_all = pickle.load(file_pickle, encoding='latin1') 32 | file_pickle.close() 33 | 34 | pickle_dict = {} 35 | for room_idx, xyz in enumerate(xyz_all): 36 | 37 | room_pt_num = xyz.shape[0] 38 | room_dict = {} 39 | 40 | room_dict["merged_label_zero"] = np.zeros((room_pt_num),dtype=int) 41 | room_dict["merged_confidence_zero"] = np.zeros((room_pt_num),dtype=float) 42 | room_dict["merged_label_half"] = np.zeros((room_pt_num), dtype=int) 43 | room_dict["merged_confidence_half"] = np.zeros((room_pt_num), dtype=float) 44 | room_dict["final_label"] = np.zeros((room_pt_num), dtype=int) 45 | 46 | pickle_dict[room_idx] = room_dict 47 | 48 | # load block preds and merge them to room scene 49 | for pred_file in pred_list: 50 | 51 | print("process:", os.path.join(args.datafolder, pred_file)) 52 | test_file = pred_file.replace("_pred","") 53 | 54 | # load pred .h5 55 | data_pred = h5py.File(os.path.join(args.datafolder, pred_file)) 56 | 57 | pred_labels_seg = data_pred['label_seg'][...].astype(np.int64) 58 | pred_indices = data_pred['indices_split_to_full'][...].astype(np.int64) 59 | pred_confidence = data_pred['confidence'][...].astype(np.float32) 60 | pred_data_num = data_pred['data_num'][...].astype(np.int64) 61 | 62 | 63 | if 'zero' in pred_file: 64 | for b_id in range(pred_labels_seg.shape[0]): 65 | indices_b = pred_indices[b_id] 66 | for p_id in range(pred_data_num[b_id]): 67 | room_indices = indices_b[p_id][0] 68 | inroom_indices = indices_b[p_id][1] 69 | pickle_dict[room_indices]["merged_label_zero"][inroom_indices] = pred_labels_seg[b_id][p_id] 70 | pickle_dict[room_indices]["merged_confidence_zero"][inroom_indices] = pred_confidence[b_id][p_id] 71 | else: 72 | for b_id in range(pred_labels_seg.shape[0]): 73 | indices_b = pred_indices[b_id] 74 | for p_id in range(pred_data_num[b_id]): 75 | room_indices = indices_b[p_id][0] 76 | inroom_indices = indices_b[p_id][1] 77 | pickle_dict[room_indices]["merged_label_half"][inroom_indices] = pred_labels_seg[b_id][p_id] 78 | pickle_dict[room_indices]["merged_confidence_half"][inroom_indices] = pred_confidence[b_id][p_id] 79 | 80 | for room_id in pickle_dict.keys(): 81 | 82 | final_label = pickle_dict[room_id]["final_label"] 83 | merged_label_zero = pickle_dict[room_id]["merged_label_zero"] 84 | merged_label_half = pickle_dict[room_id]["merged_label_half"] 85 | merged_confidence_zero = pickle_dict[room_id]["merged_confidence_zero"] 86 | merged_confidence_half = pickle_dict[room_id]["merged_confidence_half"] 87 | 88 | final_label[merged_confidence_zero >= merged_confidence_half] = merged_label_zero[merged_confidence_zero >= merged_confidence_half] 89 | final_label[merged_confidence_zero < merged_confidence_half] = merged_label_half[merged_confidence_zero < merged_confidence_half] 90 | 91 | # eval 92 | for room_id, pts in enumerate(xyz_all): 93 | 94 | label = labels_all[room_id] 95 | pred = pickle_dict[room_id]["final_label"] 96 | data_num = pts.shape[0] 97 | 98 | # compute pts acc (ignore label 0 which is scannet unannotated) 99 | c_accpts = np.sum(np.equal(pred,label)) 100 | c_ignore = np.sum(np.equal(label,0)) 101 | pts_acc_list.append([c_accpts, data_num - c_ignore]) 102 | 103 | # compute voxel accuracy (follow scannet and pointnet++) 104 | res = 0.0484 105 | coordmax = np.max(pts, axis=0) 106 | coordmin = np.min(pts, axis=0) 107 | nvox = np.ceil((coordmax - coordmin) / res) 108 | vidx = np.ceil((pts - coordmin) / res) 109 | vidx = vidx[:, 0] + vidx[:, 1] * nvox[0] + vidx[:, 2] * nvox[0] * nvox[1] 110 | uvidx, vpidx = np.unique(vidx, return_index=True) 111 | 112 | # compute voxel label 113 | uvlabel = np.array(label)[vpidx] 114 | 115 | # compute voxel pred (follow pointnet++ majority voting) 116 | uvpred_tp = [] 117 | label_pred_dict = {} 118 | 119 | for uidx in uvidx: 120 | label_pred_dict[int(uidx)] = [] 121 | for k, p in enumerate(pred): 122 | label_pred_dict[int(vidx[k])].append(p) 123 | for uidx in uvidx: 124 | uvpred_tp.append(np.argmax(np.bincount(label_pred_dict[int(uidx)]))) 125 | 126 | # compute voxel accuracy (ignore label 0 which is scannet unannotated) 127 | c_accvox = np.sum(np.equal(uvpred_tp, uvlabel)) 128 | c_ignore = np.sum(np.equal(uvlabel,0)) 129 | 130 | vox_acc_list.append([c_accvox, (len(uvlabel) - c_ignore)]) 131 | 132 | # compute avg pts acc 133 | pts_acc_sum = np.sum(pts_acc_list,0) 134 | print("pts acc", pts_acc_sum[0]*1.0/pts_acc_sum[1]) 135 | 136 | #compute avg voxel acc 137 | vox_acc_sum = np.sum(vox_acc_list,0) 138 | print("voxel acc", vox_acc_sum[0]*1.0/vox_acc_sum[1]) 139 | 140 | if __name__ == '__main__': 141 | main() 142 | -------------------------------------------------------------------------------- /evaluation/eval_shapenet_seg.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/python3 2 | """Calculate IoU of part segmentation task.""" 3 | 4 | from __future__ import absolute_import 5 | from __future__ import division 6 | from __future__ import print_function 7 | 8 | import os 9 | import sys 10 | import argparse 11 | import numpy as np 12 | 13 | BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) 14 | sys.path.append(BASE_DIR) 15 | sys.path.append(os.path.join(BASE_DIR, 'utils')) 16 | sys.path.append(os.path.join(BASE_DIR, 'setting')) 17 | from utils import data_utils 18 | 19 | 20 | 21 | def main(): 22 | parser = argparse.ArgumentParser() 23 | parser.add_argument('--folder_gt', '-g', help='Path to ground truth folder', required=True) 24 | parser.add_argument('--folder_pred', '-p', help='Path to prediction folder', required=True) 25 | parser.add_argument('--folder_data', '-d', help='Path to point cloud data folder') 26 | parser.add_argument('--part_avg', '-a', action='store_true', help='Use part level average') 27 | args = parser.parse_args() 28 | print(args) 29 | 30 | category_id_to_name = { 31 | 2691156: 'Airplane', 32 | 2773838: 'Bag', 33 | 2954340: 'Cap', 34 | 2958343: 'Car', 35 | 3001627: 'Chair', 36 | 3261776: 'Earphone', 37 | 3467517: 'Guitar', 38 | 3624134: 'Knife', 39 | 3636649: 'Lamp', 40 | 3642806: 'Laptop', 41 | 3790512: 'Motorbike', 42 | 3797390: 'Mug', 43 | 3948459: 'Pistol', 44 | 4099429: 'Rocket', 45 | 4225987: 'Skateboard', 46 | 4379243: 'Table'} 47 | 48 | categories = sorted(os.listdir(args.folder_gt)) 49 | 50 | label_min = sys.maxsize 51 | for category in categories: 52 | category_folder_gt = os.path.join(args.folder_gt, category) 53 | filenames = sorted(os.listdir(category_folder_gt)) 54 | for filename in filenames: 55 | filepath_gt = os.path.join(category_folder_gt, filename) 56 | label_gt = np.loadtxt(filepath_gt).astype(np.int32) 57 | label_min = min(label_min, np.amin(label_gt)) 58 | 59 | IoU = 0.0 60 | total_num = 0 61 | mIoU = 0.0 62 | for category in categories: 63 | category_folder_gt = os.path.join(args.folder_gt, category) 64 | category_folder_pred = os.path.join(args.folder_pred, category) 65 | if args.folder_data: 66 | category_folder_data = os.path.join(args.folder_data, category) 67 | category_folder_err = os.path.join(args.folder_pred+'_err_ply', category) 68 | 69 | IoU_category = 0.0 70 | filenames = sorted(os.listdir(category_folder_gt)) 71 | for filename in filenames: 72 | filepath_gt = os.path.join(category_folder_gt, filename) 73 | filepath_pred = os.path.join(category_folder_pred, filename) 74 | label_gt = np.loadtxt(filepath_gt).astype(np.int32) - label_min 75 | label_pred = np.loadtxt(filepath_pred).astype(np.int32) 76 | 77 | if args.folder_data: 78 | filepath_data = os.path.join(category_folder_data, filename[:-3]+'pts') 79 | filepath_err = os.path.join(category_folder_err, filename[:-3] + 'ply') 80 | coordinates = [[float(value) for value in xyz.split(' ')] 81 | for xyz in open(filepath_data, 'r') if len(xyz.split(' ')) == 3] 82 | assert (label_gt.shape[0] == len(coordinates)) 83 | data_utils.save_ply_property(np.array(coordinates), (label_gt == label_pred), 6, filepath_err) 84 | 85 | if args.part_avg: 86 | label_max = np.amax(label_gt) 87 | IoU_part = 0.0 88 | for label_idx in range(label_max+1): 89 | locations_gt = (label_gt == label_idx) 90 | locations_pred = (label_pred == label_idx) 91 | I_locations = np.logical_and(locations_gt, locations_pred) 92 | U_locations = np.logical_or(locations_gt, locations_pred) 93 | I = np.sum(I_locations) + np.finfo(np.float32).eps 94 | U = np.sum(U_locations) + np.finfo(np.float32).eps 95 | IoU_part = IoU_part + I/U 96 | IoU_sample = IoU_part / (label_max+1) 97 | else: 98 | label_correct_locations = (label_gt == label_pred) 99 | IoU_sample = np.sum(label_correct_locations) / label_gt.size 100 | IoU_category = IoU_category + IoU_sample 101 | IoU = IoU + IoU_category 102 | IoU_category = IoU_category / len(filenames) 103 | if category.isdigit(): 104 | print("IoU of %s: " % (category_id_to_name[int(category)]), IoU_category) 105 | else: 106 | print("IoU of %s: " % category, IoU_category) 107 | mIoU = mIoU + IoU_category 108 | total_num = total_num + len(filenames) 109 | IoU = IoU / total_num 110 | print("IoU: ", IoU) 111 | mIoU = mIoU / len(categories) 112 | print("mIoU: ", mIoU) 113 | 114 | if __name__ == '__main__': 115 | main() 116 | -------------------------------------------------------------------------------- /evaluation/s3dis_merge.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/python3 2 | 3 | from __future__ import absolute_import 4 | from __future__ import division 5 | from __future__ import print_function 6 | 7 | import os,sys 8 | import plyfile 9 | import numpy as np 10 | import argparse 11 | import h5py 12 | 13 | def main(): 14 | parser = argparse.ArgumentParser() 15 | parser.add_argument('--datafolder', '-d', help='Path to input *_pred.h5', default='../../data/3DIS/prepare_label_rgb/Area_2/') 16 | args = parser.parse_args() 17 | print(args) 18 | 19 | 20 | categories_list = os.listdir(args.datafolder) 21 | 22 | for category in categories_list: 23 | output_path = os.path.join(args.datafolder,category,"pred.npy") 24 | label_length = np.load(os.path.join(args.datafolder,category,"label.npy")).shape[0] 25 | 26 | merged_label_zero = np.zeros((label_length),dtype=int) 27 | merged_confidence_zero = np.zeros((label_length),dtype=float) 28 | merged_label_half = np.zeros((label_length), dtype=int) 29 | merged_confidence_half = np.zeros((label_length), dtype=float) 30 | #merged_label = np.zeros((label_length,2)) 31 | 32 | final_label = np.zeros((label_length), dtype=int) 33 | pred_list = [pred for pred in os.listdir(os.path.join(args.datafolder,category)) if pred.split(".")[-1] == "h5" and "pred" in pred] 34 | for pred_file in pred_list: 35 | print(os.path.join(args.datafolder,category, pred_file)) 36 | data = h5py.File(os.path.join(args.datafolder,category, pred_file)) 37 | labels_seg = data['label_seg'][...].astype(np.int64) 38 | indices = data['indices_split_to_full'][...].astype(np.int64) 39 | confidence = data['confidence'][...].astype(np.float32) 40 | data_num = data['data_num'][...].astype(np.int64) 41 | 42 | if 'zero' in pred_file: 43 | for i in range(labels_seg.shape[0]): 44 | merged_label_zero[indices[i][:data_num[i]]] = labels_seg[i][:data_num[i]] 45 | merged_confidence_zero[indices[i][:data_num[i]]] = confidence[i][:data_num[i]] 46 | else: 47 | for i in range(labels_seg.shape[0]): 48 | merged_label_half[indices[i][:data_num[i]]] = labels_seg[i][:data_num[i]] 49 | merged_confidence_half[indices[i][:data_num[i]]] = confidence[i][:data_num[i]] 50 | 51 | final_label[merged_confidence_zero >= merged_confidence_half] = merged_label_zero[merged_confidence_zero >= merged_confidence_half] 52 | final_label[merged_confidence_zero < merged_confidence_half] = merged_label_half[merged_confidence_zero < merged_confidence_half] 53 | 54 | np.savetxt(output_path,final_label,fmt='%d') 55 | print("saved to ",output_path) 56 | 57 | 58 | if __name__ == '__main__': 59 | main() 60 | -------------------------------------------------------------------------------- /evaluation/semantic3d_merge.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/python3 2 | 3 | from __future__ import absolute_import 4 | from __future__ import division 5 | from __future__ import print_function 6 | 7 | import os,sys 8 | import plyfile 9 | import numpy as np 10 | import argparse 11 | import h5py 12 | 13 | reduced_length_dict = {"MarketplaceFeldkirch":[10538633,"marketsquarefeldkirch4-reduced"], 14 | "StGallenCathedral":[14608690,"stgallencathedral6-reduced"], 15 | "sg27":[28931322,"sg27_10-reduced"], 16 | "sg28":[24620684,"sg28_2-reduced"]} 17 | 18 | full_length_dict = {"stgallencathedral_station1":[31179769,"stgallencathedral1"], 19 | "stgallencathedral_station3":[31643853,"stgallencathedral3"], 20 | "stgallencathedral_station6":[32486227,"stgallencathedral6"], 21 | "marketplacefeldkirch_station1":[26884140,"marketsquarefeldkirch1"], 22 | "marketplacefeldkirch_station4":[23137668,"marketsquarefeldkirch4"], 23 | "marketplacefeldkirch_station7":[23419114,"marketsquarefeldkirch7"], 24 | "birdfountain_station1":[40133912,"birdfountain1"], 25 | "castleblatten_station1":[31806225,"castleblatten1"], 26 | "castleblatten_station5":[49152311,"castleblatten5"], 27 | "sg27_station3":[422445052,"sg27_3"], 28 | "sg27_station6":[226790878,"sg27_6"], 29 | "sg27_station8":[429615314,"sg27_8"], 30 | "sg27_station10":[285579196,"sg27_10"], 31 | "sg28_station2":[170158281,"sg28_2"], 32 | "sg28_station5":[267520082,"sg28_5"]} 33 | 34 | def main(): 35 | parser = argparse.ArgumentParser() 36 | parser.add_argument('--datafolder', '-d', help='Path to input *_pred.h5', default='../../data/semantic3d/raw/test_reduced/') 37 | parser.add_argument('--version', '-v', help='full or reduced', type=str, default='reduced') 38 | args = parser.parse_args() 39 | print(args) 40 | 41 | if args.version == 'full': 42 | length_dict = full_length_dict 43 | else: 44 | length_dict = reduced_length_dict 45 | 46 | categories_list = [category for category in length_dict] 47 | print(categories_list) 48 | 49 | for category in categories_list: 50 | output_path = os.path.join(args.datafolder,"results",length_dict[category][1]+".labels") 51 | if not os.path.exists(os.path.join(args.datafolder,"results")): 52 | os.makedirs(os.path.join(args.datafolder,"results")) 53 | pred_list = [pred for pred in os.listdir(args.datafolder) 54 | if category in pred and pred.split(".")[0].split("_")[-1] == 'pred'] 55 | 56 | label_length = length_dict[category][0] 57 | merged_label = np.zeros((label_length),dtype=int) 58 | merged_confidence = np.zeros((label_length),dtype=float) 59 | 60 | for pred_file in pred_list: 61 | print(os.path.join(args.datafolder, pred_file)) 62 | data = h5py.File(os.path.join(args.datafolder, pred_file)) 63 | labels_seg = data['label_seg'][...].astype(np.int64) 64 | indices = data['indices_split_to_full'][...].astype(np.int64) 65 | confidence = data['confidence'][...].astype(np.float32) 66 | data_num = data['data_num'][...].astype(np.int64) 67 | 68 | for i in range(labels_seg.shape[0]): 69 | temp_label = np.zeros((data_num[i]),dtype=int) 70 | pred_confidence = confidence[i][:data_num[i]] 71 | temp_confidence = merged_confidence[indices[i][:data_num[i]]] 72 | 73 | temp_label[temp_confidence >= pred_confidence] = merged_label[indices[i][:data_num[i]]][temp_confidence >= pred_confidence] 74 | temp_label[pred_confidence > temp_confidence] = labels_seg[i][:data_num[i]][pred_confidence > temp_confidence] 75 | 76 | merged_confidence[indices[i][:data_num[i]][pred_confidence > temp_confidence]] = pred_confidence[pred_confidence > temp_confidence] 77 | merged_label[indices[i][:data_num[i]]] = temp_label 78 | 79 | np.savetxt(output_path,merged_label+1,fmt='%d') 80 | 81 | if __name__ == '__main__': 82 | main() -------------------------------------------------------------------------------- /log/cls/.gitignore: -------------------------------------------------------------------------------- 1 | # Ignore everything in this directory 2 | * 3 | # Except this file 4 | !.gitignore -------------------------------------------------------------------------------- /log/seg/.gitignore: -------------------------------------------------------------------------------- 1 | # Ignore everything in this directory 2 | * 3 | # Except this file 4 | !.gitignore -------------------------------------------------------------------------------- /setting/cls_modelnet40.py: -------------------------------------------------------------------------------- 1 | filelist = '../data/modelnet/train_files.txt' 2 | filelist_val = '../data/modelnet/test_files.txt' 3 | shape_name_file = '../data/modelnet/shape_names.txt' 4 | 5 | num_class = 40 6 | sample_num = 1024 7 | batch_size = 32 8 | batch_size_val = 4 9 | num_epochs = 250 10 | step_val = 1000 11 | momentum = 0.9 12 | 13 | learning_rate_base = 0.001 14 | decay_steps = 200000 15 | decay_rate = 0.7 16 | learning_rate_min = 1e-5 17 | 18 | weight_decay = 1e-8 # 'wd':0.0001 # weight_decay = 1e-6 to avoid overfitting 19 | 20 | ss = 16 # shell size (number of points contained in each shell) 21 | sconv_param_name = ('K', 'D', 'P', 'C') 22 | sconv_params = [dict(zip(sconv_param_name, sconv_param)) for sconv_param in 23 | [ 24 | (ss*4, 4, 512, 128), 25 | (ss*2, 2, 128, 256), 26 | (ss*1, 1, 32, 512)]] 27 | 28 | sdconv_params = None 29 | 30 | x=2 31 | 32 | fc_param_name = ('C', 'dropout_rate') 33 | fc_params = [dict(zip(fc_param_name, fc_param)) for fc_param in 34 | [(128 * x, 0.2), 35 | (64 * x, 0.5)]] 36 | sampling = 'random' # 'fps' or 'random' 37 | optimizer = 'adam' 38 | 39 | data_dim = 3 40 | with_multi=True -------------------------------------------------------------------------------- /setting/seg_s3dis.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/python3 2 | import math 3 | 4 | # modify the following path to train and test on other areas 5 | filelist = '../data/3DIS/prepare_label_rgb/train_files_for_val_on_Area_1.txt' 6 | filelist_val = '../data/3DIS/prepare_label_rgb/val_files_Area_1.txt' 7 | 8 | num_class = 13 9 | 10 | sample_num = 2048 11 | 12 | max_point_num = 8192 13 | 14 | batch_size = 16 15 | 16 | num_epochs = 200 17 | 18 | label_weights = [1.0] * num_class 19 | 20 | learning_rate_base = 0.001 21 | decay_steps = 5000 22 | step_val = 500 23 | decay_rate = 0.8 24 | learning_rate_min = 1e-6 25 | 26 | weight_decay = 1e-8 27 | 28 | jitter = 0.0 29 | jitter_val = 0.0 30 | 31 | sample_num_variance = 1 // 8 32 | sample_num_clip = 1 // 4 33 | 34 | rotation_range = [0, math.pi/72., math.pi/72., 'u'] 35 | rotation_range_val = [0, 0, 0, 'u'] 36 | rotation_order = 'rxyz' 37 | 38 | scaling_range = [0.001, 0.001, 0.001, 'g'] 39 | scaling_range_val = [0, 0, 0, 'u'] 40 | 41 | ss = 8 # shell size (number of points contained in each shell) 42 | sconv_param_name = ('K', 'D', 'P', 'C') 43 | sconv_params = [dict(zip(sconv_param_name, sconv_param)) for sconv_param in 44 | [ 45 | (ss*4, 4, 512, 128), 46 | (ss*2, 2, 128, 256), 47 | (ss*1, 1, 32, 512)]] 48 | 49 | sdconv_param_name = ('K', 'D', 'pts_layer_idx', 'qrs_layer_idx') 50 | sdconv_params = [dict(zip(sdconv_param_name, sdconv_param)) for sdconv_param in 51 | [ 52 | (ss*1, 1, 2, 1), 53 | (ss*2, 2, 1, 0), 54 | (ss*4, 4, 0, -1)]] 55 | 56 | x=1 57 | fc_param_name = ('C', 'dropout_rate') 58 | fc_params = [dict(zip(fc_param_name, fc_param)) for fc_param in 59 | [(128 * x, 0.0), 60 | (64 * x, 0.2)]] 61 | 62 | sampling = 'fps' # 'fps' or 'random' 63 | optimizer = 'adam' 64 | data_dim = 3 -------------------------------------------------------------------------------- /setting/seg_scannet.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/python3 2 | import math 3 | 4 | filelist = '../data/scannet/seg/train_files.txt' 5 | filelist_val = '../data/scannet/seg/test_files.txt' 6 | 7 | num_class = 21 8 | 9 | sample_num = 2048 10 | 11 | max_point_num = 8192 12 | 13 | batch_size = 16 14 | 15 | num_epochs = 200 16 | 17 | label_weights = [0.0] * 1 + [1.0] * (num_class - 1) 18 | 19 | learning_rate_base = 0.005 20 | decay_steps = 5000 21 | step_val = 500 22 | decay_rate = 0.8 23 | learning_rate_min = 1e-6 24 | 25 | weight_decay = 1e-8 26 | 27 | jitter = 0.001 28 | jitter_val = 0.0 29 | 30 | sample_num_variance = 1 // 8 31 | sample_num_clip = 1 // 4 32 | 33 | rotation_range = [0, math.pi/32., 0, 'u'] 34 | rotation_range_val = [0, 0, 0, 'u'] 35 | rotation_order = 'rxyz' 36 | 37 | scaling_range = [0.001, 0.001, 0.001, 'g'] 38 | scaling_range_val = [0, 0, 0, 'u'] 39 | 40 | ss = 8 # shell size (number of points contained in each shell) 41 | sconv_param_name = ('K', 'D', 'P', 'C') 42 | sconv_params = [dict(zip(sconv_param_name, sconv_param)) for sconv_param in 43 | [ 44 | (ss*4, 4, 512, 128), 45 | (ss*2, 2, 128, 256), 46 | (ss*1, 1, 32, 512)]] 47 | 48 | sdconv_param_name = ('K', 'D', 'pts_layer_idx', 'qrs_layer_idx') 49 | sdconv_params = [dict(zip(sdconv_param_name, sdconv_param)) for sdconv_param in 50 | [ 51 | (ss*1, 1, 2, 1), 52 | (ss*2, 2, 1, 0), 53 | (ss*4, 4, 0, -1)]] 54 | 55 | x=1 56 | fc_param_name = ('C', 'dropout_rate') 57 | fc_params = [dict(zip(fc_param_name, fc_param)) for fc_param in 58 | [(128 * x, 0.0), 59 | (64 * x, 0.2)]] 60 | 61 | sampling = 'fps' # 'fps' or 'random' 62 | 63 | optimizer = 'adam' 64 | 65 | data_dim = 3 -------------------------------------------------------------------------------- /setting/seg_semantic3d.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/python3 2 | import math 3 | 4 | filelist = '../data/semantic3d/downsampled/train_data_files.txt' 5 | filelist_val = '../data/semantic3d/downsampled/val_data_files.txt' 6 | filelist_test = '../data/semantic3d/raw/test_reduced_files.txt' 7 | 8 | num_class = 8 9 | 10 | sample_num = 2048 11 | 12 | max_point_num = 8192 13 | 14 | batch_size = 16 15 | 16 | num_epochs = 200 17 | 18 | label_weights = [] 19 | for c in range(num_class): 20 | label_weights.append(1.0) 21 | 22 | learning_rate_base = 0.001 23 | decay_steps = 20000 24 | decay_rate = 0.7 25 | learning_rate_min = 1e-6 26 | 27 | step_val = 500 28 | 29 | weight_decay = 1e-8 30 | 31 | jitter = 0.0 32 | jitter_val = 0.0 33 | 34 | rotation_range = [0, 0, math.pi/72., 'u'] 35 | rotation_range_val = [0, 0, 0, 'u'] 36 | rotation_order = 'rxyz' 37 | 38 | scaling_range = [0.001, 0.001, 0.001, 'g'] 39 | scaling_range_val = [0, 0, 0, 'u'] 40 | 41 | sample_num_variance = 1 // 8 42 | sample_num_clip = 1 // 4 43 | 44 | ss = 8 # shell size (number of points contained in each shell) 45 | sconv_param_name = ('K', 'D', 'P', 'C') 46 | sconv_params = [dict(zip(sconv_param_name, sconv_param)) for sconv_param in 47 | [ 48 | (ss*4, 4, 512, 128), 49 | (ss*2, 2, 128, 256), 50 | (ss*1, 1, 32, 512)]] 51 | 52 | sdconv_param_name = ('K', 'D', 'pts_layer_idx', 'qrs_layer_idx') 53 | sdconv_params = [dict(zip(sdconv_param_name, sdconv_param)) for sdconv_param in 54 | [ 55 | (ss*1, 1, 2, 1), 56 | (ss*2, 2, 1, 0), 57 | (ss*4, 4, 0, -1)]] 58 | 59 | x=1 60 | fc_param_name = ('C', 'dropout_rate') 61 | fc_params = [dict(zip(fc_param_name, fc_param)) for fc_param in 62 | [(128 * x, 0.0), 63 | (64 * x, 0.0)]] 64 | 65 | sampling = 'fps' # 'fps' or 'random' 66 | 67 | optimizer = 'adam' 68 | epsilon = 1e-3 69 | 70 | data_dim = 3 -------------------------------------------------------------------------------- /setting/seg_shapenet.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/python3 2 | import math 3 | 4 | filelist = '../data/shapenet_partseg/train_val_files.txt' 5 | filelist_val = '../data/shapenet_partseg/test_files.txt' 6 | 7 | # 'Path to category list file (.txt)' 8 | category = '../data/shapenet_partseg/categories.txt' 9 | 10 | # 'Path to *.pts directory' 11 | data_folder = '../data/shapenet_partseg/test_data' 12 | 13 | num_class = 50 14 | 15 | sample_num = 2048 16 | 17 | batch_size = 16 18 | 19 | num_epochs = 500 20 | 21 | label_weights = [1.0] * num_class 22 | 23 | learning_rate_base = 0.001 24 | decay_steps = 20000 25 | step_val = 500 26 | decay_rate = 0.7 27 | learning_rate_min = 0.000001 28 | 29 | weight_decay = 0.0 30 | 31 | jitter = 0.001 32 | jitter_val = 0.0 33 | 34 | rotation_range = [0, math.pi/32., 0, 'u'] 35 | rotation_range_val = [0, 0, 0, 'u'] 36 | rotation_order = 'rxyz' 37 | 38 | scaling_range = [0.0, 0.0, 0.0, 'g'] 39 | scaling_range_val = [0, 0, 0, 'u'] 40 | 41 | sample_num_variance = 1 // 8 42 | sample_num_clip = 1 // 4 43 | 44 | ss = 8 # shell size (number of points contained in each shell) 45 | sconv_param_name = ('K', 'D', 'P', 'C') 46 | sconv_params = [dict(zip(sconv_param_name, sconv_param)) for sconv_param in 47 | [ 48 | (ss*4, 4, 512, 128), 49 | (ss*2, 2, 128, 256), 50 | (ss*1, 1, 32, 512)]] 51 | 52 | sdconv_param_name = ('K', 'D', 'pts_layer_idx', 'qrs_layer_idx') 53 | sdconv_params = [dict(zip(sdconv_param_name, sdconv_param)) for sdconv_param in 54 | [ 55 | (ss*1, 1, 2, 1), 56 | (ss*2, 2, 1, 0), 57 | (ss*4, 4, 0, -1)]] 58 | 59 | x=1 60 | fc_param_name = ('C', 'dropout_rate') 61 | fc_params = [dict(zip(fc_param_name, fc_param)) for fc_param in 62 | [(128 * x, 0.2), 63 | (64 * x, 0.2)]] 64 | 65 | sampling = 'fps' # 'fps' or 'random' 66 | 67 | optimizer = 'adam' 68 | epsilon = 1e-3 69 | 70 | data_dim = 3 71 | 72 | keep_remainder = True 73 | 74 | sorting_method = None 75 | with_global = True 76 | with_local = True 77 | -------------------------------------------------------------------------------- /shellconv.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import numpy as np 3 | import math 4 | import sys 5 | import os 6 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 7 | sys.path.append(BASE_DIR) 8 | sys.path.append(os.path.join(BASE_DIR, 'utils')) 9 | sys.path.append(os.path.join(BASE_DIR, 'setting')) 10 | import pointfly as pf 11 | 12 | def dense(input, output, is_training, name, bn_decay=None, with_bn=True, activation=tf.nn.relu): 13 | if with_bn: 14 | input = tf.layers.batch_normalization(input, momentum=0.98, training=is_training, name=name+'bn') 15 | 16 | dense = tf.layers.dense(input, output, activation=activation, name=name) 17 | 18 | return dense 19 | 20 | def conv2d(input, output, name, is_training, kernel_size, bn_decay=None, 21 | reuse=None, with_bn=True, activation=tf.nn.relu): 22 | if with_bn: 23 | input = tf.layers.batch_normalization(input, momentum=0.98, training=is_training, name=name+'bn') 24 | 25 | conv2d = tf.layers.conv2d(input, output, kernel_size=kernel_size, strides=(1, 1), padding='VALID', 26 | activation=activation, reuse=reuse, name=name, use_bias=not with_bn) 27 | return conv2d 28 | 29 | def placeholder_inputs(batch_size, num_point): 30 | pointclouds_pl = tf.placeholder(tf.float32, shape=(batch_size, num_point, 3)) 31 | labels_pl = tf.placeholder(tf.int32, shape=(batch_size)) 32 | return pointclouds_pl, labels_pl 33 | 34 | def shellconv(pts, fts_prev, qrs, is_training, tag, K, D, P, C, with_local, bn_decay=None): 35 | indices = pf.knn_indices_general(qrs, pts, K, True) 36 | 37 | nn_pts = tf.gather_nd(pts, indices, name=tag + 'nn_pts') # (N, P, K, 3) 38 | nn_pts_center = tf.expand_dims(qrs, axis=2, name=tag + 'nn_pts_center') # (N, P, 1, 3) 39 | nn_pts_local = tf.subtract(nn_pts, nn_pts_center, name=tag + 'nn_pts_local') # (N, P, K, 3) 40 | 41 | [N,P,K,dim] = nn_pts_local.shape # (N, P, K, 3) 42 | nn_fts_local = None 43 | C_pts_fts = 64 44 | if with_local: 45 | nn_fts_local = dense(nn_pts_local, C_pts_fts // 2, is_training, tag + 'nn_fts_from_pts_0',bn_decay=bn_decay) 46 | nn_fts_local = dense(nn_fts_local, C_pts_fts, is_training, tag + 'nn_fts_from_pts',bn_decay=bn_decay) 47 | else: 48 | nn_fts_local = nn_pts_local 49 | 50 | if fts_prev is not None: 51 | fts_prev = tf.gather_nd(fts_prev, indices, name=tag + 'fts_prev') # (N, P, K, 3) 52 | pts_X_0 = tf.concat([nn_fts_local,fts_prev], axis=-1) 53 | else: 54 | pts_X_0 = nn_fts_local 55 | 56 | s = int(K.value/D) # no. of divisions 57 | feat_max = tf.layers.max_pooling2d(pts_X_0, [1,s], strides=[1,s], padding='valid', name=tag+'maxpool_0') 58 | 59 | fts_X = conv2d(feat_max, C, name=tag+'conv', is_training=is_training, kernel_size=[1,feat_max.shape[-2].value]) 60 | 61 | fts_X = tf.squeeze(fts_X, axis=-2) 62 | return fts_X 63 | 64 | def get_model(layer_pts, is_training, sconv_params, sdconv_params, fc_params, sampling='random', weight_decay=0.0, bn_decay=None, part_num=8): 65 | if sampling == 'fps': 66 | sys.path.append(os.path.join(BASE_DIR, 'tf_ops/sampling')) 67 | from tf_sampling import farthest_point_sample, gather_point 68 | 69 | layer_fts_list = [None] 70 | layer_pts_list = [layer_pts] 71 | for layer_idx, layer_param in enumerate(sconv_params): 72 | tag = 'sconv_' + str(layer_idx + 1) + '_' 73 | K = layer_param['K'] 74 | D = layer_param['D'] 75 | P = layer_param['P'] 76 | C = layer_param['C'] 77 | if P == -1: 78 | qrs = layer_pts 79 | else: 80 | if sampling == 'fps': 81 | qrs = gather_point(layer_pts, farthest_point_sample(P, layer_pts)) 82 | elif sampling == 'random': 83 | qrs = tf.slice(layer_pts, (0, 0, 0), (-1, P, -1), name=tag + 'qrs') # (N, P, 3) 84 | else: 85 | print('Unknown sampling method!') 86 | exit() 87 | 88 | layer_fts= shellconv(layer_pts_list[-1], layer_fts_list[-1], qrs, is_training, tag, K, D, P, C, True, bn_decay) 89 | 90 | layer_pts = qrs 91 | layer_pts_list.append(qrs) 92 | layer_fts_list.append(layer_fts) 93 | 94 | if sdconv_params is not None: 95 | fts = layer_fts_list[-1] 96 | for layer_idx, layer_param in enumerate(sdconv_params): 97 | tag = 'sdconv_' + str(layer_idx + 1) + '_' 98 | K = layer_param['K'] 99 | D = layer_param['D'] 100 | pts_layer_idx = layer_param['pts_layer_idx'] # 2 1 0 101 | qrs_layer_idx = layer_param['qrs_layer_idx'] # 1 0 -1 102 | 103 | pts = layer_pts_list[pts_layer_idx + 1] 104 | qrs = layer_pts_list[qrs_layer_idx + 1] 105 | fts_qrs = layer_fts_list[qrs_layer_idx + 1] 106 | 107 | C = fts_qrs.get_shape()[-1].value if fts_qrs is not None else C//2 108 | P = qrs.get_shape()[1].value 109 | 110 | layer_fts= shellconv(pts, fts, qrs, is_training, tag, K, D, P, C, True, bn_decay) 111 | if fts_qrs is not None: # this is for last layer 112 | fts_concat = tf.concat([layer_fts, fts_qrs], axis=-1, name=tag + 'fts_concat') 113 | fts = dense(fts_concat, C, is_training, tag + 'mlp', bn_decay=bn_decay) 114 | 115 | for layer_idx, layer_param in enumerate(fc_params): 116 | C = layer_param['C'] 117 | dropout_rate = layer_param['dropout_rate'] 118 | layer_fts = dense(layer_fts, C, is_training, name='fc{:d}'.format(layer_idx), bn_decay=bn_decay) 119 | layer_fts = tf.layers.dropout(layer_fts, rate=dropout_rate, name='fc{:d}_dropout'.format(layer_idx)) 120 | 121 | logits = dense(layer_fts, part_num, is_training, name='logits', activation=None) 122 | return logits 123 | 124 | if __name__ == '__main__': 125 | with tf.Graph().as_default(): 126 | with tf.device('/gpu:0'): 127 | sconv_param_name = ('K', 'D', 'P', 'C', 'links') 128 | sconv_params = [dict(zip(sconv_param_name, sconv_param)) for sconv_param in 129 | [(32, 4, 512, 128, []), 130 | (16, 2, 256, 256, []), 131 | (8, 1, 128, 512, [])]] 132 | 133 | sdconv_param_name = ('K', 'D', 'pts_layer_idx', 'qrs_layer_idx') 134 | sdconv_params = [dict(zip(sdconv_param_name, sdconv_param)) for sdconv_param in 135 | [(8, 1, 2, 1), 136 | (16, 2, 1, 0), 137 | (32, 4, 0, -1)]] 138 | 139 | x = 2 140 | fc_param_name = ('C', 'dropout_rate') 141 | fc_params = [dict(zip(fc_param_name, fc_param)) for fc_param in 142 | [(128 * x, 0), 143 | (64 * x, 0.5)]] 144 | inputs = tf.random_uniform([32, 1024, 3], minval=0, maxval=10,dtype=tf.float32) 145 | 146 | with_local = True 147 | 148 | outputs = get_model(inputs, tf.constant(True), sconv_params, sdconv_params, fc_params) 149 | print(outputs) 150 | -------------------------------------------------------------------------------- /test_cls_modelnet40.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | import argparse 4 | import importlib 5 | import numpy as np 6 | from time import time 7 | import tensorflow as tf 8 | from datetime import datetime 9 | 10 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 11 | sys.path.append(BASE_DIR) 12 | sys.path.append(os.path.join(BASE_DIR, 'utils')) 13 | sys.path.append(os.path.join(BASE_DIR, 'setting')) 14 | from utils import provider 15 | 16 | def main(): 17 | parser = argparse.ArgumentParser() 18 | parser.add_argument('--load_ckpt', '-l', default='log/cls/pretrained/ckpts/', help='Path to a check point file for load') 19 | parser.add_argument('--model', '-m', default='shellconv', help='Model to use') 20 | parser.add_argument('--setting', '-x', default='cls_modelnet40', help='Setting to use') 21 | args = parser.parse_args() 22 | 23 | setting_path = os.path.join(os.path.dirname(__file__), args.model) 24 | sys.path.append(setting_path) 25 | setting = importlib.import_module(args.setting) 26 | 27 | batch_size_val = setting.batch_size_val 28 | sample_num = setting.sample_num 29 | 30 | # Prepare inputs 31 | print('{}-Preparing datasets...'.format(datetime.now())) 32 | [data_val, label_val] = provider.load_cls_files(setting.filelist_val) 33 | data_val = data_val[:,0:sample_num,:] 34 | label_val = np.squeeze(label_val) 35 | num_val = label_val.shape[0] 36 | num_batch_val = int(num_val / batch_size_val) 37 | 38 | # load shape names 39 | shape_names = [line.rstrip() for line in \ 40 | open(setting.shape_name_file)] 41 | 42 | ckpt = tf.train.get_checkpoint_state(args.load_ckpt) 43 | saver = tf.train.import_meta_graph(ckpt.model_checkpoint_path + '.meta') 44 | graph = tf.get_default_graph() 45 | ops = {'data_pl': graph.get_tensor_by_name("data_pl:0"), 46 | 'is_training_pl': graph.get_tensor_by_name("is_training:0"), 47 | 'predictions_op': graph.get_tensor_by_name("predictions:0")} 48 | 49 | config = tf.ConfigProto() 50 | config.gpu_options.allow_growth = True 51 | config.allow_soft_placement = True 52 | config.log_device_placement = False 53 | sess = tf.Session(config=config) 54 | 55 | saver.restore(sess, ckpt.model_checkpoint_path) 56 | 57 | # testing 58 | total_correct = 0 59 | total_seen = 0 60 | total_seen_class = [0 for _ in range(setting.num_class)] 61 | total_correct_class = [0 for _ in range(setting.num_class)] 62 | 63 | eval_start_time = time() # eval start time 64 | for batch_idx in range(num_batch_val): 65 | start_idx = batch_idx * batch_size_val 66 | end_idx = (batch_idx+1) * batch_size_val 67 | 68 | feed_dict = {ops['data_pl']: data_val[start_idx:end_idx, :, :], 69 | ops['is_training_pl']: False,} 70 | 71 | # infer_start_time = time() 72 | predictions_val = sess.run(ops['predictions_op'], feed_dict=feed_dict) 73 | # print('infer time : %f' % (time() - infer_start_time)) 74 | 75 | correct = np.sum(predictions_val == label_val[start_idx:end_idx]) 76 | total_correct += correct 77 | total_seen += batch_size_val 78 | for i in range(start_idx, end_idx): 79 | l = label_val[i] 80 | total_seen_class[l] += 1 81 | total_correct_class[l] += (predictions_val[i-start_idx] == l) 82 | 83 | class_accuracies = np.array(total_correct_class)/np.array(total_seen_class,dtype=np.float) 84 | acc_mean_cls = np.mean(np.array(total_correct_class)/np.array(total_seen_class,dtype=np.float)) 85 | acc_oa = total_correct / float(total_seen) 86 | 87 | 88 | print('eval acc (oa): %f ---- eval acc (mean class): %f ---- time cost: %f' % \ 89 | (acc_oa, acc_mean_cls, time() - eval_start_time)) 90 | 91 | print('per-class accuracies:') 92 | for i, name in enumerate(shape_names): 93 | print('%10s:\t%0.3f' % (name, class_accuracies[i])) 94 | 95 | if __name__ == '__main__': 96 | main() -------------------------------------------------------------------------------- /test_seg_3dis.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/python3 2 | """Testing On Segmentation Task.""" 3 | 4 | from __future__ import absolute_import 5 | from __future__ import division 6 | from __future__ import print_function 7 | 8 | import os 9 | import sys 10 | import math 11 | import h5py 12 | import argparse 13 | import importlib 14 | import numpy as np 15 | import tensorflow as tf 16 | from datetime import datetime 17 | 18 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 19 | sys.path.append(BASE_DIR) 20 | sys.path.append(os.path.join(BASE_DIR, 'utils')) 21 | sys.path.append(os.path.join(BASE_DIR, 'setting')) 22 | from utils import provider 23 | from utils import data_utils 24 | 25 | def main(): 26 | parser = argparse.ArgumentParser() 27 | parser.add_argument('--load_ckpt', '-l', default='log/seg//shellconv_seg_s3dis_area_2/ckpts/epoch-19', help='Path to a check point file for load') 28 | parser.add_argument('--model', '-m', default='shellconv', help='Model to use') 29 | parser.add_argument('--setting', '-x', default='seg_s3dis', help='Setting to use') 30 | parser.add_argument('--repeat_num', '-r', help='Repeat number', type=int, default=1) 31 | parser.add_argument('--save_ply', '-s', help='Save results as ply', default=False) 32 | args = parser.parse_args( 33 | print(args) 34 | 35 | 36 | model = importlib.import_module(args.model) 37 | setting_path = os.path.join(os.path.dirname(__file__), args.model) 38 | sys.path.append(setting_path) 39 | setting = importlib.import_module(args.setting) 40 | 41 | sample_num = setting.sample_num 42 | max_point_num = setting.max_point_num 43 | batch_size = args.repeat_num * math.ceil(max_point_num / sample_num) 44 | 45 | ###################################################################### 46 | # Placeholders 47 | indices = tf.placeholder(tf.int32, shape=(batch_size, None, 2), name="indices") 48 | is_training = tf.placeholder(tf.bool, name='is_training') 49 | pointclouds = tf.placeholder(tf.float32, shape=(batch_size, max_point_num, setting.data_dim), name='points') 50 | ###################################################################### 51 | 52 | ###################################################################### 53 | points_sampled = tf.gather_nd(pointclouds, indices=indices, name='pts_sampled') 54 | 55 | 56 | logits_op = model.get_model(points_sampled, is_training, setting.sconv_params, setting.sdconv_params, setting.fc_params, 57 | sampling=setting.sampling, 58 | weight_decay=setting.weight_decay, 59 | bn_decay = None, 60 | part_num=setting.num_class) 61 | probs_op = tf.nn.softmax(logits_op, name='probs_op') 62 | 63 | # for restore model 64 | saver = tf.train.Saver() 65 | 66 | parameter_num = np.sum([np.prod(v.shape.as_list()) for v in tf.trainable_variables()]) 67 | print('{}-Parameter number: {:d}.'.format(datetime.now(), parameter_num)) 68 | 69 | with tf.Session() as sess: 70 | # Load the model 71 | saver.restore(sess, args.load_ckpt) 72 | print('{}-Checkpoint loaded from {}!'.format(datetime.now(), args.load_ckpt)) 73 | 74 | indices_batch_indices = np.tile(np.reshape(np.arange(batch_size), (batch_size, 1, 1)), (1, sample_num, 1)) 75 | 76 | folder = os.path.dirname(setting.filelist_val) 77 | filenames = [os.path.join(folder, line.strip()) for line in open(setting.filelist_val)] 78 | for filename in filenames: 79 | print('{}-Reading {}...'.format(datetime.now(), filename)) 80 | data_h5 = h5py.File(filename) 81 | data = data_h5['data'][...].astype(np.float32) 82 | data_num = data_h5['data_num'][...].astype(np.int32) 83 | batch_num = data.shape[0] 84 | if data.shape[-1] > 3: 85 | data = data[:,:,:3] 86 | 87 | labels_pred = np.full((batch_num, max_point_num), -1, dtype=np.int32) 88 | confidences_pred = np.zeros((batch_num, max_point_num), dtype=np.float32) 89 | 90 | print('{}-{:d} testing batches.'.format(datetime.now(), batch_num)) 91 | for batch_idx in range(batch_num): 92 | if batch_idx % 10 == 0: 93 | print('{}-Processing {} of {} batches.'.format(datetime.now(), batch_idx, batch_num)) 94 | points_batch = data[[batch_idx] * batch_size, ...] 95 | point_num = data_num[batch_idx] 96 | 97 | tile_num = math.ceil((sample_num * batch_size) / point_num) 98 | indices_shuffle = np.tile(np.arange(point_num), tile_num)[0:sample_num * batch_size] 99 | np.random.shuffle(indices_shuffle) 100 | indices_batch_shuffle = np.reshape(indices_shuffle, (batch_size, sample_num, 1)) 101 | indices_batch = np.concatenate((indices_batch_indices, indices_batch_shuffle), axis=2) 102 | 103 | seg_probs = sess.run([probs_op], 104 | feed_dict={ 105 | pointclouds: points_batch, 106 | indices: indices_batch, 107 | is_training: False, 108 | }) 109 | probs_2d = np.reshape(seg_probs, (sample_num * batch_size, -1)) 110 | 111 | predictions = [(-1, 0.0)] * point_num 112 | for idx in range(sample_num * batch_size): 113 | point_idx = indices_shuffle[idx] 114 | probs = probs_2d[idx, :] 115 | confidence = np.amax(probs) 116 | label = np.argmax(probs) 117 | if confidence > predictions[point_idx][1]: 118 | predictions[point_idx] = [label, confidence] 119 | labels_pred[batch_idx, 0:point_num] = np.array([label for label, _ in predictions]) 120 | confidences_pred[batch_idx, 0:point_num] = np.array([confidence for _, confidence in predictions]) 121 | 122 | filename_pred = filename[:-3] + '_pred.h5' 123 | print('{}-Saving {}...'.format(datetime.now(), filename_pred)) 124 | file = h5py.File(filename_pred, 'w') 125 | file.create_dataset('data_num', data=data_num) 126 | file.create_dataset('label_seg', data=labels_pred) 127 | file.create_dataset('confidence', data=confidences_pred) 128 | has_indices = 'indices_split_to_full' in data_h5 129 | if has_indices: 130 | file.create_dataset('indices_split_to_full', data=data_h5['indices_split_to_full'][...]) 131 | file.close() 132 | 133 | if args.save_ply: 134 | print('{}-Saving ply of {}...'.format(datetime.now(), filename_pred)) 135 | filepath_label_ply = os.path.join(filename_pred[:-3] + 'ply_label') 136 | data_utils.save_ply_property_batch(data[:, :, 0:3], labels_pred[...], 137 | filepath_label_ply, data_num[...], setting.num_class) 138 | ###################################################################### 139 | print('{}-Done!'.format(datetime.now())) 140 | 141 | 142 | if __name__ == '__main__': 143 | main() -------------------------------------------------------------------------------- /test_seg_scannet.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/python3 2 | """Testing On Segmentation Task.""" 3 | 4 | from __future__ import absolute_import 5 | from __future__ import division 6 | from __future__ import print_function 7 | 8 | import os 9 | import sys 10 | import math 11 | import h5py 12 | import argparse 13 | import importlib 14 | import numpy as np 15 | import tensorflow as tf 16 | from datetime import datetime 17 | 18 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 19 | sys.path.append(BASE_DIR) 20 | sys.path.append(os.path.join(BASE_DIR, 'utils')) 21 | sys.path.append(os.path.join(BASE_DIR, 'setting')) 22 | from utils import provider 23 | from utils import data_utils 24 | 25 | 26 | def main(): 27 | parser = argparse.ArgumentParser() 28 | parser.add_argument('--load_ckpt', '-l', default='log/seg/shellconv_seg_scannet_xxxx/ckpts/epoch-31', help='Path to a check point file for load') 29 | parser.add_argument('--model', '-m', default='shellconv', help='Model to use') 30 | parser.add_argument('--setting', '-x', default='seg_scannet', help='Setting to use') 31 | parser.add_argument('--repeat_num', '-r', help='Repeat number', type=int, default=1) # 1 is good enough 32 | parser.add_argument('--save_ply', '-s', help='Save results as ply', default=False) 33 | args = parser.parse_args() 34 | print(args) 35 | 36 | model = importlib.import_module(args.model) 37 | setting_path = os.path.join(os.path.dirname(__file__), args.model) 38 | sys.path.append(setting_path) 39 | setting = importlib.import_module(args.setting) 40 | 41 | sample_num = setting.sample_num 42 | max_point_num = setting.max_point_num 43 | batch_size = args.repeat_num * math.ceil(max_point_num / sample_num) 44 | 45 | ###################################################################### 46 | # Placeholders 47 | indices = tf.placeholder(tf.int32, shape=(batch_size, None, 2), name="indices") 48 | is_training = tf.placeholder(tf.bool, name='is_training') 49 | pointclouds = tf.placeholder(tf.float32, shape=(batch_size, max_point_num, setting.data_dim), name='points') 50 | ###################################################################### 51 | 52 | ###################################################################### 53 | points_sampled = tf.gather_nd(pointclouds, indices=indices, name='pts_sampled') 54 | 55 | 56 | logits_op = model.get_model(points_sampled, is_training, setting.sconv_params, setting.sdconv_params, setting.fc_params, 57 | sampling=setting.sampling, 58 | weight_decay=setting.weight_decay, 59 | bn_decay = None, 60 | part_num=setting.num_class) 61 | probs_op = tf.nn.softmax(logits_op, name='probs_op') 62 | 63 | # for restore model 64 | saver = tf.train.Saver() 65 | 66 | parameter_num = np.sum([np.prod(v.shape.as_list()) for v in tf.trainable_variables()]) 67 | print('{}-Parameter number: {:d}.'.format(datetime.now(), parameter_num)) 68 | 69 | with tf.Session() as sess: 70 | # Load the model 71 | saver.restore(sess, args.load_ckpt) 72 | print('{}-Checkpoint loaded from {}!'.format(datetime.now(), args.load_ckpt)) 73 | 74 | indices_batch_indices = np.tile(np.reshape(np.arange(batch_size), (batch_size, 1, 1)), (1, sample_num, 1)) 75 | 76 | folder = os.path.dirname(setting.filelist_val) 77 | filenames = [os.path.join(folder, line.strip()) for line in open(setting.filelist_val)] 78 | for filename in filenames: 79 | print('{}-Reading {}...'.format(datetime.now(), filename)) 80 | data_h5 = h5py.File(filename) 81 | data = data_h5['data'][...].astype(np.float32) 82 | data_num = data_h5['data_num'][...].astype(np.int32) 83 | batch_num = data.shape[0] 84 | if data.shape[-1] > 3: 85 | data = data[:,:,:3] 86 | 87 | labels_pred = np.full((batch_num, max_point_num), -1, dtype=np.int32) 88 | confidences_pred = np.zeros((batch_num, max_point_num), dtype=np.float32) 89 | 90 | print('{}-{:d} testing batches.'.format(datetime.now(), batch_num)) 91 | for batch_idx in range(batch_num): 92 | if batch_idx % 10 == 0: 93 | print('{}-Processing {} of {} batches.'.format(datetime.now(), batch_idx, batch_num)) 94 | points_batch = data[[batch_idx] * batch_size, ...] 95 | point_num = data_num[batch_idx] 96 | 97 | tile_num = math.ceil((sample_num * batch_size) / point_num) 98 | indices_shuffle = np.tile(np.arange(point_num), tile_num)[0:sample_num * batch_size] 99 | np.random.shuffle(indices_shuffle) 100 | indices_batch_shuffle = np.reshape(indices_shuffle, (batch_size, sample_num, 1)) 101 | indices_batch = np.concatenate((indices_batch_indices, indices_batch_shuffle), axis=2) 102 | 103 | seg_probs = sess.run([probs_op], 104 | feed_dict={ 105 | pointclouds: points_batch, 106 | indices: indices_batch, 107 | is_training: False, 108 | }) 109 | probs_2d = np.reshape(seg_probs, (sample_num * batch_size, -1)) 110 | 111 | predictions = [(-1, 0.0)] * point_num 112 | for idx in range(sample_num * batch_size): 113 | point_idx = indices_shuffle[idx] 114 | probs = probs_2d[idx, :] 115 | confidence = np.amax(probs) 116 | label = np.argmax(probs) 117 | if confidence > predictions[point_idx][1]: 118 | predictions[point_idx] = [label, confidence] 119 | labels_pred[batch_idx, 0:point_num] = np.array([label for label, _ in predictions]) 120 | confidences_pred[batch_idx, 0:point_num] = np.array([confidence for _, confidence in predictions]) 121 | 122 | filename_pred = filename[:-3] + '_pred.h5' 123 | print('{}-Saving {}...'.format(datetime.now(), filename_pred)) 124 | file = h5py.File(filename_pred, 'w') 125 | file.create_dataset('data_num', data=data_num) 126 | file.create_dataset('label_seg', data=labels_pred) 127 | file.create_dataset('confidence', data=confidences_pred) 128 | has_indices = 'indices_split_to_full' in data_h5 129 | if has_indices: 130 | file.create_dataset('indices_split_to_full', data=data_h5['indices_split_to_full'][...]) 131 | file.close() 132 | 133 | if args.save_ply: 134 | print('{}-Saving ply of {}...'.format(datetime.now(), filename_pred)) 135 | filepath_label_ply = os.path.join(filename_pred[:-3] + 'ply_label') 136 | data_utils.save_ply_property_batch(data[:, :, 0:3], labels_pred[...], 137 | filepath_label_ply, data_num[...], setting.num_class) 138 | ###################################################################### 139 | print('{}-Done!'.format(datetime.now())) 140 | 141 | 142 | if __name__ == '__main__': 143 | main() -------------------------------------------------------------------------------- /test_seg_semantic3d.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/python3 2 | """Testing On Segmentation Task.""" 3 | 4 | from __future__ import absolute_import 5 | from __future__ import division 6 | from __future__ import print_function 7 | 8 | import os 9 | import sys 10 | import math 11 | import h5py 12 | import argparse 13 | import importlib 14 | import numpy as np 15 | import tensorflow as tf 16 | from datetime import datetime 17 | 18 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 19 | sys.path.append(BASE_DIR) 20 | sys.path.append(os.path.join(BASE_DIR, 'utils')) 21 | sys.path.append(os.path.join(BASE_DIR, 'setting')) 22 | from utils import provider 23 | from utils import data_utils 24 | 25 | def main(): 26 | parser = argparse.ArgumentParser() 27 | parser.add_argument('--load_ckpt', '-l', default='log/seg/shellconv_seg_semantic3d_2019-08-04-15-52-37/ckpts/epoch-79', help='Path to a check point file for load') 28 | parser.add_argument('--model', '-m', default='shellconv', help='Model to use') 29 | parser.add_argument('--setting', '-x', default='seg_semantic3d', help='Setting to use') 30 | parser.add_argument('--repeat_num', '-r', help='Repeat number', type=int, default=1) 31 | parser.add_argument('--save_ply', '-s', help='Save results as ply', default=True) 32 | args = parser.parse_args() 33 | print(args) 34 | 35 | 36 | model = importlib.import_module(args.model) 37 | setting_path = os.path.join(os.path.dirname(__file__), args.model) 38 | sys.path.append(setting_path) 39 | setting = importlib.import_module(args.setting) 40 | 41 | sample_num = setting.sample_num 42 | max_point_num = setting.max_point_num 43 | batch_size = args.repeat_num * math.ceil(max_point_num / sample_num) 44 | 45 | ###################################################################### 46 | # Placeholders 47 | indices = tf.placeholder(tf.int32, shape=(batch_size, None, 2), name="indices") 48 | is_training = tf.placeholder(tf.bool, name='is_training') 49 | pointclouds = tf.placeholder(tf.float32, shape=(batch_size, max_point_num, setting.data_dim), name='points') 50 | ###################################################################### 51 | 52 | ###################################################################### 53 | points_sampled = tf.gather_nd(pointclouds, indices=indices, name='pts_sampled') 54 | 55 | 56 | logits_op = model.get_model(points_sampled, is_training, setting.sconv_params, setting.sdconv_params, setting.fc_params, 57 | sampling=setting.sampling, 58 | weight_decay=setting.weight_decay, 59 | bn_decay = None, 60 | part_num=setting.num_class) 61 | probs_op = tf.nn.softmax(logits_op, name='probs_op') 62 | 63 | # for restore model 64 | saver = tf.train.Saver() 65 | 66 | parameter_num = np.sum([np.prod(v.shape.as_list()) for v in tf.trainable_variables()]) 67 | print('{}-Parameter number: {:d}.'.format(datetime.now(), parameter_num)) 68 | 69 | with tf.Session() as sess: 70 | # Load the model 71 | saver.restore(sess, args.load_ckpt) 72 | print('{}-Checkpoint loaded from {}!'.format(datetime.now(), args.load_ckpt)) 73 | 74 | indices_batch_indices = np.tile(np.reshape(np.arange(batch_size), (batch_size, 1, 1)), (1, sample_num, 1)) 75 | 76 | folder = os.path.dirname(setting.filelist_test) 77 | filenames = [os.path.join(folder, line.strip()) for line in open(setting.filelist_test)] 78 | for filename in filenames: 79 | filename_pred = filename[:-3] + '_pred.h5' 80 | if os.path.isfile(filename_pred): 81 | continue 82 | print('{}-Reading {}...'.format(datetime.now(), filename)) 83 | data_h5 = h5py.File(filename) 84 | data = data_h5['data'][...].astype(np.float32) 85 | data_num = data_h5['data_num'][...].astype(np.int32) 86 | batch_num = data.shape[0] 87 | if data.shape[-1] > 3: 88 | data = data[:,:,:3] 89 | 90 | labels_pred = np.full((batch_num, max_point_num), -1, dtype=np.int32) 91 | confidences_pred = np.zeros((batch_num, max_point_num), dtype=np.float32) 92 | 93 | print('{}-{:d} testing batches.'.format(datetime.now(), batch_num)) 94 | for batch_idx in range(batch_num): 95 | if batch_idx % 10 == 0: 96 | print('{}-Processing {} of {} batches.'.format(datetime.now(), batch_idx, batch_num)) 97 | points_batch = data[[batch_idx] * batch_size, ...] 98 | point_num = data_num[batch_idx] 99 | 100 | tile_num = math.ceil((sample_num * batch_size) / point_num) 101 | indices_shuffle = np.tile(np.arange(point_num), tile_num)[0:sample_num * batch_size] 102 | np.random.shuffle(indices_shuffle) 103 | indices_batch_shuffle = np.reshape(indices_shuffle, (batch_size, sample_num, 1)) 104 | indices_batch = np.concatenate((indices_batch_indices, indices_batch_shuffle), axis=2) 105 | 106 | seg_probs = sess.run([probs_op], 107 | feed_dict={ 108 | pointclouds: points_batch, 109 | indices: indices_batch, 110 | is_training: False, 111 | }) 112 | probs_2d = np.reshape(seg_probs, (sample_num * batch_size, -1)) 113 | 114 | predictions = [(-1, 0.0)] * point_num 115 | for idx in range(sample_num * batch_size): 116 | point_idx = indices_shuffle[idx] 117 | probs = probs_2d[idx, :] 118 | confidence = np.amax(probs) 119 | label = np.argmax(probs) 120 | if confidence > predictions[point_idx][1]: 121 | predictions[point_idx] = [label, confidence] 122 | labels_pred[batch_idx, 0:point_num] = np.array([label for label, _ in predictions]) 123 | confidences_pred[batch_idx, 0:point_num] = np.array([confidence for _, confidence in predictions]) 124 | 125 | filename_pred = filename[:-3] + '_pred.h5' 126 | print('{}-Saving {}...'.format(datetime.now(), filename_pred)) 127 | file = h5py.File(filename_pred, 'w') 128 | file.create_dataset('data_num', data=data_num) 129 | file.create_dataset('label_seg', data=labels_pred) 130 | file.create_dataset('confidence', data=confidences_pred) 131 | has_indices = 'indices_split_to_full' in data_h5 132 | if has_indices: 133 | file.create_dataset('indices_split_to_full', data=data_h5['indices_split_to_full'][...]) 134 | file.close() 135 | 136 | # if args.save_ply: 137 | # print('{}-Saving ply of {}...'.format(datetime.now(), filename_pred)) 138 | # filepath_label_ply = os.path.join(filename_pred[:-3] + 'ply_label') 139 | # data_utils.save_ply_property_batch(data[:, :, 0:3], labels_pred[...], 140 | # filepath_label_ply, data_num[...], setting.num_class) 141 | ###################################################################### 142 | print('{}-Done!'.format(datetime.now())) 143 | 144 | 145 | if __name__ == '__main__': 146 | main() -------------------------------------------------------------------------------- /test_seg_shapenet.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/python3 2 | """Testing On ShapeNet Parts Segmentation Task.""" 3 | 4 | from __future__ import absolute_import 5 | from __future__ import division 6 | from __future__ import print_function 7 | 8 | import os 9 | import sys 10 | import math 11 | import argparse 12 | import importlib 13 | import numpy as np 14 | import tensorflow as tf 15 | from datetime import datetime 16 | 17 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 18 | sys.path.append(BASE_DIR) 19 | sys.path.append(os.path.join(BASE_DIR, 'utils')) 20 | sys.path.append(os.path.join(BASE_DIR, 'setting')) 21 | from utils import provider 22 | from utils import data_utils 23 | 24 | 25 | def main(): 26 | parser = argparse.ArgumentParser() 27 | parser.add_argument('--load_ckpt', '-l', default='log/seg/shellconv_seg_shapenet_2019-08-06-14-42-34/ckpts/epoch-326', help='Path to a check point file for load') 28 | parser.add_argument('--model', '-m', default='shellconv', help='Model to use') 29 | parser.add_argument('--setting', '-x', default='seg_shapenet', help='Setting to use') 30 | parser.add_argument('--repeat_num', '-r', help='Repeat number', type=int, default=1) 31 | parser.add_argument('--save_ply', '-s', help='Save results as ply', default=False) 32 | args = parser.parse_args() 33 | print(args) 34 | 35 | model = importlib.import_module(args.model) 36 | setting_path = os.path.join(os.path.dirname(__file__), args.model) 37 | sys.path.append(setting_path) 38 | setting = importlib.import_module(args.setting) 39 | 40 | sample_num = setting.sample_num 41 | 42 | output_folder = setting.data_folder + '_pred_shellnet_' + str(args.repeat_num) 43 | category_list = [(category, int(label_num)) for (category, label_num) in 44 | [line.split() for line in open(setting.category, 'r')]] 45 | offset = 0 46 | category_range = dict() 47 | for category, category_label_seg_max in category_list: 48 | category_range[category] = (offset, offset + category_label_seg_max) 49 | offset = offset + category_label_seg_max 50 | folder = os.path.join(output_folder, category) 51 | if not os.path.exists(folder): 52 | os.makedirs(folder) 53 | 54 | input_filelist = [] 55 | output_filelist = [] 56 | output_ply_filelist = [] 57 | for category in sorted(os.listdir(setting.data_folder)): 58 | data_category_folder = os.path.join(setting.data_folder, category) 59 | for filename in sorted(os.listdir(data_category_folder)): 60 | input_filelist.append(os.path.join(setting.data_folder, category, filename)) 61 | output_filelist.append(os.path.join(output_folder, category, filename[0:-3] + 'seg')) 62 | output_ply_filelist.append(os.path.join(output_folder + '_ply', category, filename[0:-3] + 'ply')) 63 | 64 | # Prepare inputs 65 | print('{}-Preparing datasets...'.format(datetime.now())) 66 | data, label, data_num, _, _ = data_utils.load_seg(setting.filelist_val) 67 | 68 | batch_num = data.shape[0] 69 | max_point_num = data.shape[1] 70 | batch_size = args.repeat_num * math.ceil(data.shape[1] / sample_num) 71 | 72 | print('{}-{:d} testing batches.'.format(datetime.now(), batch_num)) 73 | 74 | ###################################################################### 75 | # Placeholders 76 | indices = tf.placeholder(tf.int32, shape=(batch_size, None, 2), name="indices") 77 | is_training = tf.placeholder(tf.bool, name='is_training') 78 | pts_fts = tf.placeholder(tf.float32, shape=(None, max_point_num, setting.data_dim), name='pts_fts') 79 | ###################################################################### 80 | 81 | ###################################################################### 82 | pts_fts_sampled = tf.gather_nd(pts_fts, indices=indices, name='pts_fts_sampled') 83 | if setting.data_dim > 3: 84 | points_sampled, _ = tf.split(pts_fts_sampled, 85 | [3, setting.data_dim - 3], 86 | axis=-1, 87 | name='split_points_features') 88 | else: 89 | points_sampled = pts_fts_sampled 90 | 91 | logits_op = model.get_model(points_sampled, is_training, setting.sconv_params, setting.sdconv_params, setting.fc_params, 92 | sampling=setting.sampling, 93 | weight_decay=setting.weight_decay, 94 | bn_decay = None, 95 | part_num=setting.num_class) 96 | 97 | probs_op = tf.nn.softmax(logits_op, name='probs') 98 | 99 | saver = tf.train.Saver() 100 | 101 | parameter_num = np.sum([np.prod(v.shape.as_list()) for v in tf.trainable_variables()]) 102 | print('{}-Parameter number: {:d}.'.format(datetime.now(), parameter_num)) 103 | 104 | with tf.Session() as sess: 105 | # Load the model 106 | saver.restore(sess, args.load_ckpt) 107 | print('{}-Checkpoint loaded from {}!'.format(datetime.now(), args.load_ckpt)) 108 | 109 | indices_batch_indices = np.tile(np.reshape(np.arange(batch_size), (batch_size, 1, 1)), (1, sample_num, 1)) 110 | for batch_idx in range(batch_num): 111 | points_batch = data[[batch_idx] * batch_size, ...] 112 | object_label = label[batch_idx] 113 | point_num = data_num[batch_idx] 114 | category = category_list[object_label][0] 115 | label_start, label_end = category_range[category] 116 | 117 | tile_num = math.ceil((sample_num * batch_size) / point_num) 118 | indices_shuffle = np.tile(np.arange(point_num), tile_num)[0:sample_num * batch_size] 119 | np.random.shuffle(indices_shuffle) 120 | indices_batch_shuffle = np.reshape(indices_shuffle, (batch_size, sample_num, 1)) 121 | indices_batch = np.concatenate((indices_batch_indices, indices_batch_shuffle), axis=2) 122 | 123 | probs = sess.run([probs_op], 124 | feed_dict={ 125 | pts_fts: points_batch, 126 | indices: indices_batch, 127 | is_training: False, 128 | }) 129 | probs_2d = np.reshape(probs, (sample_num * batch_size, -1)) 130 | predictions = [(-1, 0.0)] * point_num 131 | for idx in range(sample_num * batch_size): 132 | point_idx = indices_shuffle[idx] 133 | probs = probs_2d[idx, label_start:label_end] 134 | confidence = np.amax(probs) 135 | seg_idx = np.argmax(probs) 136 | if confidence > predictions[point_idx][1]: 137 | predictions[point_idx] = (seg_idx, confidence) 138 | 139 | labels = [] 140 | with open(output_filelist[batch_idx], 'w') as file_seg: 141 | for seg_idx, _ in predictions: 142 | file_seg.write('%d\n' % (seg_idx)) 143 | labels.append(seg_idx) 144 | 145 | # read the coordinates from the txt file for verification 146 | coordinates = [[float(value) for value in xyz.split(' ')] 147 | for xyz in open(input_filelist[batch_idx], 'r') if len(xyz.split(' ')) == 3] 148 | assert (point_num == len(coordinates)) 149 | if args.save_ply: 150 | data_utils.save_ply_property(np.array(coordinates), np.array(labels), 6, output_ply_filelist[batch_idx]) 151 | 152 | print('{}-[Testing]-Iter: {:06d} saved to {}'.format(datetime.now(), batch_idx, output_filelist[batch_idx])) 153 | sys.stdout.flush() 154 | ###################################################################### 155 | print('{}-Done!'.format(datetime.now())) 156 | 157 | 158 | if __name__ == '__main__': 159 | main() -------------------------------------------------------------------------------- /tf_ops/sampling/LICENSE: -------------------------------------------------------------------------------- 1 | PointNet++: Deep Hierarchical Feature Learning on Point Sets in a Metric Space. 2 | 3 | Copyright (c) 2017, Geometric Computation Group of Stanford University 4 | 5 | The MIT License (MIT) 6 | 7 | Copyright (c) 2017 Charles R. Qi 8 | 9 | Permission is hereby granted, free of charge, to any person obtaining a copy 10 | of this software and associated documentation files (the "Software"), to deal 11 | in the Software without restriction, including without limitation the rights 12 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 13 | copies of the Software, and to permit persons to whom the Software is 14 | furnished to do so, subject to the following conditions: 15 | 16 | The above copyright notice and this permission notice shall be included in all 17 | copies or substantial portions of the Software. 18 | 19 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 20 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 21 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 22 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 23 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 24 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 25 | SOFTWARE. 26 | -------------------------------------------------------------------------------- /tf_ops/sampling/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hkust-vgd/shellnet/fe5cf5f3806e618c23e161b0fb8a57387268df78/tf_ops/sampling/__init__.py -------------------------------------------------------------------------------- /tf_ops/sampling/tf_sampling.cpp: -------------------------------------------------------------------------------- 1 | /* Furthest point sampling 2 | * Original author: Haoqiang Fan 3 | * Modified by Charles R. Qi 4 | * All Rights Reserved. 2017. 5 | */ 6 | #include "tensorflow/core/framework/op.h" 7 | #include "tensorflow/core/framework/op_kernel.h" 8 | #include "tensorflow/core/framework/shape_inference.h" 9 | #include "tensorflow/core/framework/common_shape_fns.h" 10 | #include 11 | 12 | using namespace tensorflow; 13 | 14 | REGISTER_OP("ProbSample") 15 | .Input("inp: float32") 16 | .Input("inpr: float32") 17 | .Output("out: int32") 18 | .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) { 19 | ::tensorflow::shape_inference::ShapeHandle dims1; // batch_size * ncategory 20 | c->WithRank(c->input(0), 2, &dims1); 21 | ::tensorflow::shape_inference::ShapeHandle dims2; // batch_size * npoints 22 | c->WithRank(c->input(1), 2, &dims2); 23 | // batch_size * npoints 24 | ::tensorflow::shape_inference::ShapeHandle output = c->MakeShape({c->Dim(dims2, 0), c->Dim(dims2, 1)}); 25 | c->set_output(0, output); 26 | return Status::OK(); 27 | }); 28 | REGISTER_OP("FarthestPointSample") 29 | .Attr("npoint: int") 30 | .Input("inp: float32") 31 | .Output("out: int32") 32 | .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) { 33 | ::tensorflow::shape_inference::ShapeHandle dims1; // batch_size * npoint * 3 34 | c->WithRank(c->input(0), 3, &dims1); 35 | int npoint; 36 | TF_RETURN_IF_ERROR(c->GetAttr("npoint", &npoint)); 37 | ::tensorflow::shape_inference::ShapeHandle output = c->MakeShape({c->Dim(dims1, 0), npoint}); 38 | c->set_output(0, output); 39 | return Status::OK(); 40 | }); 41 | REGISTER_OP("GatherPoint") 42 | .Input("inp: float32") 43 | .Input("idx: int32") 44 | .Output("out: float32") 45 | .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) { 46 | ::tensorflow::shape_inference::ShapeHandle dims1; // batch_size * ndataset * 3 47 | c->WithRank(c->input(0), 3, &dims1); 48 | ::tensorflow::shape_inference::ShapeHandle dims2; // batch_size * npoints 49 | c->WithRank(c->input(1), 2, &dims2); 50 | // batch_size * npoints * 3 51 | ::tensorflow::shape_inference::ShapeHandle output = c->MakeShape({c->Dim(dims1, 0), c->Dim(dims2, 1), c->Dim(dims1, 2)}); 52 | c->set_output(0, output); 53 | return Status::OK(); 54 | }); 55 | REGISTER_OP("GatherPointGrad") 56 | .Input("inp: float32") 57 | .Input("idx: int32") 58 | .Input("out_g: float32") 59 | .Output("inp_g: float32") 60 | .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) { 61 | c->set_output(0, c->input(0)); 62 | return Status::OK(); 63 | }); 64 | 65 | void probsampleLauncher(int b,int n,int m,const float * inp_p,const float * inp_r,float * temp,int * out); 66 | class ProbSampleGpuOp: public OpKernel{ 67 | public: 68 | explicit ProbSampleGpuOp(OpKernelConstruction* context):OpKernel(context){} 69 | void Compute(OpKernelContext * context)override{ 70 | const Tensor& inp_tensor=context->input(0); 71 | const Tensor& inpr_tensor=context->input(1); 72 | auto inp_flat=inp_tensor.flat(); 73 | auto inpr_flat=inpr_tensor.flat(); 74 | const float * inp=&(inp_flat(0)); 75 | const float * inpr=&(inpr_flat(0)); 76 | OP_REQUIRES(context,inp_tensor.dims()==2,errors::InvalidArgument("ProbSample expects (batch_size,num_choices) inp shape")); 77 | int b=inp_tensor.shape().dim_size(0); 78 | int n=inp_tensor.shape().dim_size(1); 79 | OP_REQUIRES(context,inpr_tensor.dims()==2 && inpr_tensor.shape().dim_size(0)==b,errors::InvalidArgument("ProbSample expects (batch_size,num_points) inpr shape")); 80 | int m=inpr_tensor.shape().dim_size(1); 81 | Tensor * out_tensor=NULL; 82 | OP_REQUIRES_OK(context,context->allocate_output(0,TensorShape{b,m},&out_tensor)); 83 | auto out_flat=out_tensor->flat(); 84 | int * out=&(out_flat(0)); 85 | Tensor temp_tensor; 86 | OP_REQUIRES_OK(context,context->allocate_temp(DataTypeToEnum::value,TensorShape{b,n},&temp_tensor)); 87 | auto temp_flat=temp_tensor.flat(); 88 | float * temp=&(temp_flat(0)); 89 | probsampleLauncher(b,n,m,inp,inpr,temp,out); 90 | } 91 | }; 92 | REGISTER_KERNEL_BUILDER(Name("ProbSample").Device(DEVICE_GPU), ProbSampleGpuOp); 93 | 94 | void farthestpointsamplingLauncher(int b,int n,int m,const float * inp,float * temp,int * out); 95 | class FarthestPointSampleGpuOp: public OpKernel{ 96 | public: 97 | explicit FarthestPointSampleGpuOp(OpKernelConstruction* context):OpKernel(context) { 98 | OP_REQUIRES_OK(context, context->GetAttr("npoint", &npoint_)); 99 | OP_REQUIRES(context, npoint_ > 0, errors::InvalidArgument("FarthestPointSample expects positive npoint")); 100 | } 101 | void Compute(OpKernelContext * context)override{ 102 | int m = npoint_; 103 | 104 | const Tensor& inp_tensor=context->input(0); 105 | OP_REQUIRES(context,inp_tensor.dims()==3 && inp_tensor.shape().dim_size(2)==3,errors::InvalidArgument("FarthestPointSample expects (batch_size,num_points,3) inp shape")); 106 | int b=inp_tensor.shape().dim_size(0); 107 | int n=inp_tensor.shape().dim_size(1); 108 | auto inp_flat=inp_tensor.flat(); 109 | const float * inp=&(inp_flat(0)); 110 | Tensor * out_tensor; 111 | OP_REQUIRES_OK(context,context->allocate_output(0,TensorShape{b,m},&out_tensor)); 112 | auto out_flat=out_tensor->flat(); 113 | int * out=&(out_flat(0)); 114 | Tensor temp_tensor; 115 | OP_REQUIRES_OK(context,context->allocate_temp(DataTypeToEnum::value,TensorShape{32,n},&temp_tensor)); 116 | auto temp_flat=temp_tensor.flat(); 117 | float * temp=&(temp_flat(0)); 118 | farthestpointsamplingLauncher(b,n,m,inp,temp,out); 119 | } 120 | private: 121 | int npoint_; 122 | }; 123 | REGISTER_KERNEL_BUILDER(Name("FarthestPointSample").Device(DEVICE_GPU),FarthestPointSampleGpuOp); 124 | 125 | void gatherpointLauncher(int b,int n,int m,const float * inp,const int * idx,float * out); 126 | class GatherPointGpuOp: public OpKernel{ 127 | public: 128 | explicit GatherPointGpuOp(OpKernelConstruction * context):OpKernel(context){} 129 | void Compute(OpKernelContext * context)override{ 130 | const Tensor& inp_tensor=context->input(0); 131 | OP_REQUIRES(context,inp_tensor.dims()==3 && inp_tensor.shape().dim_size(2)==3,errors::InvalidArgument("GatherPoint expects (batch_size,num_points,3) inp shape")); 132 | int b=inp_tensor.shape().dim_size(0); 133 | int n=inp_tensor.shape().dim_size(1); 134 | const Tensor& idx_tensor=context->input(1); 135 | OP_REQUIRES(context,idx_tensor.dims()==2 && idx_tensor.shape().dim_size(0)==b,errors::InvalidArgument("GatherPoint expects (batch_size,num_result) idx shape")); 136 | int m=idx_tensor.shape().dim_size(1); 137 | auto inp_flat=inp_tensor.flat(); 138 | const float * inp=&(inp_flat(0)); 139 | auto idx_flat=idx_tensor.flat(); 140 | const int * idx=&(idx_flat(0)); 141 | Tensor * out_tensor=NULL; 142 | OP_REQUIRES_OK(context,context->allocate_output(0,TensorShape{b,m,3},&out_tensor)); 143 | auto out_flat=out_tensor->flat(); 144 | float * out=&(out_flat(0)); 145 | gatherpointLauncher(b,n,m,inp,idx,out); 146 | } 147 | }; 148 | REGISTER_KERNEL_BUILDER(Name("GatherPoint").Device(DEVICE_GPU),GatherPointGpuOp); 149 | 150 | void scatteraddpointLauncher(int b,int n,int m,const float * out_g,const int * idx,float * inp_g); 151 | class GatherPointGradGpuOp: public OpKernel{ 152 | public: 153 | explicit GatherPointGradGpuOp(OpKernelConstruction * context):OpKernel(context){} 154 | void Compute(OpKernelContext * context)override{ 155 | const Tensor& inp_tensor=context->input(0); 156 | OP_REQUIRES(context,inp_tensor.dims()==3 && inp_tensor.shape().dim_size(2)==3,errors::InvalidArgument("GatherPointGradGpuOp expects (batch_size,num_points,3) inp")); 157 | int b=inp_tensor.shape().dim_size(0); 158 | int n=inp_tensor.shape().dim_size(1); 159 | const Tensor& idx_tensor=context->input(1); 160 | OP_REQUIRES(context,idx_tensor.dims()==2 && idx_tensor.shape().dim_size(0)==b,errors::InvalidArgument("GatherPointGradGpuOp expects (batch_size,num_result) idx shape")); 161 | int m=idx_tensor.shape().dim_size(1); 162 | auto inp_flat=inp_tensor.flat(); 163 | const float * inp=&(inp_flat(0)); 164 | auto idx_flat=idx_tensor.flat(); 165 | const int * idx=&(idx_flat(0)); 166 | const Tensor& out_g_tensor=context->input(2); 167 | OP_REQUIRES(context,out_g_tensor.dims()==3 && out_g_tensor.shape().dim_size(0)==b && out_g_tensor.shape().dim_size(1)==m && out_g_tensor.shape().dim_size(2)==3,errors::InvalidArgument("GatherPointGradGpuOp expects (batch_size,num_result,3) out_g shape")); 168 | auto out_g_flat=out_g_tensor.flat(); 169 | const float * out_g=&(out_g_flat(0)); 170 | Tensor * inp_g_tensor=NULL; 171 | OP_REQUIRES_OK(context,context->allocate_output(0,TensorShape{b,n,3},&inp_g_tensor)); 172 | auto inp_g_flat=inp_g_tensor->flat(); 173 | float * inp_g=&(inp_g_flat(0)); 174 | cudaMemset(inp_g,0,b*n*3*4); 175 | scatteraddpointLauncher(b,n,m,out_g,idx,inp_g); 176 | } 177 | }; 178 | REGISTER_KERNEL_BUILDER(Name("GatherPointGrad").Device(DEVICE_GPU),GatherPointGradGpuOp); 179 | 180 | -------------------------------------------------------------------------------- /tf_ops/sampling/tf_sampling.py: -------------------------------------------------------------------------------- 1 | ''' Furthest point sampling 2 | Original author: Haoqiang Fan 3 | Modified by Charles R. Qi 4 | All Rights Reserved. 2017. 5 | ''' 6 | import tensorflow as tf 7 | from tensorflow.python.framework import ops 8 | import sys 9 | import os 10 | import numpy as np 11 | import pickle as pickle 12 | 13 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 14 | sys.path.append(BASE_DIR) 15 | sampling_module=tf.load_op_library(os.path.join(BASE_DIR, 'tf_sampling_so.so')) 16 | def prob_sample(inp,inpr): 17 | ''' 18 | input: 19 | batch_size * ncategory float32 20 | batch_size * npoints float32 21 | returns: 22 | batch_size * npoints int32 23 | ''' 24 | return sampling_module.prob_sample(inp,inpr) 25 | ops.NoGradient('ProbSample') 26 | # TF1.0 API requires set shape in C++ 27 | #@tf.RegisterShape('ProbSample') 28 | #def _prob_sample_shape(op): 29 | # shape1=op.inputs[0].get_shape().with_rank(2) 30 | # shape2=op.inputs[1].get_shape().with_rank(2) 31 | # return [tf.TensorShape([shape2.dims[0],shape2.dims[1]])] 32 | def gather_point(inp,idx): 33 | ''' 34 | input: 35 | batch_size * ndataset * 3 float32 36 | batch_size * npoints int32 37 | returns: 38 | batch_size * npoints * 3 float32 39 | ''' 40 | return sampling_module.gather_point(inp,idx) 41 | #@tf.RegisterShape('GatherPoint') 42 | #def _gather_point_shape(op): 43 | # shape1=op.inputs[0].get_shape().with_rank(3) 44 | # shape2=op.inputs[1].get_shape().with_rank(2) 45 | # return [tf.TensorShape([shape1.dims[0],shape2.dims[1],shape1.dims[2]])] 46 | @tf.RegisterGradient('GatherPoint') 47 | def _gather_point_grad(op,out_g): 48 | inp=op.inputs[0] 49 | idx=op.inputs[1] 50 | return [sampling_module.gather_point_grad(inp,idx,out_g),None] 51 | def farthest_point_sample(npoint,inp): 52 | ''' 53 | input: 54 | int32 55 | batch_size * ndataset * 3 float32 56 | returns: 57 | batch_size * npoint int32 58 | ''' 59 | return sampling_module.farthest_point_sample(inp, npoint) 60 | ops.NoGradient('FarthestPointSample') 61 | 62 | 63 | if __name__=='__main__': 64 | 65 | 66 | batch_size = 3 67 | 68 | #np.random.seed(100) 69 | triangles=np.random.rand(batch_size,5,3,3).astype('float32') 70 | #pts=np.random.rand(batch_size,1024,3,3).astype('float32') 71 | 72 | inp=tf.constant(triangles) 73 | tria=inp[:,:,0,:] 74 | trib=inp[:,:,1,:] 75 | tric=inp[:,:,2,:] 76 | 77 | areas=tf.sqrt(tf.reduce_sum(tf.cross(trib-tria,tric-tria)**2,2)+1e-9) 78 | randomnumbers=tf.random_uniform((batch_size,8192))#(N,8192) 79 | triids=prob_sample(areas,randomnumbers) 80 | tria_sample=gather_point(tria,triids) 81 | trib_sample=gather_point(trib,triids) 82 | tric_sample=gather_point(tric,triids) 83 | us=tf.random_uniform((batch_size,8192)) 84 | vs=tf.random_uniform((batch_size,8192)) 85 | uplusv=1-tf.abs(us+vs-1) 86 | uminusv=us-vs 87 | us=(uplusv+uminusv)*0.5 88 | vs=(uplusv-uminusv)*0.5 89 | pt_sample=tria_sample+(trib_sample-tria_sample)*tf.expand_dims(us,-1)+(tric_sample-tria_sample)*tf.expand_dims(vs,-1) 90 | test = farthest_point_sample(1024,pt_sample) 91 | reduced_sample=gather_point(pt_sample,farthest_point_sample(1024,pt_sample)) 92 | 93 | with tf.Session() as sess: 94 | ret=sess.run(reduced_sample) 95 | pt = sess.run(pt_sample) 96 | 97 | print("tria:",tria.shape) 98 | print("areas:",areas.shape) 99 | print("triids:",triids.shape) 100 | print("tria_sample:",tria_sample.shape) 101 | print("pt_sample:",pt.shape,pt.dtype) 102 | print("test:",test.shape) 103 | print("reduced_sample",ret.shape,ret.dtype) 104 | 105 | 106 | #pickle.dump(ret,open('1.pkl','wb'),-1) 107 | print("done") 108 | -------------------------------------------------------------------------------- /tf_ops/sampling/tf_sampling_compile.sh: -------------------------------------------------------------------------------- 1 | #/bin/bash 2 | PYTHON=python3 3 | CUDA_PATH=/usr/local/cuda 4 | TF_LIB=$($PYTHON -c 'import tensorflow as tf; print(tf.sysconfig.get_lib())') 5 | PYTHON_VERSION=$($PYTHON -c 'import sys; print("%d.%d"%(sys.version_info[0], sys.version_info[1]))') 6 | TF_PATH=$TF_LIB/include 7 | $CUDA_PATH/bin/nvcc tf_sampling_g.cu -o tf_sampling_g.cu.o -c -O2 -DGOOGLE_CUDA=1 -x cu -Xcompiler -fPIC 8 | g++ -std=c++11 tf_sampling.cpp tf_sampling_g.cu.o -o tf_sampling_so.so -shared -fPIC -L$TF_LIB -ltensorflow_framework -I $TF_PATH/external/nsync/public/ -I $TF_PATH -I $CUDA_PATH/include -lcudart -L $CUDA_PATH/lib64/ -O2 -D_GLIBCXX_USE_CXX11_ABI=0 9 | -------------------------------------------------------------------------------- /tf_ops/sampling/tf_sampling_g.cu: -------------------------------------------------------------------------------- 1 | /* Furthest point sampling GPU implementation 2 | * Original author: Haoqiang Fan 3 | * Modified by Charles R. Qi 4 | * All Rights Reserved. 2017. 5 | */ 6 | 7 | __global__ void cumsumKernel(int b,int n,const float * __restrict__ inp,float * __restrict__ out){ 8 | const int BlockSize=2048; 9 | const int paddingLevel=5; 10 | __shared__ float buffer4[BlockSize*4]; 11 | __shared__ float buffer[BlockSize+(BlockSize>>paddingLevel)]; 12 | for (int i=blockIdx.x;i>2; 18 | for (int k=threadIdx.x*4;k>2)+(k>>(2+paddingLevel))]=v4; 33 | }else{ 34 | float v=0; 35 | for (int k2=k;k2>2)+(k>>(2+paddingLevel))]=v; 43 | } 44 | } 45 | int u=0; 46 | for (;(2<>(u+1));k+=blockDim.x){ 49 | int i1=(((k<<1)+2)<>paddingLevel; 52 | i2+=i2>>paddingLevel; 53 | buffer[i1]+=buffer[i2]; 54 | } 55 | } 56 | u--; 57 | for (;u>=0;u--){ 58 | __syncthreads(); 59 | for (int k=threadIdx.x;k>(u+1));k+=blockDim.x){ 60 | int i1=(((k<<1)+3)<>paddingLevel; 63 | i2+=i2>>paddingLevel; 64 | buffer[i1]+=buffer[i2]; 65 | } 66 | } 67 | __syncthreads(); 68 | for (int k=threadIdx.x*4;k>2)-1)+(((k>>2)-1)>>paddingLevel); 71 | buffer4[k]+=buffer[k2]; 72 | buffer4[k+1]+=buffer[k2]; 73 | buffer4[k+2]+=buffer[k2]; 74 | buffer4[k+3]+=buffer[k2]; 75 | } 76 | } 77 | __syncthreads(); 78 | for (int k=threadIdx.x;k>paddingLevel)]+runningsum2; 82 | float r2=runningsum+t; 83 | runningsum2=t-(r2-runningsum); 84 | runningsum=r2; 85 | __syncthreads(); 86 | } 87 | } 88 | } 89 | 90 | __global__ void binarysearchKernel(int b,int n,int m,const float * __restrict__ dataset,const float * __restrict__ query, int * __restrict__ result){ 91 | int base=1; 92 | while (base=1;k>>=1) 99 | if (r>=k && dataset[i*n+r-k]>=q) 100 | r-=k; 101 | result[i*m+j]=r; 102 | } 103 | } 104 | } 105 | __global__ void farthestpointsamplingKernel(int b,int n,int m,const float * __restrict__ dataset,float * __restrict__ temp,int * __restrict__ idxs){ 106 | if (m<=0) 107 | return; 108 | const int BlockSize=512; 109 | __shared__ float dists[BlockSize]; 110 | __shared__ int dists_i[BlockSize]; 111 | const int BufferSize=3072; 112 | __shared__ float buf[BufferSize*3]; 113 | for (int i=blockIdx.x;ibest){ 147 | best=d2; 148 | besti=k; 149 | } 150 | } 151 | dists[threadIdx.x]=best; 152 | dists_i[threadIdx.x]=besti; 153 | for (int u=0;(1<>(u+1))){ 156 | int i1=(threadIdx.x*2)<>>(b,n,inp,out); 196 | } 197 | //require b*n working space 198 | void probsampleLauncher(int b,int n,int m,const float * inp_p,const float * inp_r,float * temp,int * out){ 199 | cumsumKernel<<<32,512>>>(b,n,inp_p,temp); 200 | binarysearchKernel<<>>(b,n,m,temp,inp_r,out); 201 | } 202 | //require 32*n working space 203 | void farthestpointsamplingLauncher(int b,int n,int m,const float * inp,float * temp,int * out){ 204 | farthestpointsamplingKernel<<<32,512>>>(b,n,m,inp,temp,out); 205 | } 206 | void gatherpointLauncher(int b,int n,int m,const float * inp,const int * idx,float * out){ 207 | gatherpointKernel<<>>(b,n,m,inp,idx,out); 208 | } 209 | void scatteraddpointLauncher(int b,int n,int m,const float * out_g,const int * idx,float * inp_g){ 210 | scatteraddpointKernel<<>>(b,n,m,out_g,idx,inp_g); 211 | } 212 | 213 | -------------------------------------------------------------------------------- /train_val_cls.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | import argparse 4 | import importlib 5 | import numpy as np 6 | from time import time 7 | import tensorflow as tf 8 | from datetime import datetime 9 | 10 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 11 | sys.path.append(BASE_DIR) 12 | sys.path.append(os.path.join(BASE_DIR, 'utils')) 13 | sys.path.append(os.path.join(BASE_DIR, 'setting')) 14 | from utils import provider 15 | 16 | LOG_FOUT = None 17 | def log_string(out_str): 18 | LOG_FOUT.write(out_str+'\n') 19 | LOG_FOUT.flush() 20 | print(out_str) 21 | 22 | def main(): 23 | parser = argparse.ArgumentParser() 24 | parser.add_argument('--save_folder', '-s', default='log/cls/', help='Path to folder for saving check points and summary') 25 | parser.add_argument('--model', '-m', default='shellconv', help='Model to use') 26 | parser.add_argument('--setting', '-x', default='cls_modelnet40', help='Setting to use') 27 | parser.add_argument('--log', help='Log to FILE in save folder; use - for stdout (default is log.txt)', metavar='FILE', default='log.txt') 28 | args = parser.parse_args() 29 | 30 | time_string = datetime.now().strftime('%Y-%m-%d-%H-%M-%S') 31 | root_folder = os.path.join(args.save_folder, '%s_%s_%s' % (args.model, args.setting, time_string)) 32 | 33 | if not os.path.exists(root_folder): 34 | os.makedirs(root_folder) 35 | 36 | global LOG_FOUT 37 | if args.log != '-': 38 | LOG_FOUT = open(os.path.join(root_folder, args.log), 'w') 39 | 40 | model = importlib.import_module(args.model) 41 | setting_path = os.path.join(os.path.dirname(__file__), args.model) 42 | sys.path.append(setting_path) 43 | setting = importlib.import_module(args.setting) 44 | 45 | num_epochs = setting.num_epochs 46 | batch_size = setting.batch_size 47 | batch_size_val = setting.batch_size_val 48 | sample_num = setting.sample_num 49 | 50 | # Prepare inputs 51 | print('{}-Preparing datasets...'.format(datetime.now())) 52 | [data_train, label_train] = provider.load_cls_files(setting.filelist) 53 | [data_val, label_val] = provider.load_cls_files(setting.filelist_val) 54 | data_train = data_train[:,0:sample_num,:] 55 | data_val = data_val[:,0:sample_num,:] 56 | label_train = np.squeeze(label_train) 57 | label_val = np.squeeze(label_val) 58 | num_train = label_train.shape[0] 59 | num_val = label_val.shape[0] 60 | num_batch_train = int((num_train+batch_size) / batch_size) 61 | num_batch_val = int(num_val / batch_size_val) 62 | 63 | num_extra = num_batch_train * batch_size - num_train 64 | choices = np.random.choice(num_train, num_extra, replace=False) 65 | data_extra = data_train[choices, :, :] 66 | label_extra = label_train[choices] 67 | data_train = np.concatenate((data_train, data_extra), 0) 68 | label_train = np.concatenate((label_train, label_extra), 0) 69 | 70 | ###################################################################### 71 | 72 | global_step = tf.Variable(0, trainable=False, name='global_step') 73 | is_training_pl = tf.placeholder(tf.bool, name='is_training') 74 | 75 | data_pl = tf.placeholder(tf.float32, [None, sample_num, setting.data_dim], name='data_pl') 76 | label_pl = tf.placeholder(tf.int64, [None], name='label_pl') 77 | 78 | bn_decay_exp_op = tf.train.exponential_decay(0.5, global_step*batch_size, setting.decay_steps, 79 | 0.5, staircase=True) 80 | bn_decay_op = tf.minimum(0.99, 1 - bn_decay_exp_op) 81 | logits_op = model.get_model(data_pl, is_training_pl, setting.sconv_params, setting.sdconv_params, setting.fc_params, 82 | weight_decay=setting.weight_decay, 83 | bn_decay = bn_decay_op, 84 | part_num=setting.num_class) 85 | 86 | # compute loss 87 | if setting.with_multi: 88 | labels_2d = tf.expand_dims(label_pl, axis=-1, name='labels_2d') 89 | label_pl = tf.tile(labels_2d, (1, logits_op.shape[1]), name='labels_tile') 90 | predictions_op = tf.argmax(tf.reduce_mean(logits_op, axis = -2), axis=-1, name='predictions') 91 | else: 92 | predictions_op = tf.argmax(logits_op, axis=-1, name='predictions') 93 | 94 | loss_op = tf.losses.sparse_softmax_cross_entropy(labels=label_pl, logits=logits_op) 95 | 96 | lr_exp_op = tf.train.exponential_decay(setting.learning_rate_base, global_step*batch_size, setting.decay_steps, 97 | setting.decay_rate, staircase=True) 98 | lr_clip_op = tf.maximum(lr_exp_op, setting.learning_rate_min) 99 | _ = tf.summary.scalar('learning_rate', tensor=lr_clip_op, collections=['train']) 100 | 101 | if setting.optimizer == 'adam': 102 | optimizer = tf.train.AdamOptimizer(learning_rate=lr_clip_op) 103 | elif setting.optimizer == 'momentum': 104 | optimizer = tf.train.MomentumOptimizer(learning_rate=lr_clip_op, momentum=setting.momentum, use_nesterov=True) 105 | 106 | update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) 107 | with tf.control_dependencies(update_ops): 108 | train_op = optimizer.minimize(loss_op, global_step=global_step) 109 | 110 | saver = tf.train.Saver(max_to_keep=None) 111 | 112 | folder_ckpt = os.path.join(root_folder, 'ckpts') 113 | if not os.path.exists(folder_ckpt): 114 | os.makedirs(folder_ckpt) 115 | 116 | folder_summary = os.path.join(root_folder, 'summary') 117 | if not os.path.exists(folder_summary): 118 | os.makedirs(folder_summary) 119 | 120 | parameter_num = np.sum([np.prod(v.shape.as_list()) for v in tf.trainable_variables()]) 121 | print('{}-Parameter number: {:d}.'.format(datetime.now(), parameter_num)) 122 | 123 | shape_names = [line.rstrip() for line in \ 124 | open(setting.shape_name_file)] 125 | 126 | config = tf.ConfigProto() 127 | config.gpu_options.allow_growth = True 128 | config.allow_soft_placement = True 129 | config.log_device_placement = False 130 | # set config to save more gpu memory 131 | with tf.Session(config=config) as sess: 132 | summaries_op = tf.summary.merge_all('train') 133 | summary_writer = tf.summary.FileWriter(folder_summary, sess.graph) 134 | 135 | sess.run(tf.global_variables_initializer()) 136 | 137 | best_acc = 0 138 | best_epoch = 0 139 | start_time = time() 140 | for epoch in range(num_epochs): 141 | log_string('\n----------------------------- EPOCH %03d -----------------------------' % (epoch)) 142 | 143 | # train 144 | total_correct = 0 145 | total_seen = 0 146 | loss_avg = 0. 147 | data_train, label_train = provider.shuffle_data(data_train, label_train) 148 | for batch_idx in range(num_batch_train): 149 | start_time_sub = time() 150 | 151 | start_idx = batch_idx * batch_size 152 | end_idx = (batch_idx+1) * batch_size 153 | if setting.with_multi: 154 | labels_2d = np.expand_dims(label_train[start_idx:end_idx], axis=-1) 155 | labels_cur = np.tile(labels_2d, (1, setting.sconv_params[-1]['P'])) 156 | else: 157 | labels_cur = label_train[start_idx:end_idx] 158 | 159 | feed_dict = {data_pl: data_train[start_idx:end_idx, :, :], 160 | label_pl: labels_cur, 161 | is_training_pl: True,} 162 | 163 | _, loss, predictions, summaries, step = sess.run([train_op, loss_op, predictions_op, summaries_op, global_step], feed_dict=feed_dict) 164 | 165 | summary_writer.add_summary(summaries, step) 166 | 167 | correct = np.sum(predictions == label_train[start_idx:end_idx]) 168 | total_correct += correct 169 | total_seen += batch_size 170 | 171 | loss_avg = loss_avg*batch_idx/(batch_idx+1) + loss/(batch_idx+1) 172 | print('[ep:%d/%d: %d/%d] train acc(oa): %f -- loss: %f -- time cost: %f' % \ 173 | (epoch, num_epochs, batch_idx, num_batch_train, total_correct/float(total_seen), loss_avg, (time() - start_time_sub))) 174 | 175 | log_string('training -- acc(oa): %f -- loss: %f ' % (total_correct / float(total_seen), loss_avg)) 176 | log_string("Epoch %d, time cost: %f" % (epoch, time() - start_time)) 177 | 178 | # Validation 179 | total_correct = 0 180 | total_seen = 0 181 | loss_sum = 0 182 | total_seen_class = [0 for _ in range(setting.num_class)] 183 | total_correct_class = [0 for _ in range(setting.num_class)] 184 | 185 | eval_start_time = time() # eval start time 186 | for batch_idx in range(num_batch_val): 187 | start_idx = batch_idx * batch_size_val 188 | end_idx = (batch_idx+1) * batch_size_val 189 | if setting.with_multi: 190 | labels_2d = np.expand_dims(label_val[start_idx:end_idx], axis=-1) 191 | labels_cur = np.tile(labels_2d, (1, setting.sconv_params[-1]['P'])) 192 | else: 193 | labels_cur = label_val[start_idx:end_idx] 194 | 195 | feed_dict = {data_pl: data_val[start_idx:end_idx, :, :], 196 | label_pl: labels_cur, 197 | is_training_pl: False,} 198 | 199 | loss, predictions_val = sess.run([loss_op, predictions_op], feed_dict=feed_dict) 200 | 201 | correct = np.sum(predictions_val == label_val[start_idx:end_idx]) 202 | total_correct += correct 203 | total_seen += batch_size_val 204 | loss_sum += (loss*batch_size_val) 205 | for i in range(start_idx, end_idx): 206 | l = label_val[i] 207 | total_seen_class[l] += 1 208 | total_correct_class[l] += (predictions_val[i-start_idx] == l) 209 | 210 | class_accuracies = np.array(total_correct_class)/np.array(total_seen_class,dtype=np.float) 211 | acc_mean_cls = np.mean(np.array(total_correct_class)/np.array(total_seen_class,dtype=np.float)) 212 | acc_oa = total_correct / float(total_seen) 213 | loss_val = loss_sum / float(total_seen) 214 | 215 | log_string('eval acc (oa): %f ---- eval acc (mean class): %f ---- loss: %f' % \ 216 | (acc_oa, acc_mean_cls, loss_val)) 217 | 218 | print('eval time: %f' % (time() - eval_start_time)) 219 | if acc_oa > best_acc: 220 | save_path = saver.save(sess, os.path.join(folder_ckpt, "best_model_epoch.ckpt")) 221 | best_acc = acc_oa 222 | best_epoch = epoch 223 | 224 | with open(os.path.join(root_folder,"class_accuracies.txt"), 'w') as the_file: 225 | the_file.write('best epoch: %f \n' % (best_epoch)) 226 | for i, name in enumerate(shape_names): 227 | print('%10s:\t%0.3f' % (name, class_accuracies[i])) 228 | the_file.write('%10s:\t%0.3f\n' % (name, class_accuracies[i])) 229 | the_file.write('%10s:\t%0.3f\n' % ('mean class acc: ', acc_mean_cls)) 230 | 231 | if __name__ == '__main__': 232 | main() 233 | if LOG_FOUT is not None: 234 | LOG_FOUT.close() -------------------------------------------------------------------------------- /utils/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hkust-vgd/shellnet/fe5cf5f3806e618c23e161b0fb8a57387268df78/utils/__init__.py -------------------------------------------------------------------------------- /utils/data/.gitignore: -------------------------------------------------------------------------------- 1 | # Ignore everything in this directory 2 | * 3 | # Except this file 4 | !.gitignore -------------------------------------------------------------------------------- /utils/data_utils.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from __future__ import division 3 | from __future__ import print_function 4 | 5 | import os 6 | import h5py 7 | import plyfile 8 | import numpy as np 9 | from matplotlib import cm 10 | 11 | 12 | def save_ply(points, filename, colors=None, normals=None): 13 | vertex = np.core.records.fromarrays(points.transpose(), names='x, y, z', formats='f4, f4, f4') 14 | n = len(vertex) 15 | desc = vertex.dtype.descr 16 | 17 | if normals is not None: 18 | vertex_normal = np.core.records.fromarrays(normals.transpose(), names='nx, ny, nz', formats='f4, f4, f4') 19 | assert len(vertex_normal) == n 20 | desc = desc + vertex_normal.dtype.descr 21 | 22 | if colors is not None: 23 | vertex_color = np.core.records.fromarrays(colors.transpose() * 255, names='red, green, blue', 24 | formats='u1, u1, u1') 25 | assert len(vertex_color) == n 26 | desc = desc + vertex_color.dtype.descr 27 | 28 | vertex_all = np.empty(n, dtype=desc) 29 | 30 | for prop in vertex.dtype.names: 31 | vertex_all[prop] = vertex[prop] 32 | 33 | if normals is not None: 34 | for prop in vertex_normal.dtype.names: 35 | vertex_all[prop] = vertex_normal[prop] 36 | 37 | if colors is not None: 38 | for prop in vertex_color.dtype.names: 39 | vertex_all[prop] = vertex_color[prop] 40 | 41 | ply = plyfile.PlyData([plyfile.PlyElement.describe(vertex_all, 'vertex')], text=False) 42 | if not os.path.exists(os.path.dirname(filename)): 43 | os.makedirs(os.path.dirname(filename)) 44 | ply.write(filename) 45 | 46 | 47 | def save_ply_property(points, property, property_max, filename, cmap_name='tab20'): 48 | point_num = points.shape[0] 49 | colors = np.full(points.shape, 0.5) 50 | cmap = cm.get_cmap(cmap_name) 51 | for point_idx in range(point_num): 52 | if property[point_idx] == 0: 53 | colors[point_idx] = np.array([0, 0, 0]) 54 | else: 55 | colors[point_idx] = cmap(property[point_idx] / property_max)[:3] 56 | save_ply(points, filename, colors) 57 | 58 | 59 | def save_ply_batch(points_batch, file_path, points_num=None): 60 | batch_size = points_batch.shape[0] 61 | if type(file_path) != list: 62 | basename = os.path.splitext(file_path)[0] 63 | ext = '.ply' 64 | for batch_idx in range(batch_size): 65 | point_num = points_batch.shape[1] if points_num is None else points_num[batch_idx] 66 | if type(file_path) == list: 67 | save_ply(points_batch[batch_idx][:point_num], file_path[batch_idx]) 68 | else: 69 | save_ply(points_batch[batch_idx][:point_num], '%s_%04d%s' % (basename, batch_idx, ext)) 70 | 71 | 72 | def save_ply_color_batch(points_batch, colors_batch, file_path, points_num=None): 73 | batch_size = points_batch.shape[0] 74 | if type(file_path) != list: 75 | basename = os.path.splitext(file_path)[0] 76 | ext = '.ply' 77 | for batch_idx in range(batch_size): 78 | point_num = points_batch.shape[1] if points_num is None else points_num[batch_idx] 79 | if type(file_path) == list: 80 | save_ply(points_batch[batch_idx][:point_num], file_path[batch_idx], colors_batch[batch_idx][:point_num]) 81 | else: 82 | save_ply(points_batch[batch_idx][:point_num], '%s_%04d%s' % (basename, batch_idx, ext), 83 | colors_batch[batch_idx][:point_num]) 84 | 85 | 86 | def save_ply_property_batch(points_batch, property_batch, file_path, points_num=None, property_max=None, 87 | cmap_name='tab20'): 88 | batch_size = points_batch.shape[0] 89 | if type(file_path) != list: 90 | basename = os.path.splitext(file_path)[0] 91 | ext = '.ply' 92 | property_max = np.max(property_batch) if property_max is None else property_max 93 | for batch_idx in range(batch_size): 94 | point_num = points_batch.shape[1] if points_num is None else points_num[batch_idx] 95 | if type(file_path) == list: 96 | save_ply_property(points_batch[batch_idx][:point_num], property_batch[batch_idx][:point_num], 97 | property_max, file_path[batch_idx], cmap_name) 98 | else: 99 | save_ply_property(points_batch[batch_idx][:point_num], property_batch[batch_idx][:point_num], 100 | property_max, '%s_%04d%s' % (basename, batch_idx, ext), cmap_name) 101 | 102 | 103 | def save_ply_point_with_normal(data_sample, folder): 104 | for idx, sample in enumerate(data_sample): 105 | filename_pts = os.path.join(folder, '{:08d}.ply'.format(idx)) 106 | save_ply(sample[..., :3], filename_pts, normals=sample[..., 3:]) 107 | 108 | 109 | def grouped_shuffle(inputs): 110 | for idx in range(len(inputs) - 1): 111 | assert (len(inputs[idx]) == len(inputs[idx + 1])) 112 | 113 | shuffle_indices = np.arange(inputs[0].shape[0]) 114 | np.random.shuffle(shuffle_indices) 115 | outputs = [] 116 | for idx in range(len(inputs)): 117 | outputs.append(inputs[idx][shuffle_indices, ...]) 118 | 119 | 120 | return outputs 121 | 122 | 123 | def load_cls(filelist): 124 | points = [] 125 | labels = [] 126 | 127 | folder = os.path.dirname(filelist) 128 | for line in open(filelist): 129 | filename = os.path.basename(line.rstrip()) 130 | data = h5py.File(os.path.join(folder, filename)) 131 | if 'normal' in data: 132 | points.append(np.concatenate([data['data'][...], data['normal'][...]], axis=-1).astype(np.float32)) 133 | else: 134 | points.append(data['data'][...].astype(np.float32)) 135 | labels.append(np.squeeze(data['label'][:]).astype(np.int64)) 136 | return (np.concatenate(points, axis=0), 137 | np.concatenate(labels, axis=0)) 138 | 139 | 140 | def load_cls_train_val(filelist, filelist_val): 141 | data_train, label_train = grouped_shuffle(load_cls(filelist)) 142 | data_val, label_val = load_cls(filelist_val) 143 | return data_train, label_train, data_val, label_val 144 | 145 | 146 | def is_h5_list(filelist): 147 | return all([line.strip()[-3:] == '.h5' for line in open(filelist)]) 148 | 149 | 150 | def load_seg_list(filelist): 151 | folder = os.path.dirname(filelist) 152 | return [os.path.join(folder, line.strip()) for line in open(filelist)] 153 | 154 | 155 | def load_seg(filelist): 156 | points = [] 157 | labels = [] 158 | point_nums = [] 159 | labels_seg = [] 160 | indices_split_to_full = [] 161 | 162 | folder = os.path.dirname(filelist) 163 | for line in open(filelist): 164 | data = h5py.File(os.path.join(folder, line.strip())) 165 | points.append(data['data'][...].astype(np.float32)) 166 | labels.append(data['label'][...].astype(np.int64)) 167 | point_nums.append(data['data_num'][...].astype(np.int32)) 168 | labels_seg.append(data['label_seg'][...].astype(np.int64)) 169 | if 'indices_split_to_full' in data: 170 | indices_split_to_full.append(data['indices_split_to_full'][...].astype(np.int64)) 171 | 172 | return (np.concatenate(points, axis=0), 173 | np.concatenate(labels, axis=0), 174 | np.concatenate(point_nums, axis=0), 175 | np.concatenate(labels_seg, axis=0), 176 | np.concatenate(indices_split_to_full, axis=0) if indices_split_to_full else None) 177 | 178 | 179 | def balance_classes(labels): 180 | _, inverse, counts = np.unique(labels, return_inverse=True, return_counts=True) 181 | counts_max = np.amax(counts) 182 | repeat_num_avg_unique = counts_max / counts 183 | repeat_num_avg = repeat_num_avg_unique[inverse] 184 | repeat_num_floor = np.floor(repeat_num_avg) 185 | repeat_num_probs = repeat_num_avg - repeat_num_floor 186 | repeat_num = repeat_num_floor + (np.random.rand(repeat_num_probs.shape[0]) < repeat_num_probs) 187 | 188 | return repeat_num.astype(np.int64) 189 | -------------------------------------------------------------------------------- /utils/provider.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | import numpy as np 4 | import h5py 5 | 6 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 7 | sys.path.append(BASE_DIR) 8 | 9 | # Download dataset for point cloud classification 10 | DATA_DIR = os.path.join(BASE_DIR, 'data') 11 | if not os.path.exists(DATA_DIR): 12 | os.mkdir(DATA_DIR) 13 | 14 | def shuffle_data_seg(data, labels): 15 | """ Shuffle data and labels. 16 | Input: 17 | data: B,N,... numpy array 18 | label: B,... numpy array 19 | Return: 20 | shuffled data, label and shuffle indices 21 | """ 22 | idx = np.arange(len(labels)) 23 | np.random.shuffle(idx) 24 | return data[idx, ...], labels[idx], idx 25 | 26 | def shuffle_data(data, labels): 27 | """ Shuffle data and labels. 28 | Input: 29 | data: B,N,... numpy array 30 | label: B,... numpy array 31 | Return: 32 | shuffled data, label and shuffle indices 33 | """ 34 | point_num = data.shape[1] 35 | idx = np.arange(point_num) 36 | np.random.shuffle(idx) 37 | data = data[:,idx,:] 38 | 39 | idx = np.arange(len(labels)) 40 | np.random.shuffle(idx) 41 | 42 | return data[idx, ...], labels[idx] 43 | 44 | 45 | 46 | def jitter_point_cloud(batch_data, sigma=0.01, clip=0.05): 47 | """ Randomly jitter points. jittering is per point. 48 | Input: 49 | BxNx3 array, original batch of point clouds 50 | Return: 51 | BxNx3 array, jittered batch of point clouds 52 | """ 53 | B, N, C = batch_data.shape 54 | assert(clip > 0) 55 | jittered_data = np.clip(sigma * np.random.randn(B, N, C), -1*clip, clip) 56 | jittered_data += batch_data 57 | return jittered_data 58 | 59 | def random_point_dropout(batch_pc, max_dropout_ratio=0.875): 60 | ''' batch_pc: BxNx3 ''' 61 | for b in range(batch_pc.shape[0]): 62 | dropout_ratio = np.random.random()*max_dropout_ratio # 0~0.875 63 | drop_idx = np.where(np.random.random((batch_pc.shape[1]))<=dropout_ratio)[0] 64 | if len(drop_idx)>0: 65 | batch_pc[b,drop_idx,:] = batch_pc[b,0,:] # set to the first point 66 | return batch_pc 67 | 68 | def getDataFiles(list_filename): 69 | return [line.rstrip() for line in open(list_filename)] 70 | 71 | def load_cls_files(filelist): 72 | points = [] 73 | labels = [] 74 | 75 | folder = os.path.dirname(filelist) 76 | for line in open(filelist): 77 | [pts, label]= load_h5(os.path.join(folder, os.path.basename(line.strip()))) 78 | points.append(pts) 79 | labels.append(label) 80 | 81 | return (np.concatenate(points, axis=0), 82 | np.concatenate(labels, axis=0)) 83 | 84 | def load_h5(h5_filename): 85 | f = h5py.File(h5_filename) 86 | data = f['data'][:] 87 | label = f['label'][:] 88 | return (data, label) 89 | 90 | def loadDataFile(filename): 91 | return load_h5(filename) 92 | 93 | 94 | 95 | def load_h5_data_label_seg(h5_filename): 96 | f = h5py.File(h5_filename) 97 | data = f['data'][:] 98 | label = f['label'][:] 99 | seg = f['pid'][:] 100 | return (data, label, seg) 101 | 102 | 103 | def loadDataFile_with_seg(filename): 104 | return load_h5_data_label_seg(filename) 105 | 106 | 107 | 108 | def loadDataFile_with_normal(filename): 109 | f = h5py.File(filename) 110 | data = f['data'][:] 111 | label = f['label'][:] 112 | normal = f['normal'][:] 113 | return (data, label, normal) 114 | 115 | def loadDataFile_without_normal(filename): 116 | f = h5py.File(filename) 117 | data = f['data'][:] 118 | label = f['label'][:] 119 | return (data, label) 120 | 121 | def point_cloud_label_to_surface_voxel_label(point_cloud, label, res=0.0484): 122 | coordmax = np.max(point_cloud,axis=0) 123 | coordmin = np.min(point_cloud,axis=0) 124 | nvox = np.ceil((coordmax-coordmin)/res) 125 | vidx = np.ceil((point_cloud-coordmin)/res) 126 | vidx = vidx[:,0]+vidx[:,1]*nvox[0]+vidx[:,2]*nvox[0]*nvox[1] 127 | uvidx = np.unique(vidx) 128 | if label.ndim==1: 129 | uvlabel = [np.argmax(np.bincount(label[vidx==uv].astype(np.uint32))) for uv in uvidx] 130 | else: 131 | assert(label.ndim==2) 132 | uvlabel = np.zeros(len(uvidx),label.shape[1]) 133 | for i in range(label.shape[1]): 134 | uvlabel[:,i] = np.array([np.argmax(np.bincount(label[vidx==uv,i].astype(np.uint32))) for uv in uvidx]) 135 | return uvidx, uvlabel, nvox 136 | 137 | def point_cloud_label_to_surface_voxel_label_fast(point_cloud, label, res=0.0484): 138 | coordmax = np.max(point_cloud,axis=0) 139 | coordmin = np.min(point_cloud,axis=0) 140 | nvox = np.ceil((coordmax-coordmin)/res) 141 | vidx = np.ceil((point_cloud-coordmin)/res) 142 | vidx = vidx[:,0]+vidx[:,1]*nvox[0]+vidx[:,2]*nvox[0]*nvox[1] 143 | uvidx, vpidx = np.unique(vidx,return_index=True) 144 | if label.ndim==1: 145 | uvlabel = label[vpidx] 146 | else: 147 | assert(label.ndim==2) 148 | uvlabel = label[vpidx,:] 149 | return uvidx, uvlabel, nvox 150 | 151 | def rotate_point_cloud(batch_data, rotation_axis="z"): 152 | """ Randomly rotate the point clouds to augument the dataset 153 | rotation is per shape based along up direction 154 | Input: 155 | BxNx3 array, original batch of point clouds 156 | Return: 157 | BxNx3 array, rotated batch of point clouds 158 | """ 159 | if np.ndim(batch_data) != 3: 160 | raise ValueError("np.ndim(batch_data) != 3, must be (b, n, 3)") 161 | if batch_data.shape[2] != 3: 162 | raise ValueError("batch_data.shape[2] != 3, must be (x, y, z)") 163 | 164 | rotated_data = np.zeros(batch_data.shape, dtype=np.float32) 165 | for k in range(batch_data.shape[0]): 166 | rotation_angle = np.random.uniform() * 2 * np.pi 167 | cosval = np.cos(rotation_angle) 168 | sinval = np.sin(rotation_angle) 169 | if rotation_axis == "x": 170 | rotation_matrix = np.array( 171 | [[1, 0, 0], [0, cosval, sinval], [0, -sinval, cosval]] 172 | ) 173 | elif rotation_axis == "y": 174 | rotation_matrix = np.array( 175 | [[cosval, 0, sinval], [0, 1, 0], [-sinval, 0, cosval]] 176 | ) 177 | elif rotation_axis == "z": 178 | rotation_matrix = np.array( 179 | [[cosval, sinval, 0], [-sinval, cosval, 0], [0, 0, 1]] 180 | ) 181 | else: 182 | raise ValueError("Wrong rotation axis") 183 | shape_pc = batch_data[k, ...] 184 | rotated_data[k, ...] = np.dot(shape_pc.reshape((-1, 3)), rotation_matrix) 185 | return rotated_data 186 | 187 | 188 | def rotate_feature_point_cloud(batch_data, feature_size=3, rotation_axis="z"): 189 | """ Randomly rotate the point clouds to augument the dataset 190 | rotation is per shape based along up direction 191 | Input: 192 | BxNx3 array, original batch of point clouds 193 | Return: 194 | BxNx3 array, rotated batch of point clouds 195 | """ 196 | rotated_data = np.zeros(batch_data.shape, dtype=np.float32) 197 | rotated_data[:, :, 3 : 3 + feature_size] = batch_data[:, :, 3 : 3 + feature_size] 198 | for k in range(batch_data.shape[0]): 199 | rotation_angle = np.random.uniform() * 2 * np.pi 200 | cosval = np.cos(rotation_angle) 201 | sinval = np.sin(rotation_angle) 202 | if rotation_axis == "x": 203 | rotation_matrix = np.array( 204 | [[1, 0, 0], [0, cosval, sinval], [0, -sinval, cosval]] 205 | ) 206 | elif rotation_axis == "y": 207 | rotation_matrix = np.array( 208 | [[cosval, 0, sinval], [0, 1, 0], [-sinval, 0, cosval]] 209 | ) 210 | elif rotation_axis == "z": 211 | rotation_matrix = np.array( 212 | [[cosval, sinval, 0], [-sinval, cosval, 0], [0, 0, 1]] 213 | ) 214 | else: 215 | raise ValueError("Wrong rotation axis") 216 | shape_pc = batch_data[k, :, 0:3] 217 | rotated_data[k, :, 0:3] = np.dot(shape_pc.reshape((-1, 3)), rotation_matrix) 218 | return rotated_data 219 | 220 | 221 | def rotate_point_cloud_with_normal(batch_xyz_normal): 222 | """ Randomly rotate XYZ, normal point cloud. 223 | Input: 224 | batch_xyz_normal: B,N,6, first three channels are XYZ, last 3 all normal 225 | Output: 226 | B,N,6, rotated XYZ, normal point cloud 227 | """ 228 | for k in range(batch_xyz_normal.shape[0]): 229 | rotation_angle = np.random.uniform() * 2 * np.pi 230 | cosval = np.cos(rotation_angle) 231 | sinval = np.sin(rotation_angle) 232 | rotation_matrix = np.array( 233 | [[cosval, 0, sinval], [0, 1, 0], [-sinval, 0, cosval]] 234 | ) 235 | shape_pc = batch_xyz_normal[k, :, 0:3] 236 | shape_normal = batch_xyz_normal[k, :, 3:6] 237 | batch_xyz_normal[k, :, 0:3] = np.dot(shape_pc.reshape((-1, 3)), rotation_matrix) 238 | batch_xyz_normal[k, :, 3:6] = np.dot( 239 | shape_normal.reshape((-1, 3)), rotation_matrix 240 | ) 241 | return batch_xyz_normal 242 | 243 | 244 | def rotate_perturbation_point_cloud_with_normal( 245 | batch_data, angle_sigma=0.06, angle_clip=0.18 246 | ): 247 | """ Randomly perturb the point clouds by small rotations 248 | Input: 249 | BxNx6 array, original batch of point clouds and point normals 250 | Return: 251 | BxNx3 array, rotated batch of point clouds 252 | """ 253 | rotated_data = np.zeros(batch_data.shape, dtype=np.float32) 254 | for k in range(batch_data.shape[0]): 255 | angles = np.clip(angle_sigma * np.random.randn(3), -angle_clip, angle_clip) 256 | Rx = np.array( 257 | [ 258 | [1, 0, 0], 259 | [0, np.cos(angles[0]), -np.sin(angles[0])], 260 | [0, np.sin(angles[0]), np.cos(angles[0])], 261 | ] 262 | ) 263 | Ry = np.array( 264 | [ 265 | [np.cos(angles[1]), 0, np.sin(angles[1])], 266 | [0, 1, 0], 267 | [-np.sin(angles[1]), 0, np.cos(angles[1])], 268 | ] 269 | ) 270 | Rz = np.array( 271 | [ 272 | [np.cos(angles[2]), -np.sin(angles[2]), 0], 273 | [np.sin(angles[2]), np.cos(angles[2]), 0], 274 | [0, 0, 1], 275 | ] 276 | ) 277 | R = np.dot(Rz, np.dot(Ry, Rx)) 278 | shape_pc = batch_data[k, :, 0:3] 279 | shape_normal = batch_data[k, :, 3:6] 280 | rotated_data[k, :, 0:3] = np.dot(shape_pc.reshape((-1, 3)), R) 281 | rotated_data[k, :, 3:6] = np.dot(shape_normal.reshape((-1, 3)), R) 282 | return rotated_data 283 | 284 | 285 | def rotate_point_cloud_by_angle(batch_data, rotation_angle): 286 | """ Rotate the point cloud along up direction with certain angle. 287 | Input: 288 | BxNx3 array, original batch of point clouds 289 | Return: 290 | BxNx3 array, rotated batch of point clouds 291 | """ 292 | rotated_data = np.zeros(batch_data.shape, dtype=np.float32) 293 | for k in range(batch_data.shape[0]): 294 | # rotation_angle = np.random.uniform() * 2 * np.pi 295 | cosval = np.cos(rotation_angle) 296 | sinval = np.sin(rotation_angle) 297 | rotation_matrix = np.array( 298 | [[cosval, 0, sinval], [0, 1, 0], [-sinval, 0, cosval]] 299 | ) 300 | shape_pc = batch_data[k, :, 0:3] 301 | rotated_data[k, :, 0:3] = np.dot(shape_pc.reshape((-1, 3)), rotation_matrix) 302 | return rotated_data 303 | 304 | 305 | def rotate_point_cloud_by_angle_with_normal(batch_data, rotation_angle): 306 | """ Rotate the point cloud along up direction with certain angle. 307 | Input: 308 | BxNx3 array, original batch of point clouds 309 | Return: 310 | BxNx3 array, rotated batch of point clouds 311 | """ 312 | rotated_data = np.zeros(batch_data.shape, dtype=np.float32) 313 | for k in range(batch_data.shape[0]): 314 | # rotation_angle = np.random.uniform() * 2 * np.pi 315 | cosval = np.cos(rotation_angle) 316 | sinval = np.sin(rotation_angle) 317 | rotation_matrix = np.array( 318 | [[cosval, 0, sinval], [0, 1, 0], [-sinval, 0, cosval]] 319 | ) 320 | shape_pc = batch_data[k, ...] 321 | shape_normal = batch_data[k, :, 3:6] 322 | rotated_data[k, ...] = np.dot(shape_pc.reshape((-1, 3)), rotation_matrix) 323 | rotated_data[k, :, 3:6] = np.dot(shape_normal.reshape((-1, 3)), rotation_matrix) 324 | return rotated_data 325 | 326 | 327 | def rotate_perturbation_point_cloud(batch_data, angle_sigma=0.06, angle_clip=0.18): 328 | """ Randomly perturb the point clouds by small rotations 329 | Input: 330 | BxNx3 array, original batch of point clouds 331 | Return: 332 | BxNx3 array, rotated batch of point clouds 333 | """ 334 | rotated_data = np.zeros(batch_data.shape, dtype=np.float32) 335 | for k in range(batch_data.shape[0]): 336 | angles = np.clip(angle_sigma * np.random.randn(3), -angle_clip, angle_clip) 337 | Rx = np.array( 338 | [ 339 | [1, 0, 0], 340 | [0, np.cos(angles[0]), -np.sin(angles[0])], 341 | [0, np.sin(angles[0]), np.cos(angles[0])], 342 | ] 343 | ) 344 | Ry = np.array( 345 | [ 346 | [np.cos(angles[1]), 0, np.sin(angles[1])], 347 | [0, 1, 0], 348 | [-np.sin(angles[1]), 0, np.cos(angles[1])], 349 | ] 350 | ) 351 | Rz = np.array( 352 | [ 353 | [np.cos(angles[2]), -np.sin(angles[2]), 0], 354 | [np.sin(angles[2]), np.cos(angles[2]), 0], 355 | [0, 0, 1], 356 | ] 357 | ) 358 | R = np.dot(Rz, np.dot(Ry, Rx)) 359 | shape_pc = batch_data[k, ...] 360 | rotated_data[k, ...] = np.dot(shape_pc.reshape((-1, 3)), R) 361 | return rotated_data 362 | 363 | 364 | def shift_point_cloud(batch_data, shift_range=0.1): 365 | """ Randomly shift point cloud. Shift is per point cloud. 366 | Input: 367 | BxNx3 array, original batch of point clouds 368 | Return: 369 | BxNx3 array, shifted batch of point clouds 370 | """ 371 | B, N, C = batch_data.shape 372 | shifts = np.random.uniform(-shift_range, shift_range, (B, 3)) 373 | for batch_index in range(B): 374 | batch_data[batch_index, :, :] += shifts[batch_index, :] 375 | return batch_data 376 | 377 | 378 | def random_scale_point_cloud(batch_data, scale_low=0.8, scale_high=1.25): 379 | """ Randomly scale the point cloud. Scale is per point cloud. 380 | Input: 381 | BxNx3 array, original batch of point clouds 382 | Return: 383 | BxNx3 array, scaled batch of point clouds 384 | """ 385 | B, N, C = batch_data.shape 386 | scales = np.random.uniform(scale_low, scale_high, B) 387 | for batch_index in range(B): 388 | batch_data[batch_index, :, :] *= scales[batch_index] 389 | return batch_data 390 | 391 | 392 | def getDataFiles(list_filename): 393 | return [line.rstrip() for line in open(list_filename)] 394 | 395 | 396 | def load_h5(h5_filename): 397 | import h5py 398 | 399 | f = h5py.File(h5_filename) 400 | data = f["data"][:] 401 | label = f["label"][:] 402 | return (data, label) 403 | 404 | 405 | def loadDataFile(filename): 406 | return load_h5(filename) 407 | --------------------------------------------------------------------------------