├── LICENSE ├── README.md ├── SCFNet.py ├── compile_op.sh ├── helper_dp.py ├── helper_ply.py ├── helper_tf_util.py ├── img ├── abstract.png ├── architecture.png ├── s3dis.png ├── s3dis_vis.png ├── semantic3d.png └── semantic3d_vis.png ├── requirements.txt ├── s3dis_6fold.sh ├── s3dis_6fold_eval.py ├── s3dis_main.py ├── s3dis_test.py ├── semantic3d_main.py ├── semantic3d_test.py ├── utils ├── cpp_wrappers │ ├── compile_wrappers.sh │ ├── cpp_subsampling │ │ ├── grid_subsampling │ │ │ ├── grid_subsampling.cpp │ │ │ └── grid_subsampling.h │ │ ├── setup.py │ │ └── wrapper.cpp │ └── cpp_utils │ │ ├── cloud │ │ ├── cloud.cpp │ │ └── cloud.h │ │ └── nanoflann │ │ └── nanoflann.hpp ├── download_semantic3d.sh ├── meta │ ├── anno_paths.txt │ └── class_names.txt ├── nearest_neighbors │ ├── KDTreeTableAdaptor.h │ ├── knn.pyx │ ├── knn_.cxx │ ├── knn_.h │ ├── nanoflann.hpp │ ├── setup.py │ └── test.py ├── s3dis_dp.py └── semantic3d_dp.py └── visualization.py /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2020 Siqi Fan 4 | All rights reserved. 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. 23 | 24 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # SCF-Net: Learning Spatial Contextual Features for Large-Scale Point Cloud Segmentation 2 | 3 | [![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/scf-net-learning-spatial-contextual-features/semantic-segmentation-on-semantic3d)](https://paperswithcode.com/sota/semantic-segmentation-on-semantic3d?p=scf-net-learning-spatial-contextual-features) 4 | [![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/scf-net-learning-spatial-contextual-features/semantic-segmentation-on-s3dis)](https://paperswithcode.com/sota/semantic-segmentation-on-s3dis?p=scf-net-learning-spatial-contextual-features) 5 | [![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/scf-net-learning-spatial-contextual-features/3d-semantic-segmentation-on-stpls3d)](https://paperswithcode.com/sota/3d-semantic-segmentation-on-stpls3d?p=scf-net-learning-spatial-contextual-features) 6 | 7 | ![architecture](./img/architecture.png) 8 | 9 | > For technical details, please refer to: 10 | > 11 | > [SCF-Net: Learning Spatial Contextual Features for Large-Scale Point Cloud Segmentation](https://openaccess.thecvf.com/content/CVPR2021/html/Fan_SCF-Net_Learning_Spatial_Contextual_Features_for_Large-Scale_Point_Cloud_Segmentation_CVPR_2021_paper.html) 12 | 13 | ### (0) Abstract 14 | 15 | How to learn effective features from large-scale point clouds for semantic segmentation has attracted increasing attention in recent years. Addressing this problem, we propose a learnable module that learns Spatial Contextual Features from large-scale point clouds, called SCF in this paper. The proposed module mainly consists of three blocks, including the local polar representation block, the dual-distance attentive pooling block, and the global contextual feature block. For each 3D point, the local polar representation block is firstly explored to construct a spatial representation that is invariant to the z-axis rotation, then the dual-distance attentive pooling block is designed to utilize the representations of its neighbors for learning more discriminative local features according to both the geometric and feature distances among them, and finally, the global contextual feature block is designed to learn a global context for each 3D point by utilizing its spatial location and the volume ratio of the neighborhood to the global point cloud. The proposed module could be easily embedded into various network architectures for point cloud segmentation, naturally resulting in a new 3D semantic segmentation network with an encoder-decoder architecture, called SCF-Net in this work. Extensive experimental results on two public datasets demonstrate that the proposed SCF-Net performs better than several state-of-the-art methods in most cases. 16 | 17 | ![abstract](./img/abstract.png) 18 | 19 | ### (1) Setup 20 | 21 | This code has been tested with Python 3.5, Tensorflow 1.11, CUDA 9.0 and cuDNN v7 on Ubuntu 16.04. 22 | 23 | * Setup python environment 24 | ``` 25 | conda create -n SCF-Net python=3.5 26 | source activate SCF-Net 27 | ``` 28 | 29 | * Clone the repository 30 | 31 | * Install the requirements 32 | ``` 33 | pip install -r requirements.txt 34 | sh compile_op.sh 35 | ``` 36 | 37 | * Download pertained models 38 | 39 | | Dataset | Baidu Cloud | Google Drive | 40 | |--------------------------|----------------|----------------| 41 | | S3DIS | [link](https://pan.baidu.com/s/1QSdXmhuJYixxYFF70No4XQ): rw6m | [link](https://drive.google.com/file/d/1zZUAFLI_sdefhySIhBtd7eTW9u4dWV01/view?usp=sharing) | 42 | | Semantic3D | [link](https://pan.baidu.com/s/1rsbPDTUXH-4-XvDSb0K2-g): bzcz | [link](https://drive.google.com/file/d/1f5MumuaJ97ut4Z9kIv4F1e7JdUvxBYT4/view?usp=sharing) | 43 | 44 | ### (2) S3DIS 45 | 46 | ![S3DIS](./img/s3dis.png) 47 | 48 | * Setup the dataset 49 | * Download the "Stanford3dDataset_v1.2_Aligned_Version.zip" at [S3DIS](https://docs.google.com/forms/d/e/1FAIpQLScDimvNMCGhy_rmBA2gHfDu3naktRm6A8BPwAWWDv-Uhm6Shw/viewform?c=0&w=1), and move the upcompressed folder to /data/S3DIS 50 | * Preparing the dataset 51 | ``` 52 | python utils/s3dis_dp.py 53 | ``` 54 | * Training 55 | ``` 56 | python s3dis_main.py --gpu 0 --mode train --test_area xxx 57 | ``` 58 | * Evaluation 59 | ``` 60 | python s3dis_main.py --gpu 0 --mode test --test_area xxx --model_path='path_to_model' 61 | ``` 62 | * 6 fold cross validation 63 | ``` 64 | a) train and predict 65 | sh s3dis_6fold.sh 66 | 67 | b) evaluate 68 | python s3dis_6fold_eval.py 69 | ``` 70 | 71 | ### (3) Semantic3D 72 | 73 | ![Semantic3D](./img/semantic3d.png) 74 | 75 | * Setup the dataset 76 | * Download the dataset 77 | ``` 78 | sh utils/download_semantic3d.sh 79 | ``` 80 | * Preparing the dataset 81 | ``` 82 | python utils/semantic3d_dp.py 83 | ``` 84 | * Training 85 | ``` 86 | python semantic3d_main.py --mode train --gpu 0 87 | ``` 88 | * Testing 89 | ``` 90 | python semantic3d_main.py --mode test --gpu 0 --model_path='path_to_model' 91 | ``` 92 | 93 | ### (4) STPLS3D 94 | 95 | * Please follow the official [guidance](https://github.com/meidachen/STPLS3D/tree/main/SCF-Net) of STPLS3D. 96 | 97 | ### (5) Visualization 98 | 99 | ![S3DIS_VIS](./img/s3dis_vis.png) 100 | 101 | ![Semantic3D_VIS](./img/semantic3d_vis.png) 102 | 103 | ``` 104 | python visualization.py --dataset 'xxx' --ply_path "path_to_origin_ply" --label_path "path_to_label" 105 | ``` 106 | 107 | ### Citation 108 | 109 | If you find our work useful in your research, please consider citing: 110 | 111 | ``` 112 | @InProceedings{Fan_2021_CVPR, 113 | author = {Fan, Siqi and Dong, Qiulei and Zhu, Fenghua and Lv, Yisheng and Ye, Peijun and Wang, Fei-Yue}, 114 | title = {SCF-Net: Learning Spatial Contextual Features for Large-Scale Point Cloud Segmentation}, 115 | booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)}, 116 | month = {June}, 117 | year = {2021}, 118 | pages = {14504-14513}} 119 | ``` 120 | 121 | ### Acknowledgment 122 | 123 | Part of our code refers to the work [RandLA-Net](https://github.com/QingyongHu/RandLA-Net) 124 | 125 | 126 | -------------------------------------------------------------------------------- /SCFNet.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import numpy as np 3 | import time 4 | from os import makedirs 5 | from os.path import exists, join 6 | from sklearn.metrics import confusion_matrix 7 | 8 | from helper_dp import DataProcessing as DP 9 | import helper_tf_util 10 | 11 | 12 | def log_out(out_str, f_out): 13 | f_out.write(out_str + '\n') 14 | f_out.flush() 15 | print(out_str) 16 | 17 | 18 | class Network: 19 | def __init__(self, dataset, config): 20 | flat_inputs = dataset.flat_inputs 21 | self.config = config 22 | 23 | # Path of the result folder 24 | if self.config.saving: 25 | if self.config.saving_path is None: 26 | self.saving_path = time.strftime('results/Log_%Y-%m-%d_%H-%M-%S', time.gmtime()) 27 | else: 28 | self.saving_path = self.config.saving_path 29 | makedirs(self.saving_path) if not exists(self.saving_path) else None 30 | 31 | with tf.variable_scope('inputs'): 32 | self.inputs = dict() 33 | num_layers = self.config.num_layers 34 | self.inputs['xyz'] = flat_inputs[:num_layers] 35 | self.inputs['neigh_idx'] = flat_inputs[num_layers: 2 * num_layers] 36 | self.inputs['sub_idx'] = flat_inputs[2 * num_layers:3 * num_layers] 37 | self.inputs['interp_idx'] = flat_inputs[3 * num_layers:4 * num_layers] 38 | self.inputs['features'] = flat_inputs[4 * num_layers] 39 | self.inputs['labels'] = flat_inputs[4 * num_layers + 1] 40 | self.inputs['input_inds'] = flat_inputs[4 * num_layers + 2] 41 | self.inputs['cloud_inds'] = flat_inputs[4 * num_layers + 3] 42 | 43 | self.labels = self.inputs['labels'] 44 | self.is_training = tf.placeholder(tf.bool, shape=()) 45 | self.training_step = 1 46 | self.training_epoch = 0 47 | self.correct_prediction = 0 48 | self.accuracy = 0 49 | self.mIou_list = [0] 50 | self.class_weights = DP.get_class_weights(dataset.name) 51 | 52 | log_time = time.strftime('_%Y-%m-%d_%H-%M-%S', time.gmtime()) 53 | self.Log_file = open('log_train_' + dataset.name + str(dataset.val_split) + log_time + '.txt', 'a') 54 | 55 | with tf.variable_scope('layers'): 56 | self.logits = self.inference(self.inputs, self.is_training) 57 | 58 | # Ignore the invalid point (unlabeled) when calculating the loss 59 | with tf.variable_scope('loss'): 60 | self.logits = tf.reshape(self.logits, [-1, config.num_classes]) 61 | self.labels = tf.reshape(self.labels, [-1]) 62 | 63 | # Boolean mask of points that should be ignored 64 | ignored_bool = tf.zeros_like(self.labels, dtype=tf.bool) 65 | for ign_label in self.config.ignored_label_inds: 66 | ignored_bool = tf.logical_or(ignored_bool, tf.equal(self.labels, ign_label)) 67 | 68 | # Collect logits and labels that are not ignored 69 | valid_idx = tf.squeeze(tf.where(tf.logical_not(ignored_bool))) 70 | valid_logits = tf.gather(self.logits, valid_idx, axis=0) 71 | valid_labels_init = tf.gather(self.labels, valid_idx, axis=0) 72 | 73 | # Reduce label values in the range of logit shape 74 | reducing_list = tf.range(self.config.num_classes, dtype=tf.int32) 75 | inserted_value = tf.zeros((1,), dtype=tf.int32) 76 | for ign_label in self.config.ignored_label_inds: 77 | reducing_list = tf.concat([reducing_list[:ign_label], inserted_value, reducing_list[ign_label:]], 0) 78 | valid_labels = tf.gather(reducing_list, valid_labels_init) 79 | 80 | self.loss = self.get_loss(valid_logits, valid_labels, self.class_weights) 81 | 82 | with tf.variable_scope('optimizer'): 83 | self.learning_rate = tf.Variable(config.learning_rate, trainable=False, name='learning_rate') 84 | self.train_op = tf.train.AdamOptimizer(self.learning_rate).minimize(self.loss) 85 | self.extra_update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) 86 | 87 | with tf.variable_scope('results'): 88 | self.correct_prediction = tf.nn.in_top_k(valid_logits, valid_labels, 1) 89 | self.accuracy = tf.reduce_mean(tf.cast(self.correct_prediction, tf.float32)) 90 | self.prob_logits = tf.nn.softmax(self.logits) 91 | 92 | tf.summary.scalar('learning_rate', self.learning_rate) 93 | tf.summary.scalar('loss', self.loss) 94 | tf.summary.scalar('accuracy', self.accuracy) 95 | 96 | my_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES) 97 | self.saver = tf.train.Saver(my_vars, max_to_keep=100) 98 | c_proto = tf.ConfigProto() 99 | c_proto.gpu_options.allow_growth = True 100 | self.sess = tf.Session(config=c_proto) 101 | self.merged = tf.summary.merge_all() 102 | self.train_writer = tf.summary.FileWriter(config.train_sum_dir, self.sess.graph) 103 | self.sess.run(tf.global_variables_initializer()) 104 | 105 | def inference(self, inputs, is_training): 106 | 107 | d_out = self.config.d_out 108 | feature = inputs['features'] 109 | feature = tf.layers.dense(feature, 8, activation=None, name='fc0') 110 | feature = tf.nn.leaky_relu(tf.layers.batch_normalization(feature, -1, 0.99, 1e-6, training=is_training)) 111 | feature = tf.expand_dims(feature, axis=2) 112 | 113 | # Encoder 114 | f_encoder_list = [] 115 | for i in range(self.config.num_layers): 116 | # SCF 117 | f_encoder_i = self.scf_module(feature, inputs['xyz'][i], inputs['neigh_idx'][i], d_out[i], 'Encoder_layer_' + str(i), is_training) 118 | # RS 119 | f_sampled_i = self.random_sample(f_encoder_i, inputs['sub_idx'][i]) 120 | feature = f_sampled_i 121 | if i == 0: 122 | f_encoder_list.append(f_encoder_i) 123 | f_encoder_list.append(f_sampled_i) 124 | # 125 | 126 | feature = helper_tf_util.conv2d(f_encoder_list[-1], f_encoder_list[-1].get_shape()[3].value, [1, 1], 'decoder_0', [1, 1], 'VALID', True, is_training) 127 | 128 | # Decoder 129 | f_decoder_list = [] 130 | for j in range(self.config.num_layers): 131 | f_interp_i = self.nearest_interpolation(feature, inputs['interp_idx'][-j - 1]) 132 | f_decoder_i = helper_tf_util.conv2d_transpose(tf.concat([f_encoder_list[-j - 2], f_interp_i], axis=3), f_encoder_list[-j - 2].get_shape()[-1].value, 133 | [1, 1], 'Decoder_layer_' + str(j), [1, 1], 'VALID', bn=True, is_training=is_training) 134 | feature = f_decoder_i 135 | f_decoder_list.append(f_decoder_i) 136 | # 137 | 138 | f_layer_fc1 = helper_tf_util.conv2d(f_decoder_list[-1], 64, [1, 1], 'fc1', [1, 1], 'VALID', True, is_training) 139 | f_layer_fc2 = helper_tf_util.conv2d(f_layer_fc1, 32, [1, 1], 'fc2', [1, 1], 'VALID', True, is_training) 140 | f_layer_drop = helper_tf_util.dropout(f_layer_fc2, keep_prob=0.5, is_training=is_training, scope='dp1') 141 | f_layer_fc3 = helper_tf_util.conv2d(f_layer_drop, self.config.num_classes, [1, 1], 'fc', [1, 1], 'VALID', False, is_training, activation_fn=None) 142 | f_out = tf.squeeze(f_layer_fc3, [2]) 143 | return f_out 144 | 145 | def train(self, dataset): 146 | log_out('****EPOCH {}****'.format(self.training_epoch), self.Log_file) 147 | 148 | self.sess.run(dataset.train_init_op) 149 | 150 | while self.training_epoch < self.config.max_epoch: 151 | t_start = time.time() 152 | try: 153 | ops = [self.train_op, 154 | self.extra_update_ops, 155 | self.merged, 156 | self.loss, 157 | self.logits, 158 | self.labels, 159 | self.accuracy] 160 | _, _, summary, l_out, probs, labels, acc = self.sess.run(ops, {self.is_training: True}) 161 | self.train_writer.add_summary(summary, self.training_step) 162 | t_end = time.time() 163 | if self.training_step % 50 == 0: 164 | message = 'Step {:08d} L_out={:5.3f} Acc={:4.2f} ''---{:8.2f} ms/batch' 165 | log_out(message.format(self.training_step, l_out, acc, 1000 * (t_end - t_start)), self.Log_file) 166 | self.training_step += 1 167 | except tf.errors.OutOfRangeError: 168 | 169 | m_iou = self.evaluate(dataset) 170 | if m_iou > np.max(self.mIou_list): 171 | # Save the best model 172 | snapshot_directory = join(self.saving_path, 'snapshots') 173 | makedirs(snapshot_directory) if not exists(snapshot_directory) else None 174 | self.saver.save(self.sess, snapshot_directory + '/snap', global_step=self.training_step) 175 | self.mIou_list.append(m_iou) 176 | log_out('Best m_IoU is: {:5.3f}'.format(max(self.mIou_list)), self.Log_file) 177 | 178 | self.training_epoch += 1 179 | self.sess.run(dataset.train_init_op) 180 | # Update learning rate 181 | op = self.learning_rate.assign(tf.multiply(self.learning_rate, self.config.lr_decays[self.training_epoch])) 182 | self.sess.run(op) 183 | log_out('****EPOCH {}****'.format(self.training_epoch), self.Log_file) 184 | 185 | except tf.errors.InvalidArgumentError as e: 186 | print('Caught a NaN error :') 187 | print(e.error_code) 188 | print(e.message) 189 | print(e.op) 190 | print(e.op.name) 191 | print([t.name for t in e.op.inputs]) 192 | print([t.name for t in e.op.outputs]) 193 | 194 | a = 1 / 0 195 | 196 | print('finished') 197 | self.sess.close() 198 | 199 | def evaluate(self, dataset): 200 | 201 | # Initialise iterator with validation data 202 | self.sess.run(dataset.val_init_op) 203 | 204 | gt_classes = [0 for _ in range(self.config.num_classes)] 205 | positive_classes = [0 for _ in range(self.config.num_classes)] 206 | true_positive_classes = [0 for _ in range(self.config.num_classes)] 207 | val_total_correct = 0 208 | val_total_seen = 0 209 | 210 | for step_id in range(self.config.val_steps): 211 | if step_id % 50 == 0: 212 | print(str(step_id) + ' / ' + str(self.config.val_steps)) 213 | try: 214 | ops = (self.prob_logits, self.labels, self.accuracy) 215 | stacked_prob, labels, acc = self.sess.run(ops, {self.is_training: False}) 216 | pred = np.argmax(stacked_prob, 1) 217 | if not self.config.ignored_label_inds: 218 | pred_valid = pred 219 | labels_valid = labels 220 | else: 221 | invalid_idx = np.where(labels == self.config.ignored_label_inds)[0] 222 | labels_valid = np.delete(labels, invalid_idx) 223 | labels_valid = labels_valid - 1 224 | pred_valid = np.delete(pred, invalid_idx) 225 | 226 | correct = np.sum(pred_valid == labels_valid) 227 | val_total_correct += correct 228 | val_total_seen += len(labels_valid) 229 | 230 | conf_matrix = confusion_matrix(labels_valid, pred_valid, np.arange(0, self.config.num_classes, 1)) 231 | gt_classes += np.sum(conf_matrix, axis=1) 232 | positive_classes += np.sum(conf_matrix, axis=0) 233 | true_positive_classes += np.diagonal(conf_matrix) 234 | 235 | except tf.errors.OutOfRangeError: 236 | break 237 | 238 | iou_list = [] 239 | for n in range(0, self.config.num_classes, 1): 240 | iou = true_positive_classes[n] / float(gt_classes[n] + positive_classes[n] - true_positive_classes[n]) 241 | iou_list.append(iou) 242 | mean_iou = sum(iou_list) / float(self.config.num_classes) 243 | 244 | log_out('eval accuracy: {}'.format(val_total_correct / float(val_total_seen)), self.Log_file) 245 | log_out('mean IOU:{}'.format(mean_iou), self.Log_file) 246 | 247 | mean_iou = 100 * mean_iou 248 | log_out('Mean IoU = {:.1f}%'.format(mean_iou), self.Log_file) 249 | s = '{:5.2f} | '.format(mean_iou) 250 | for IoU in iou_list: 251 | s += '{:5.2f} '.format(100 * IoU) 252 | log_out('-' * len(s), self.Log_file) 253 | log_out(s, self.Log_file) 254 | log_out('-' * len(s) + '\n', self.Log_file) 255 | return mean_iou 256 | 257 | def get_loss(self, logits, labels, pre_cal_weights): 258 | # calculate the weighted cross entropy according to the inverse frequency 259 | class_weights = tf.convert_to_tensor(pre_cal_weights, dtype=tf.float32) 260 | one_hot_labels = tf.one_hot(labels, depth=self.config.num_classes) 261 | weights = tf.reduce_sum(class_weights * one_hot_labels, axis=1) 262 | unweighted_losses = tf.nn.softmax_cross_entropy_with_logits(logits=logits, labels=one_hot_labels) 263 | weighted_losses = unweighted_losses * weights 264 | output_loss = tf.reduce_mean(weighted_losses) 265 | return output_loss 266 | 267 | def scf_module(self, feature, xyz, neigh_idx, d_out, name, is_training): 268 | """ 269 | SCF 270 | """ 271 | # Local Contextual Features 272 | # MLP 1 273 | f_pc = helper_tf_util.conv2d(feature, d_out // 2, [1, 1], name + 'mlp1', [1, 1], 'VALID', True, is_training) 274 | # Local Context Learning (LPR + DDAP) 275 | f_lc, lg_volume_ratio = self.local_context_learning(xyz, f_pc, neigh_idx, d_out, name + 'LFA', is_training) 276 | # MLP 2 277 | f_lc = helper_tf_util.conv2d(f_lc, d_out * 2, [1, 1], name + 'mlp2', [1, 1], 'VALID', True, is_training, activation_fn=None) 278 | # MLP Shotcut 279 | shortcut = helper_tf_util.conv2d(feature, d_out * 2, [1, 1], name + 'shortcut', [1, 1], 'VALID', activation_fn=None, bn=True, is_training=is_training) 280 | 281 | # Global Contextual Features 282 | f_gc = tf.expand_dims(tf.concat([xyz, lg_volume_ratio], axis=-1), -2) 283 | f_gc = helper_tf_util.conv2d(f_gc, d_out * 2, [1, 1], name + 'lg', [1, 1], 'VALID', activation_fn=None, bn=True, is_training=is_training) 284 | 285 | return tf.nn.leaky_relu(tf.concat([f_lc + shortcut, f_gc], axis=-1)) 286 | 287 | def local_context_learning(self, xyz, feature, neigh_idx, d_out, name, is_training): 288 | """ 289 | (LPR + DDAP) * 2 290 | """ 291 | d_in = feature.get_shape()[-1].value 292 | # LPR 293 | local_rep, g_dis, lg_volume_ratio = self.local_polar_representation(xyz, neigh_idx) 294 | 295 | # 1 296 | local_rep = helper_tf_util.conv2d(local_rep, d_in, [1, 1], name + 'mlp1', [1, 1], 'VALID', True, is_training) 297 | f_neighbours = self.gather_neighbour(tf.squeeze(feature, axis=2), neigh_idx) 298 | f_concat = tf.concat([f_neighbours, local_rep], axis=-1) 299 | f_dis = self.cal_feature_dis(tf.squeeze(feature, axis=2), f_neighbours) 300 | f_lc = self.dualdis_att_pool(f_concat, f_dis, g_dis, d_out // 2, name + 'dis_att_pooling_1', is_training) 301 | 302 | # 2 303 | local_rep = helper_tf_util.conv2d(local_rep, d_out // 2, [1, 1], name + 'mlp2', [1, 1], 'VALID', True, is_training) 304 | f_neighbours = self.gather_neighbour(tf.squeeze(f_lc, axis=2), neigh_idx) 305 | f_concat = tf.concat([f_neighbours, local_rep], axis=-1) 306 | f_dis = self.cal_feature_dis(tf.squeeze(f_lc, axis=2), f_neighbours) 307 | f_lc = self.dualdis_att_pool(f_concat, f_dis, g_dis, d_out, name + 'dis_att_pooling_2', is_training) 308 | 309 | return f_lc, lg_volume_ratio 310 | 311 | def local_polar_representation(self, xyz, neigh_idx): 312 | """ 313 | LPR & Calculate volume ratio for GCF 314 | """ 315 | # Get neighbor xyz 316 | neighbor_xyz = self.gather_neighbour(xyz, neigh_idx) 317 | 318 | # Relative position transforming 319 | relative_info, relative_alpha, relative_beta, geometric_dis, local_volume = self.relative_pos_transforming(xyz, neigh_idx, neighbor_xyz) 320 | 321 | # Local direction calculation (angle) 322 | neighbor_mean = tf.reduce_mean(neighbor_xyz, axis=-2) 323 | direction = xyz - neighbor_mean 324 | direction_tile = tf.tile(tf.expand_dims(direction, axis=2), [1, 1, tf.shape(neigh_idx)[-1], 1]) 325 | 326 | direction_alpha = tf.expand_dims(tf.atan2(direction_tile[:,:,:,1], direction_tile[:,:,:,0]), axis=-1) 327 | direction_xydis = tf.sqrt(tf.reduce_sum(tf.square(direction_tile[:,:,:,:2]), axis=-1)) 328 | direction_beta = tf.expand_dims(tf.atan2(direction_tile[:,:,:,2], direction_xydis), axis=-1) 329 | 330 | # Polar angle updating 331 | angle_alpha = relative_alpha - direction_alpha 332 | angle_beta = relative_beta - direction_beta 333 | angle_updated = tf.concat([angle_alpha, angle_beta], axis=-1) 334 | 335 | # Generate local spatial representation 336 | local_rep = tf.concat([angle_updated, relative_info], axis=-1) 337 | 338 | # Calculate volume ratio for GCF 339 | global_dis = tf.sqrt(tf.reduce_sum(tf.square(xyz), axis=-1, keepdims=True)) 340 | global_volume = tf.pow(tf.reduce_max(global_dis, axis=-1), 3) 341 | lg_volume_ratio = tf.expand_dims(local_volume / global_volume, -1) 342 | 343 | return local_rep, geometric_dis, lg_volume_ratio 344 | 345 | 346 | def relative_pos_transforming(self, xyz, neigh_idx, neighbor_xyz): 347 | 348 | xyz_tile = tf.tile(tf.expand_dims(xyz, axis=2), [1, 1, tf.shape(neigh_idx)[-1], 1]) 349 | 350 | relative_xyz = xyz_tile - neighbor_xyz 351 | 352 | relative_alpha = tf.expand_dims(tf.atan2(relative_xyz[:,:,:,1], relative_xyz[:,:,:,0]), axis=-1) 353 | relative_xydis = tf.sqrt(tf.reduce_sum(tf.square(relative_xyz[:,:,:,:2]), axis=-1)) 354 | relative_beta = tf.expand_dims(tf.atan2(relative_xyz[:,:,:,2], relative_xydis), axis=-1) 355 | relative_dis = tf.sqrt(tf.reduce_sum(tf.square(relative_xyz), axis=-1, keepdims=True)) 356 | 357 | relative_info = tf.concat([relative_dis, xyz_tile, neighbor_xyz], axis=-1) 358 | 359 | # negative exp of geometric distance 360 | exp_dis = tf.exp(-relative_dis) 361 | 362 | # volume of local region 363 | local_volume = tf.pow(tf.reduce_max(tf.reduce_max(relative_dis, -1), -1), 3) 364 | 365 | return relative_info, relative_alpha, relative_beta, exp_dis, local_volume 366 | 367 | def cal_feature_dis(self, feature, f_neighbours): 368 | """ 369 | Calculate the feature distance 370 | """ 371 | feature_tile = tf.tile(tf.expand_dims(feature, axis=2), [1, 1, tf.shape(f_neighbours)[-2], 1]) 372 | feature_dist = feature_tile - f_neighbours 373 | feature_dist = tf.expand_dims(tf.reduce_mean(tf.abs(feature_dist), axis=-1), axis=-1) 374 | feature_dist = tf.exp(-feature_dist) 375 | 376 | return feature_dist 377 | 378 | @staticmethod 379 | def random_sample(feature, pool_idx): 380 | """ 381 | :param feature: [B, N, d] input features matrix 382 | :param pool_idx: [B, N', max_num] N' < N, N' is the selected position after pooling 383 | :return: pool_features = [B, N', d] pooled features matrix 384 | """ 385 | feature = tf.squeeze(feature, axis=2) 386 | num_neigh = tf.shape(pool_idx)[-1] 387 | d = feature.get_shape()[-1] 388 | batch_size = tf.shape(pool_idx)[0] 389 | pool_idx = tf.reshape(pool_idx, [batch_size, -1]) 390 | pool_features = tf.batch_gather(feature, pool_idx) 391 | pool_features = tf.reshape(pool_features, [batch_size, -1, num_neigh, d]) 392 | pool_features = tf.reduce_max(pool_features, axis=2, keepdims=True) 393 | return pool_features 394 | 395 | @staticmethod 396 | def nearest_interpolation(feature, interp_idx): 397 | """ 398 | :param feature: [B, N, d] input features matrix 399 | :param interp_idx: [B, up_num_points, 1] nearest neighbour index 400 | :return: [B, up_num_points, d] interpolated features matrix 401 | """ 402 | feature = tf.squeeze(feature, axis=2) 403 | batch_size = tf.shape(interp_idx)[0] 404 | up_num_points = tf.shape(interp_idx)[1] 405 | interp_idx = tf.reshape(interp_idx, [batch_size, up_num_points]) 406 | interpolated_features = tf.batch_gather(feature, interp_idx) 407 | interpolated_features = tf.expand_dims(interpolated_features, axis=2) 408 | return interpolated_features 409 | 410 | @staticmethod 411 | def gather_neighbour(pc, neighbor_idx): 412 | """ 413 | gather the coordinates or features of neighboring points 414 | """ 415 | batch_size = tf.shape(pc)[0] 416 | num_points = tf.shape(pc)[1] 417 | d = pc.get_shape()[2].value 418 | 419 | index_input = tf.reshape(neighbor_idx, shape=[batch_size, -1]) 420 | features = tf.batch_gather(pc, index_input) 421 | features = tf.reshape(features, [batch_size, num_points, tf.shape(neighbor_idx)[-1], d]) 422 | return features 423 | 424 | @staticmethod 425 | def dualdis_att_pool(feature_set, f_dis, g_dis, d_out, name, is_training): 426 | """ 427 | DDAP 428 | """ 429 | batch_size = tf.shape(feature_set)[0] 430 | num_points = tf.shape(feature_set)[1] 431 | num_neigh = tf.shape(feature_set)[2] 432 | d = feature_set.get_shape()[3].value 433 | d_dis = g_dis.get_shape()[3].value 434 | 435 | f_reshaped = tf.reshape(feature_set, shape=[-1, num_neigh, d]) 436 | f_dis_reshaped = tf.reshape(f_dis, shape=[-1, num_neigh, d_dis]) * 0.1 437 | g_dis_reshaped = tf.reshape(g_dis, shape=[-1, num_neigh, d_dis]) 438 | concat = tf.concat([g_dis_reshaped, f_dis_reshaped, f_reshaped], axis=-1) 439 | 440 | # weight learning 441 | att_activation = tf.layers.dense(concat, d, activation=None, use_bias=False, name=name + 'dis_self_fc') 442 | att_scores = tf.nn.softmax(att_activation, axis=1) 443 | # dot product 444 | f_lc = f_reshaped * att_scores 445 | # sum 446 | f_lc = tf.reduce_sum(f_lc, axis=1) 447 | f_lc = tf.reshape(f_lc, [batch_size, num_points, 1, d]) 448 | # shared MLP 449 | f_lc = helper_tf_util.conv2d(f_lc, d_out, [1, 1], name + 'mlp', [1, 1], 'VALID', True, is_training) 450 | 451 | return f_lc 452 | 453 | -------------------------------------------------------------------------------- /compile_op.sh: -------------------------------------------------------------------------------- 1 | cd utils/nearest_neighbors 2 | python setup.py install --home="." 3 | cd ../../ 4 | 5 | cd utils/cpp_wrappers 6 | sh compile_wrappers.sh 7 | cd ../../../ -------------------------------------------------------------------------------- /helper_dp.py: -------------------------------------------------------------------------------- 1 | import pandas as pd 2 | import numpy as np 3 | import colorsys, random, os, sys 4 | from os.path import join 5 | 6 | os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' 7 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 8 | sys.path.append(BASE_DIR) 9 | sys.path.append(os.path.join(BASE_DIR, 'utils')) 10 | import cpp_wrappers.cpp_subsampling.grid_subsampling as cpp_subsampling 11 | import nearest_neighbors.lib.python.nearest_neighbors as nearest_neighbors 12 | 13 | 14 | class DataProcessing: 15 | @staticmethod 16 | def load_pc_semantic3d(filename): 17 | pc_pd = pd.read_csv(filename, header=None, delim_whitespace=True, dtype=np.float16) 18 | pc = pc_pd.values 19 | return pc 20 | 21 | @staticmethod 22 | def load_label_semantic3d(filename): 23 | label_pd = pd.read_csv(filename, header=None, delim_whitespace=True, dtype=np.uint8) 24 | cloud_labels = label_pd.values 25 | return cloud_labels 26 | 27 | @staticmethod 28 | def knn_search(support_pts, query_pts, k): 29 | """ 30 | :param support_pts: points you have, B*N1*3 31 | :param query_pts: points you want to know the neighbour index, B*N2*3 32 | :param k: Number of neighbours in knn search 33 | :return: neighbor_idx: neighboring points indexes, B*N2*k 34 | """ 35 | 36 | neighbor_idx = nearest_neighbors.knn_batch(support_pts, query_pts, k, omp=True) 37 | return neighbor_idx.astype(np.int32) 38 | 39 | @staticmethod 40 | def data_aug(xyz, color, labels, idx, num_out): 41 | num_in = len(xyz) 42 | dup = np.random.choice(num_in, num_out - num_in) 43 | xyz_dup = xyz[dup, ...] 44 | xyz_aug = np.concatenate([xyz, xyz_dup], 0) 45 | color_dup = color[dup, ...] 46 | color_aug = np.concatenate([color, color_dup], 0) 47 | idx_dup = list(range(num_in)) + list(dup) 48 | idx_aug = idx[idx_dup] 49 | label_aug = labels[idx_dup] 50 | return xyz_aug, color_aug, idx_aug, label_aug 51 | 52 | @staticmethod 53 | def shuffle_idx(x): 54 | # random shuffle the index 55 | idx = np.arange(len(x)) 56 | np.random.shuffle(idx) 57 | return x[idx] 58 | 59 | 60 | @staticmethod 61 | def grid_sub_sampling(points, features=None, labels=None, grid_size=0.1, verbose=0): 62 | """ 63 | CPP wrapper for a grid sub_sampling (method = barycenter for points and features 64 | :param points: (N, 3) matrix of input points 65 | :param features: optional (N, d) matrix of features (floating number) 66 | :param labels: optional (N,) matrix of integer labels 67 | :param grid_size: parameter defining the size of grid voxels 68 | :param verbose: 1 to display 69 | :return: sub_sampled points, with features and/or labels depending of the input 70 | """ 71 | 72 | if (features is None) and (labels is None): 73 | return cpp_subsampling.compute(points, sampleDl=grid_size, verbose=verbose) 74 | elif labels is None: 75 | return cpp_subsampling.compute(points, features=features, sampleDl=grid_size, verbose=verbose) 76 | elif features is None: 77 | return cpp_subsampling.compute(points, classes=labels, sampleDl=grid_size, verbose=verbose) 78 | else: 79 | return cpp_subsampling.compute(points, features=features, classes=labels, sampleDl=grid_size, 80 | verbose=verbose) 81 | 82 | @staticmethod 83 | def IoU_from_confusions(confusions): 84 | """ 85 | Computes IoU from confusion matrices. 86 | :param confusions: ([..., n_c, n_c] np.int32). Can be any dimension, the confusion matrices should be described by 87 | the last axes. n_c = number of classes 88 | :return: ([..., n_c] np.float32) IoU score 89 | """ 90 | 91 | # Compute TP, FP, FN. This assume that the second to last axis counts the truths (like the first axis of a 92 | # confusion matrix), and that the last axis counts the predictions (like the second axis of a confusion matrix) 93 | TP = np.diagonal(confusions, axis1=-2, axis2=-1) 94 | TP_plus_FN = np.sum(confusions, axis=-1) 95 | TP_plus_FP = np.sum(confusions, axis=-2) 96 | 97 | # Compute IoU 98 | IoU = TP / (TP_plus_FP + TP_plus_FN - TP + 1e-6) 99 | 100 | # Compute mIoU with only the actual classes 101 | mask = TP_plus_FN < 1e-3 102 | counts = np.sum(1 - mask, axis=-1, keepdims=True) 103 | mIoU = np.sum(IoU, axis=-1, keepdims=True) / (counts + 1e-6) 104 | 105 | # If class is absent, place mIoU in place of 0 IoU to get the actual mean later 106 | IoU += mask * mIoU 107 | return IoU 108 | 109 | @staticmethod 110 | def get_class_weights(dataset_name): 111 | # pre-calculate the number of points in each category 112 | num_per_class = [] 113 | if dataset_name is 'S3DIS': 114 | num_per_class = np.array([3370714, 2856755, 4919229, 318158, 375640, 478001, 974733, 115 | 650464, 791496, 88727, 1284130, 229758, 2272837], dtype=np.int32) 116 | elif dataset_name is 'Semantic3D': 117 | num_per_class = np.array([5181602, 5012952, 6830086, 1311528, 10476365, 946982, 334860, 269353], 118 | dtype=np.int32) 119 | weight = num_per_class / float(sum(num_per_class)) 120 | ce_label_weight = 1 / (weight + 0.02) 121 | return np.expand_dims(ce_label_weight, axis=0) -------------------------------------------------------------------------------- /helper_ply.py: -------------------------------------------------------------------------------- 1 | # 2 | # 3 | # 0===============================0 4 | # | PLY files reader/writer | 5 | # 0===============================0 6 | # 7 | # 8 | # ---------------------------------------------------------------------------------------------------------------------- 9 | # 10 | # function to read/write .ply files 11 | # 12 | # ---------------------------------------------------------------------------------------------------------------------- 13 | # 14 | # Hugues THOMAS - 10/02/2017 15 | # 16 | 17 | 18 | # ---------------------------------------------------------------------------------------------------------------------- 19 | # 20 | # Imports and global variables 21 | # \**********************************/ 22 | # 23 | 24 | 25 | # Basic libs 26 | import numpy as np 27 | import sys 28 | 29 | 30 | # Define PLY types 31 | ply_dtypes = dict([ 32 | (b'int8', 'i1'), 33 | (b'char', 'i1'), 34 | (b'uint8', 'u1'), 35 | (b'uchar', 'u1'), 36 | (b'int16', 'i2'), 37 | (b'short', 'i2'), 38 | (b'uint16', 'u2'), 39 | (b'ushort', 'u2'), 40 | (b'int32', 'i4'), 41 | (b'int', 'i4'), 42 | (b'uint32', 'u4'), 43 | (b'uint', 'u4'), 44 | (b'float32', 'f4'), 45 | (b'float', 'f4'), 46 | (b'float64', 'f8'), 47 | (b'double', 'f8') 48 | ]) 49 | 50 | # Numpy reader format 51 | valid_formats = {'ascii': '', 'binary_big_endian': '>', 52 | 'binary_little_endian': '<'} 53 | 54 | 55 | # ---------------------------------------------------------------------------------------------------------------------- 56 | # 57 | # Functions 58 | # \***************/ 59 | # 60 | 61 | 62 | def parse_header(plyfile, ext): 63 | # Variables 64 | line = [] 65 | properties = [] 66 | num_points = None 67 | 68 | while b'end_header' not in line and line != b'': 69 | line = plyfile.readline() 70 | 71 | if b'element' in line: 72 | line = line.split() 73 | num_points = int(line[2]) 74 | 75 | elif b'property' in line: 76 | line = line.split() 77 | properties.append((line[2].decode(), ext + ply_dtypes[line[1]])) 78 | 79 | return num_points, properties 80 | 81 | 82 | def parse_mesh_header(plyfile, ext): 83 | # Variables 84 | line = [] 85 | vertex_properties = [] 86 | num_points = None 87 | num_faces = None 88 | current_element = None 89 | 90 | 91 | while b'end_header' not in line and line != b'': 92 | line = plyfile.readline() 93 | 94 | # Find point element 95 | if b'element vertex' in line: 96 | current_element = 'vertex' 97 | line = line.split() 98 | num_points = int(line[2]) 99 | 100 | elif b'element face' in line: 101 | current_element = 'face' 102 | line = line.split() 103 | num_faces = int(line[2]) 104 | 105 | elif b'property' in line: 106 | if current_element == 'vertex': 107 | line = line.split() 108 | vertex_properties.append((line[2].decode(), ext + ply_dtypes[line[1]])) 109 | elif current_element == 'vertex': 110 | if not line.startswith('property list uchar int'): 111 | raise ValueError('Unsupported faces property : ' + line) 112 | 113 | return num_points, num_faces, vertex_properties 114 | 115 | 116 | def read_ply(filename, triangular_mesh=False): 117 | """ 118 | Read ".ply" files 119 | 120 | Parameters 121 | ---------- 122 | filename : string 123 | the name of the file to read. 124 | 125 | Returns 126 | ------- 127 | result : array 128 | data stored in the file 129 | 130 | Examples 131 | -------- 132 | Store data in file 133 | 134 | >>> points = np.random.rand(5, 3) 135 | >>> values = np.random.randint(2, size=10) 136 | >>> write_ply('example.ply', [points, values], ['x', 'y', 'z', 'values']) 137 | 138 | Read the file 139 | 140 | >>> data = read_ply('example.ply') 141 | >>> values = data['values'] 142 | array([0, 0, 1, 1, 0]) 143 | 144 | >>> points = np.vstack((data['x'], data['y'], data['z'])).T 145 | array([[ 0.466 0.595 0.324] 146 | [ 0.538 0.407 0.654] 147 | [ 0.850 0.018 0.988] 148 | [ 0.395 0.394 0.363] 149 | [ 0.873 0.996 0.092]]) 150 | 151 | """ 152 | 153 | with open(filename, 'rb') as plyfile: 154 | 155 | 156 | # Check if the file start with ply 157 | if b'ply' not in plyfile.readline(): 158 | raise ValueError('The file does not start whith the word ply') 159 | 160 | # get binary_little/big or ascii 161 | fmt = plyfile.readline().split()[1].decode() 162 | if fmt == "ascii": 163 | raise ValueError('The file is not binary') 164 | 165 | # get extension for building the numpy dtypes 166 | ext = valid_formats[fmt] 167 | 168 | # PointCloud reader vs mesh reader 169 | if triangular_mesh: 170 | 171 | # Parse header 172 | num_points, num_faces, properties = parse_mesh_header(plyfile, ext) 173 | 174 | # Get point data 175 | vertex_data = np.fromfile(plyfile, dtype=properties, count=num_points) 176 | 177 | # Get face data 178 | face_properties = [('k', ext + 'u1'), 179 | ('v1', ext + 'i4'), 180 | ('v2', ext + 'i4'), 181 | ('v3', ext + 'i4')] 182 | faces_data = np.fromfile(plyfile, dtype=face_properties, count=num_faces) 183 | 184 | # Return vertex data and concatenated faces 185 | faces = np.vstack((faces_data['v1'], faces_data['v2'], faces_data['v3'])).T 186 | data = [vertex_data, faces] 187 | 188 | else: 189 | 190 | # Parse header 191 | num_points, properties = parse_header(plyfile, ext) 192 | 193 | # Get data 194 | data = np.fromfile(plyfile, dtype=properties, count=num_points) 195 | 196 | return data 197 | 198 | 199 | def header_properties(field_list, field_names): 200 | 201 | # List of lines to write 202 | lines = [] 203 | 204 | # First line describing element vertex 205 | lines.append('element vertex %d' % field_list[0].shape[0]) 206 | 207 | # Properties lines 208 | i = 0 209 | for fields in field_list: 210 | for field in fields.T: 211 | lines.append('property %s %s' % (field.dtype.name, field_names[i])) 212 | i += 1 213 | 214 | return lines 215 | 216 | 217 | def write_ply(filename, field_list, field_names, triangular_faces=None): 218 | """ 219 | Write ".ply" files 220 | 221 | Parameters 222 | ---------- 223 | filename : string 224 | the name of the file to which the data is saved. A '.ply' extension will be appended to the 225 | file name if it does no already have one. 226 | 227 | field_list : list, tuple, numpy array 228 | the fields to be saved in the ply file. Either a numpy array, a list of numpy arrays or a 229 | tuple of numpy arrays. Each 1D numpy array and each column of 2D numpy arrays are considered 230 | as one field. 231 | 232 | field_names : list 233 | the name of each fields as a list of strings. Has to be the same length as the number of 234 | fields. 235 | 236 | Examples 237 | -------- 238 | >>> points = np.random.rand(10, 3) 239 | >>> write_ply('example1.ply', points, ['x', 'y', 'z']) 240 | 241 | >>> values = np.random.randint(2, size=10) 242 | >>> write_ply('example2.ply', [points, values], ['x', 'y', 'z', 'values']) 243 | 244 | >>> colors = np.random.randint(255, size=(10,3), dtype=np.uint8) 245 | >>> field_names = ['x', 'y', 'z', 'red', 'green', 'blue', values'] 246 | >>> write_ply('example3.ply', [points, colors, values], field_names) 247 | 248 | """ 249 | 250 | # Format list input to the right form 251 | field_list = list(field_list) if (type(field_list) == list or type(field_list) == tuple) else list((field_list,)) 252 | for i, field in enumerate(field_list): 253 | if field.ndim < 2: 254 | field_list[i] = field.reshape(-1, 1) 255 | if field.ndim > 2: 256 | print('fields have more than 2 dimensions') 257 | return False 258 | 259 | # check all fields have the same number of data 260 | n_points = [field.shape[0] for field in field_list] 261 | if not np.all(np.equal(n_points, n_points[0])): 262 | print('wrong field dimensions') 263 | return False 264 | 265 | # Check if field_names and field_list have same nb of column 266 | n_fields = np.sum([field.shape[1] for field in field_list]) 267 | if (n_fields != len(field_names)): 268 | print('wrong number of field names') 269 | return False 270 | 271 | # Add extension if not there 272 | if not filename.endswith('.ply'): 273 | filename += '.ply' 274 | 275 | # open in text mode to write the header 276 | with open(filename, 'w') as plyfile: 277 | 278 | # First magical word 279 | header = ['ply'] 280 | 281 | # Encoding format 282 | header.append('format binary_' + sys.byteorder + '_endian 1.0') 283 | 284 | # Points properties description 285 | header.extend(header_properties(field_list, field_names)) 286 | 287 | # Add faces if needded 288 | if triangular_faces is not None: 289 | header.append('element face {:d}'.format(triangular_faces.shape[0])) 290 | header.append('property list uchar int vertex_indices') 291 | 292 | # End of header 293 | header.append('end_header') 294 | 295 | # Write all lines 296 | for line in header: 297 | plyfile.write("%s\n" % line) 298 | 299 | # open in binary/append to use tofile 300 | with open(filename, 'ab') as plyfile: 301 | 302 | # Create a structured array 303 | i = 0 304 | type_list = [] 305 | for fields in field_list: 306 | for field in fields.T: 307 | type_list += [(field_names[i], field.dtype.str)] 308 | i += 1 309 | data = np.empty(field_list[0].shape[0], dtype=type_list) 310 | i = 0 311 | for fields in field_list: 312 | for field in fields.T: 313 | data[field_names[i]] = field 314 | i += 1 315 | 316 | data.tofile(plyfile) 317 | 318 | if triangular_faces is not None: 319 | triangular_faces = triangular_faces.astype(np.int32) 320 | type_list = [('k', 'uint8')] + [(str(ind), 'int32') for ind in range(3)] 321 | data = np.empty(triangular_faces.shape[0], dtype=type_list) 322 | data['k'] = np.full((triangular_faces.shape[0],), 3, dtype=np.uint8) 323 | data['0'] = triangular_faces[:, 0] 324 | data['1'] = triangular_faces[:, 1] 325 | data['2'] = triangular_faces[:, 2] 326 | data.tofile(plyfile) 327 | 328 | return True 329 | 330 | 331 | def describe_element(name, df): 332 | """ Takes the columns of the dataframe and builds a ply-like description 333 | 334 | Parameters 335 | ---------- 336 | name: str 337 | df: pandas DataFrame 338 | 339 | Returns 340 | ------- 341 | element: list[str] 342 | """ 343 | property_formats = {'f': 'float', 'u': 'uchar', 'i': 'int'} 344 | element = ['element ' + name + ' ' + str(len(df))] 345 | 346 | if name == 'face': 347 | element.append("property list uchar int points_indices") 348 | 349 | else: 350 | for i in range(len(df.columns)): 351 | # get first letter of dtype to infer format 352 | f = property_formats[str(df.dtypes[i])[0]] 353 | element.append('property ' + f + ' ' + df.columns.values[i]) 354 | 355 | return element 356 | 357 | -------------------------------------------------------------------------------- /helper_tf_util.py: -------------------------------------------------------------------------------- 1 | """ Wrapper functions for TensorFlow layers. 2 | 3 | Author: Charles R. Qi 4 | Date: November 2016 5 | """ 6 | 7 | import numpy as np 8 | import tensorflow as tf 9 | 10 | 11 | def _variable_on_cpu(name, shape, initializer, use_fp16=False): 12 | """Helper to create a Variable stored on CPU memory. 13 | Args: 14 | name: name of the variable 15 | shape: list of ints 16 | initializer: initializer for Variable 17 | Returns: 18 | Variable Tensor 19 | """ 20 | with tf.device('/cpu:0'): 21 | dtype = tf.float16 if use_fp16 else tf.float32 22 | var = tf.get_variable(name, shape, initializer=initializer, dtype=dtype) 23 | return var 24 | 25 | 26 | def _variable_with_weight_decay(name, shape, stddev, wd, use_xavier=True): 27 | """Helper to create an initialized Variable with weight decay. 28 | 29 | Note that the Variable is initialized with a truncated normal distribution. 30 | A weight decay is added only if one is specified. 31 | 32 | Args: 33 | name: name of the variable 34 | shape: list of ints 35 | stddev: standard deviation of a truncated Gaussian 36 | wd: add L2Loss weight decay multiplied by this float. If None, weight 37 | decay is not added for this Variable. 38 | use_xavier: bool, whether to use xavier initializer 39 | 40 | Returns: 41 | Variable Tensor 42 | """ 43 | if use_xavier: 44 | initializer = tf.contrib.layers.xavier_initializer() 45 | var = _variable_on_cpu(name, shape, initializer) 46 | else: 47 | # initializer = tf.truncated_normal_initializer(stddev=stddev) 48 | with tf.device('/cpu:0'): 49 | var = tf.truncated_normal(shape, stddev=np.sqrt(2 / shape[-1])) 50 | var = tf.round(var * tf.constant(1000, dtype=tf.float32)) / tf.constant(1000, dtype=tf.float32) 51 | var = tf.Variable(var, name='weights') 52 | if wd is not None: 53 | weight_decay = tf.multiply(tf.nn.l2_loss(var), wd, name='weight_loss') 54 | tf.add_to_collection('losses', weight_decay) 55 | return var 56 | 57 | 58 | def conv1d(inputs, 59 | num_output_channels, 60 | kernel_size, 61 | scope, 62 | stride=1, 63 | padding='SAME', 64 | use_xavier=True, 65 | stddev=1e-3, 66 | weight_decay=0.0, 67 | activation_fn=tf.nn.relu, 68 | bn=False, 69 | bn_decay=None, 70 | is_training=None): 71 | """ 1D convolution with non-linear operation. 72 | 73 | Args: 74 | inputs: 3-D tensor variable BxLxC 75 | num_output_channels: int 76 | kernel_size: int 77 | scope: string 78 | stride: int 79 | padding: 'SAME' or 'VALID' 80 | use_xavier: bool, use xavier_initializer if true 81 | stddev: float, stddev for truncated_normal init 82 | weight_decay: float 83 | activation_fn: function 84 | bn: bool, whether to use batch norm 85 | bn_decay: float or float tensor variable in [0,1] 86 | is_training: bool Tensor variable 87 | 88 | Returns: 89 | Variable tensor 90 | """ 91 | with tf.variable_scope(scope) as sc: 92 | num_in_channels = inputs.get_shape()[-1].value 93 | kernel_shape = [kernel_size, 94 | num_in_channels, num_output_channels] 95 | kernel = _variable_with_weight_decay('weights', 96 | shape=kernel_shape, 97 | use_xavier=use_xavier, 98 | stddev=stddev, 99 | wd=weight_decay) 100 | outputs = tf.nn.conv1d(inputs, kernel, 101 | stride=stride, 102 | padding=padding) 103 | biases = _variable_on_cpu('biases', [num_output_channels], 104 | tf.constant_initializer(0.0)) 105 | outputs = tf.nn.bias_add(outputs, biases) 106 | 107 | if bn: 108 | outputs = batch_norm_for_conv1d(outputs, is_training, 109 | bn_decay=bn_decay, scope='bn') 110 | if activation_fn is not None: 111 | outputs = activation_fn(outputs) 112 | return outputs 113 | 114 | 115 | def conv2d(inputs, 116 | num_output_channels, 117 | kernel_size, 118 | scope, 119 | stride=[1, 1], 120 | padding='SAME', 121 | bn=False, 122 | is_training=None, 123 | use_xavier=False, 124 | stddev=1e-3, 125 | weight_decay=0.0, 126 | activation_fn=tf.nn.relu, 127 | bn_decay=None): 128 | """ 2D convolution with non-linear operation. 129 | 130 | Args: 131 | inputs: 4-D tensor variable BxHxWxC 132 | num_output_channels: int 133 | kernel_size: a list of 2 ints 134 | scope: string 135 | stride: a list of 2 ints 136 | padding: 'SAME' or 'VALID' 137 | use_xavier: bool, use xavier_initializer if true 138 | stddev: float, stddev for truncated_normal init 139 | weight_decay: float 140 | activation_fn: function 141 | bn: bool, whether to use batch norm 142 | bn_decay: float or float tensor variable in [0,1] 143 | is_training: bool Tensor variable 144 | 145 | Returns: 146 | Variable tensor 147 | """ 148 | with tf.variable_scope(scope) as sc: 149 | kernel_h, kernel_w = kernel_size 150 | num_in_channels = inputs.get_shape()[-1].value 151 | kernel_shape = [kernel_h, kernel_w, 152 | num_in_channels, num_output_channels] 153 | kernel = _variable_with_weight_decay('weights', 154 | shape=kernel_shape, 155 | use_xavier=use_xavier, 156 | stddev=stddev, 157 | wd=weight_decay) 158 | stride_h, stride_w = stride 159 | outputs = tf.nn.conv2d(inputs, kernel, 160 | [1, stride_h, stride_w, 1], 161 | padding=padding) 162 | biases = _variable_on_cpu('biases', [num_output_channels], 163 | tf.constant_initializer(0.0)) 164 | outputs = tf.nn.bias_add(outputs, biases) 165 | 166 | if bn: 167 | outputs = tf.layers.batch_normalization(outputs, momentum=0.99, epsilon=1e-6, training=is_training) 168 | if activation_fn is not None: 169 | outputs = tf.nn.leaky_relu(outputs, alpha=0.2) 170 | return outputs 171 | 172 | 173 | def conv2d_transpose(inputs, 174 | num_output_channels, 175 | kernel_size, 176 | scope, 177 | stride=[1, 1], 178 | padding='SAME', 179 | use_xavier=False, 180 | stddev=1e-3, 181 | weight_decay=0.0, 182 | activation_fn=tf.nn.relu, 183 | bn=False, 184 | bn_decay=None, 185 | is_training=None): 186 | """ 2D convolution transpose with non-linear operation. 187 | 188 | Args: 189 | inputs: 4-D tensor variable BxHxWxC 190 | num_output_channels: int 191 | kernel_size: a list of 2 ints 192 | scope: string 193 | stride: a list of 2 ints 194 | padding: 'SAME' or 'VALID' 195 | use_xavier: bool, use xavier_initializer if true 196 | stddev: float, stddev for truncated_normal init 197 | weight_decay: float 198 | activation_fn: function 199 | bn: bool, whether to use batch norm 200 | bn_decay: float or float tensor variable in [0,1] 201 | is_training: bool Tensor variable 202 | 203 | Returns: 204 | Variable tensor 205 | 206 | Note: conv2d(conv2d_transpose(a, num_out, ksize, stride), a.shape[-1], ksize, stride) == a 207 | """ 208 | with tf.variable_scope(scope) as sc: 209 | kernel_h, kernel_w = kernel_size 210 | num_in_channels = inputs.get_shape()[-1].value 211 | kernel_shape = [kernel_h, kernel_w, 212 | num_output_channels, num_in_channels] # reversed to conv2d 213 | kernel = _variable_with_weight_decay('weights', 214 | shape=kernel_shape, 215 | use_xavier=use_xavier, 216 | stddev=stddev, 217 | wd=weight_decay) 218 | stride_h, stride_w = stride 219 | 220 | # from slim.convolution2d_transpose 221 | def get_deconv_dim(dim_size, stride_size, kernel_size, padding): 222 | dim_size *= stride_size 223 | 224 | if padding == 'VALID' and dim_size is not None: 225 | dim_size += max(kernel_size - stride_size, 0) 226 | return dim_size 227 | 228 | # caculate output shape 229 | batch_size = tf.shape(inputs)[0] 230 | height = tf.shape(inputs)[1] 231 | width = tf.shape(inputs)[2] 232 | out_height = get_deconv_dim(height, stride_h, kernel_h, padding) 233 | out_width = get_deconv_dim(width, stride_w, kernel_w, padding) 234 | output_shape = tf.stack([batch_size, out_height, out_width, num_output_channels], axis=0) 235 | 236 | outputs = tf.nn.conv2d_transpose(inputs, kernel, output_shape, 237 | [1, stride_h, stride_w, 1], 238 | padding=padding) 239 | biases = _variable_on_cpu('biases', [num_output_channels], 240 | tf.constant_initializer(0.0)) 241 | outputs = tf.nn.bias_add(outputs, biases) 242 | 243 | if bn: 244 | # outputs = batch_norm_for_conv2d(outputs, is_training, 245 | # bn_decay=bn_decay, scope='bn') 246 | outputs = tf.layers.batch_normalization(outputs, momentum=0.99, epsilon=1e-6, training=is_training) 247 | if activation_fn is not None: 248 | # outputs = activation_fn(outputs) 249 | outputs = tf.nn.leaky_relu(outputs, alpha=0.2) 250 | return outputs 251 | 252 | 253 | def conv3d(inputs, 254 | num_output_channels, 255 | kernel_size, 256 | scope, 257 | stride=[1, 1, 1], 258 | padding='SAME', 259 | use_xavier=True, 260 | stddev=1e-3, 261 | weight_decay=0.0, 262 | activation_fn=tf.nn.relu, 263 | bn=False, 264 | bn_decay=None, 265 | is_training=None): 266 | """ 3D convolution with non-linear operation. 267 | 268 | Args: 269 | inputs: 5-D tensor variable BxDxHxWxC 270 | num_output_channels: int 271 | kernel_size: a list of 3 ints 272 | scope: string 273 | stride: a list of 3 ints 274 | padding: 'SAME' or 'VALID' 275 | use_xavier: bool, use xavier_initializer if true 276 | stddev: float, stddev for truncated_normal init 277 | weight_decay: float 278 | activation_fn: function 279 | bn: bool, whether to use batch norm 280 | bn_decay: float or float tensor variable in [0,1] 281 | is_training: bool Tensor variable 282 | 283 | Returns: 284 | Variable tensor 285 | """ 286 | with tf.variable_scope(scope) as sc: 287 | kernel_d, kernel_h, kernel_w = kernel_size 288 | num_in_channels = inputs.get_shape()[-1].value 289 | kernel_shape = [kernel_d, kernel_h, kernel_w, 290 | num_in_channels, num_output_channels] 291 | kernel = _variable_with_weight_decay('weights', 292 | shape=kernel_shape, 293 | use_xavier=use_xavier, 294 | stddev=stddev, 295 | wd=weight_decay) 296 | stride_d, stride_h, stride_w = stride 297 | outputs = tf.nn.conv3d(inputs, kernel, 298 | [1, stride_d, stride_h, stride_w, 1], 299 | padding=padding) 300 | biases = _variable_on_cpu('biases', [num_output_channels], 301 | tf.constant_initializer(0.0)) 302 | outputs = tf.nn.bias_add(outputs, biases) 303 | 304 | if bn: 305 | outputs = batch_norm_for_conv3d(outputs, is_training, 306 | bn_decay=bn_decay, scope='bn') 307 | 308 | if activation_fn is not None: 309 | outputs = activation_fn(outputs) 310 | return outputs 311 | 312 | 313 | def fully_connected(inputs, 314 | num_outputs, 315 | scope, 316 | use_xavier=True, 317 | stddev=1e-3, 318 | weight_decay=0.0, 319 | activation_fn=tf.nn.relu, 320 | bn=False, 321 | bn_decay=None, 322 | is_training=None): 323 | """ Fully connected layer with non-linear operation. 324 | 325 | Args: 326 | inputs: 2-D tensor BxN 327 | num_outputs: int 328 | 329 | Returns: 330 | Variable tensor of size B x num_outputs. 331 | """ 332 | with tf.variable_scope(scope) as sc: 333 | num_input_units = inputs.get_shape()[-1].value 334 | weights = _variable_with_weight_decay('weights', 335 | shape=[num_input_units, num_outputs], 336 | use_xavier=use_xavier, 337 | stddev=stddev, 338 | wd=weight_decay) 339 | outputs = tf.matmul(inputs, weights) 340 | biases = _variable_on_cpu('biases', [num_outputs], 341 | tf.constant_initializer(0.0)) 342 | outputs = tf.nn.bias_add(outputs, biases) 343 | 344 | if bn: 345 | outputs = batch_norm_for_fc(outputs, is_training, bn_decay, 'bn') 346 | 347 | if activation_fn is not None: 348 | # outputs = activation_fn(outputs) 349 | outputs = tf.nn.leaky_relu(outputs, alpha=0.2) 350 | return outputs 351 | 352 | 353 | def max_pool2d(inputs, 354 | kernel_size, 355 | scope, 356 | stride=[2, 2], 357 | padding='VALID'): 358 | """ 2D max pooling. 359 | 360 | Args: 361 | inputs: 4-D tensor BxHxWxC 362 | kernel_size: a list of 2 ints 363 | stride: a list of 2 ints 364 | 365 | Returns: 366 | Variable tensor 367 | """ 368 | with tf.variable_scope(scope) as sc: 369 | kernel_h, kernel_w = kernel_size 370 | stride_h, stride_w = stride 371 | outputs = tf.nn.max_pool(inputs, 372 | ksize=[1, kernel_h, kernel_w, 1], 373 | strides=[1, stride_h, stride_w, 1], 374 | padding=padding, 375 | name=sc.name) 376 | return outputs 377 | 378 | 379 | def avg_pool2d(inputs, 380 | kernel_size, 381 | scope, 382 | stride=[2, 2], 383 | padding='VALID'): 384 | """ 2D avg pooling. 385 | 386 | Args: 387 | inputs: 4-D tensor BxHxWxC 388 | kernel_size: a list of 2 ints 389 | stride: a list of 2 ints 390 | 391 | Returns: 392 | Variable tensor 393 | """ 394 | with tf.variable_scope(scope) as sc: 395 | kernel_h, kernel_w = kernel_size 396 | stride_h, stride_w = stride 397 | outputs = tf.nn.avg_pool(inputs, 398 | ksize=[1, kernel_h, kernel_w, 1], 399 | strides=[1, stride_h, stride_w, 1], 400 | padding=padding, 401 | name=sc.name) 402 | return outputs 403 | 404 | 405 | def max_pool3d(inputs, 406 | kernel_size, 407 | scope, 408 | stride=[2, 2, 2], 409 | padding='VALID'): 410 | """ 3D max pooling. 411 | 412 | Args: 413 | inputs: 5-D tensor BxDxHxWxC 414 | kernel_size: a list of 3 ints 415 | stride: a list of 3 ints 416 | 417 | Returns: 418 | Variable tensor 419 | """ 420 | with tf.variable_scope(scope) as sc: 421 | kernel_d, kernel_h, kernel_w = kernel_size 422 | stride_d, stride_h, stride_w = stride 423 | outputs = tf.nn.max_pool3d(inputs, 424 | ksize=[1, kernel_d, kernel_h, kernel_w, 1], 425 | strides=[1, stride_d, stride_h, stride_w, 1], 426 | padding=padding, 427 | name=sc.name) 428 | return outputs 429 | 430 | 431 | def avg_pool3d(inputs, 432 | kernel_size, 433 | scope, 434 | stride=[2, 2, 2], 435 | padding='VALID'): 436 | """ 3D avg pooling. 437 | 438 | Args: 439 | inputs: 5-D tensor BxDxHxWxC 440 | kernel_size: a list of 3 ints 441 | stride: a list of 3 ints 442 | 443 | Returns: 444 | Variable tensor 445 | """ 446 | with tf.variable_scope(scope) as sc: 447 | kernel_d, kernel_h, kernel_w = kernel_size 448 | stride_d, stride_h, stride_w = stride 449 | outputs = tf.nn.avg_pool3d(inputs, 450 | ksize=[1, kernel_d, kernel_h, kernel_w, 1], 451 | strides=[1, stride_d, stride_h, stride_w, 1], 452 | padding=padding, 453 | name=sc.name) 454 | return outputs 455 | 456 | 457 | def batch_norm_template(inputs, is_training, scope, moments_dims, bn_decay): 458 | """ Batch normalization on convolutional maps and beyond... 459 | Ref.: http://stackoverflow.com/questions/33949786/how-could-i-use-batch-normalization-in-tensorflow 460 | 461 | Args: 462 | inputs: Tensor, k-D input ... x C could be BC or BHWC or BDHWC 463 | is_training: boolean tf.Varialbe, true indicates training phase 464 | scope: string, variable scope 465 | moments_dims: a list of ints, indicating dimensions for moments calculation 466 | bn_decay: float or float tensor variable, controling moving average weight 467 | Return: 468 | normed: batch-normalized maps 469 | """ 470 | with tf.variable_scope(scope) as sc: 471 | num_channels = inputs.get_shape()[-1].value 472 | beta = tf.Variable(tf.constant(0.0, shape=[num_channels]), 473 | name='beta', trainable=True) 474 | gamma = tf.Variable(tf.constant(1.0, shape=[num_channels]), 475 | name='gamma', trainable=True) 476 | batch_mean, batch_var = tf.nn.moments(inputs, moments_dims, name='moments') 477 | decay = bn_decay if bn_decay is not None else 0.9 478 | ema = tf.train.ExponentialMovingAverage(decay=decay) 479 | # Operator that maintains moving averages of variables. 480 | ema_apply_op = tf.cond(is_training, 481 | lambda: ema.apply([batch_mean, batch_var]), 482 | lambda: tf.no_op()) 483 | 484 | # Update moving average and return current batch's avg and var. 485 | def mean_var_with_update(): 486 | with tf.control_dependencies([ema_apply_op]): 487 | return tf.identity(batch_mean), tf.identity(batch_var) 488 | 489 | # ema.average returns the Variable holding the average of var. 490 | mean, var = tf.cond(is_training, 491 | mean_var_with_update, 492 | lambda: (ema.average(batch_mean), ema.average(batch_var))) 493 | normed = tf.nn.batch_normalization(inputs, mean, var, beta, gamma, 1e-3) 494 | return normed 495 | 496 | 497 | def batch_norm_for_fc(inputs, is_training, bn_decay, scope): 498 | """ Batch normalization on FC data. 499 | 500 | Args: 501 | inputs: Tensor, 2D BxC input 502 | is_training: boolean tf.Varialbe, true indicates training phase 503 | bn_decay: float or float tensor variable, controling moving average weight 504 | scope: string, variable scope 505 | Return: 506 | normed: batch-normalized maps 507 | """ 508 | return batch_norm_template(inputs, is_training, scope, [0, ], bn_decay) 509 | 510 | 511 | def batch_norm_for_conv1d(inputs, is_training, bn_decay, scope): 512 | """ Batch normalization on 1D convolutional maps. 513 | 514 | Args: 515 | inputs: Tensor, 3D BLC input maps 516 | is_training: boolean tf.Varialbe, true indicates training phase 517 | bn_decay: float or float tensor variable, controling moving average weight 518 | scope: string, variable scope 519 | Return: 520 | normed: batch-normalized maps 521 | """ 522 | return batch_norm_template(inputs, is_training, scope, [0, 1], bn_decay) 523 | 524 | 525 | def batch_norm_for_conv2d(inputs, is_training, bn_decay, scope): 526 | """ Batch normalization on 2D convolutional maps. 527 | 528 | Args: 529 | inputs: Tensor, 4D BHWC input maps 530 | is_training: boolean tf.Varialbe, true indicates training phase 531 | bn_decay: float or float tensor variable, controling moving average weight 532 | scope: string, variable scope 533 | Return: 534 | normed: batch-normalized maps 535 | """ 536 | return batch_norm_template(inputs, is_training, scope, [0, 1, 2], bn_decay) 537 | 538 | 539 | def batch_norm_for_conv3d(inputs, is_training, bn_decay, scope): 540 | """ Batch normalization on 3D convolutional maps. 541 | 542 | Args: 543 | inputs: Tensor, 5D BDHWC input maps 544 | is_training: boolean tf.Varialbe, true indicates training phase 545 | bn_decay: float or float tensor variable, controling moving average weight 546 | scope: string, variable scope 547 | Return: 548 | normed: batch-normalized maps 549 | """ 550 | return batch_norm_template(inputs, is_training, scope, [0, 1, 2, 3], bn_decay) 551 | 552 | 553 | def dropout(inputs, 554 | is_training, 555 | scope, 556 | keep_prob=0.5, 557 | noise_shape=None): 558 | """ Dropout layer. 559 | 560 | Args: 561 | inputs: tensor 562 | is_training: boolean tf.Variable 563 | scope: string 564 | keep_prob: float in [0,1] 565 | noise_shape: list of ints 566 | 567 | Returns: 568 | tensor variable 569 | """ 570 | with tf.variable_scope(scope) as sc: 571 | outputs = tf.cond(is_training, 572 | lambda: tf.nn.dropout(inputs, keep_prob, noise_shape), 573 | lambda: inputs) 574 | return outputs 575 | -------------------------------------------------------------------------------- /img/abstract.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/leofansq/SCF-Net/a25bf06e9481b85d584ac78afc9e5114a8c49cfd/img/abstract.png -------------------------------------------------------------------------------- /img/architecture.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/leofansq/SCF-Net/a25bf06e9481b85d584ac78afc9e5114a8c49cfd/img/architecture.png -------------------------------------------------------------------------------- /img/s3dis.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/leofansq/SCF-Net/a25bf06e9481b85d584ac78afc9e5114a8c49cfd/img/s3dis.png -------------------------------------------------------------------------------- /img/s3dis_vis.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/leofansq/SCF-Net/a25bf06e9481b85d584ac78afc9e5114a8c49cfd/img/s3dis_vis.png -------------------------------------------------------------------------------- /img/semantic3d.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/leofansq/SCF-Net/a25bf06e9481b85d584ac78afc9e5114a8c49cfd/img/semantic3d.png -------------------------------------------------------------------------------- /img/semantic3d_vis.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/leofansq/SCF-Net/a25bf06e9481b85d584ac78afc9e5114a8c49cfd/img/semantic3d_vis.png -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | numpy==1.16.1 2 | h5py==2.10.0 3 | cython==0.29.15 4 | open3d-python==0.3.0 5 | pandas==0.25.3 6 | scikit-learn==0.21.3 7 | scipy==1.4.1 8 | PyYAML==5.1.2 9 | -------------------------------------------------------------------------------- /s3dis_6fold.sh: -------------------------------------------------------------------------------- 1 | python -B s3dis_main.py --gpu 0 --mode train --test_area 1 2 | python -B s3dis_main.py --gpu 0 --mode test --test_area 1 3 | python -B s3dis_main.py --gpu 0 --mode train --test_area 2 4 | python -B s3dis_main.py --gpu 0 --mode test --test_area 2 5 | python -B s3dis_main.py --gpu 0 --mode train --test_area 3 6 | python -B s3dis_main.py --gpu 0 --mode test --test_area 3 7 | python -B s3dis_main.py --gpu 0 --mode train --test_area 4 8 | python -B s3dis_main.py --gpu 0 --mode test --test_area 4 9 | python -B s3dis_main.py --gpu 0 --mode train --test_area 5 10 | python -B s3dis_main.py --gpu 0 --mode test --test_area 5 11 | python -B s3dis_main.py --gpu 0 --mode train --test_area 6 12 | python -B s3dis_main.py --gpu 0 --mode test --test_area 6 13 | 14 | 15 | -------------------------------------------------------------------------------- /s3dis_6fold_eval.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import glob, os, sys 3 | 4 | from helper_ply import read_ply 5 | 6 | if __name__ == '__main__': 7 | base_dir = './test/s3dis_preds' 8 | original_data_dir = './data/S3DIS/original_ply' 9 | data_path = glob.glob(os.path.join(base_dir, '*.ply')) 10 | data_path = np.sort(data_path) 11 | 12 | test_total_correct = 0 13 | test_total_seen = 0 14 | gt_classes = [0 for _ in range(13)] 15 | positive_classes = [0 for _ in range(13)] 16 | true_positive_classes = [0 for _ in range(13)] 17 | 18 | for file_name in data_path: 19 | pred_data = read_ply(file_name) 20 | pred = pred_data['pred'] 21 | original_data = read_ply(os.path.join(original_data_dir, file_name.split('/')[-1][:-4] + '.ply')) 22 | labels = original_data['class'] 23 | points = np.vstack((original_data['x'], original_data['y'], original_data['z'])).T 24 | 25 | correct = np.sum(pred == labels) 26 | print(str(file_name.split('/')[-1][:-4]) + '_acc:' + str(correct / float(len(labels)))) 27 | test_total_correct += correct 28 | test_total_seen += len(labels) 29 | 30 | for j in range(len(labels)): 31 | gt_l = int(labels[j]) 32 | pred_l = int(pred[j]) 33 | gt_classes[gt_l] += 1 34 | positive_classes[pred_l] += 1 35 | true_positive_classes[gt_l] += int(gt_l == pred_l) 36 | 37 | iou_list = [] 38 | for n in range(13): 39 | iou = true_positive_classes[n] / float(gt_classes[n] + positive_classes[n] - true_positive_classes[n]) 40 | iou_list.append(iou) 41 | mean_iou = sum(iou_list) / 13.0 42 | print('eval accuracy: {}'.format(test_total_correct / float(test_total_seen))) 43 | print('mean IOU:{}'.format(mean_iou)) 44 | print(iou_list) 45 | 46 | acc_list = [] 47 | for n in range(13): 48 | acc = true_positive_classes[n] / float(gt_classes[n]) 49 | acc_list.append(acc) 50 | mean_acc = sum(acc_list) / 13.0 51 | print('mAcc value is :{}'.format(mean_acc)) 52 | -------------------------------------------------------------------------------- /s3dis_main.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import numpy as np 3 | import time, pickle, argparse, glob, os 4 | from os.path import join 5 | 6 | from SCFNet import Network 7 | from s3dis_test import ModelTester 8 | from helper_ply import read_ply 9 | from helper_dp import DataProcessing as DP 10 | 11 | 12 | class cfg: 13 | k_n = 16 # KNN 14 | num_layers = 5 # Number of layers 15 | num_points = 40960 # Number of input points 16 | num_classes = 13 # Number of valid classes 17 | sub_grid_size = 0.04 # preprocess_parameter 18 | 19 | batch_size = 4 # batch_size during training 20 | val_batch_size = 20 # batch_size during validation and test 21 | train_steps = 750 # Number of steps per epochs 22 | val_steps = 100 # Number of validation steps per epoch 23 | 24 | sub_sampling_ratio = [4, 4, 4, 4, 2] # sampling ratio of random sampling at each layer 25 | d_out = [16, 64, 128, 256, 512] # feature dimension 26 | 27 | noise_init = 3.5 # noise initial parameter 28 | max_epoch = 100 # maximum epoch during training 29 | learning_rate = 1e-2 # initial learning rate 30 | lr_decays = {i: 0.95 for i in range(0, 500)} # decay rate of learning rate 31 | 32 | train_sum_dir = 'train_log' 33 | saving = True 34 | saving_path = None 35 | 36 | 37 | class S3DIS: 38 | def __init__(self, test_area_idx): 39 | self.name = 'S3DIS' 40 | self.path = './data/S3DIS' 41 | self.label_to_names = {0: 'ceiling', 42 | 1: 'floor', 43 | 2: 'wall', 44 | 3: 'beam', 45 | 4: 'column', 46 | 5: 'window', 47 | 6: 'door', 48 | 7: 'table', 49 | 8: 'chair', 50 | 9: 'sofa', 51 | 10: 'bookcase', 52 | 11: 'board', 53 | 12: 'clutter'} 54 | self.num_classes = len(self.label_to_names) 55 | self.label_values = np.sort([k for k, v in self.label_to_names.items()]) 56 | self.label_to_idx = {l: i for i, l in enumerate(self.label_values)} 57 | self.ignored_labels = np.array([]) 58 | 59 | self.val_split = 'Area_' + str(test_area_idx) 60 | self.all_files = glob.glob(join(self.path, 'original_ply', '*.ply')) 61 | 62 | # Initiate containers 63 | self.val_proj = [] 64 | self.val_labels = [] 65 | self.possibility = {} 66 | self.min_possibility = {} 67 | self.input_trees = {'training': [], 'validation': []} 68 | self.input_colors = {'training': [], 'validation': []} 69 | self.input_labels = {'training': [], 'validation': []} 70 | self.input_names = {'training': [], 'validation': []} 71 | self.load_sub_sampled_clouds(cfg.sub_grid_size) 72 | 73 | def load_sub_sampled_clouds(self, sub_grid_size): 74 | tree_path = join(self.path, 'input_{:.3f}'.format(sub_grid_size)) 75 | for i, file_path in enumerate(self.all_files): 76 | t0 = time.time() 77 | cloud_name = file_path.split('/')[-1][:-4] 78 | if self.val_split in cloud_name: 79 | cloud_split = 'validation' 80 | else: 81 | cloud_split = 'training' 82 | 83 | # Name of the input files 84 | kd_tree_file = join(tree_path, '{:s}_KDTree.pkl'.format(cloud_name)) 85 | sub_ply_file = join(tree_path, '{:s}.ply'.format(cloud_name)) 86 | 87 | data = read_ply(sub_ply_file) 88 | sub_colors = np.vstack((data['red'], data['green'], data['blue'])).T 89 | sub_labels = data['class'] 90 | 91 | # Read pkl with search tree 92 | with open(kd_tree_file, 'rb') as f: 93 | search_tree = pickle.load(f) 94 | 95 | self.input_trees[cloud_split] += [search_tree] 96 | self.input_colors[cloud_split] += [sub_colors] 97 | self.input_labels[cloud_split] += [sub_labels] 98 | self.input_names[cloud_split] += [cloud_name] 99 | 100 | size = sub_colors.shape[0] * 4 * 7 101 | print('{:s} {:.1f} MB loaded in {:.1f}s'.format(kd_tree_file.split('/')[-1], size * 1e-6, time.time() - t0)) 102 | 103 | print('\nPreparing reprojected indices for testing') 104 | 105 | # Get validation and test reprojected indices 106 | for i, file_path in enumerate(self.all_files): 107 | t0 = time.time() 108 | cloud_name = file_path.split('/')[-1][:-4] 109 | 110 | # Validation projection and labels 111 | if self.val_split in cloud_name: 112 | proj_file = join(tree_path, '{:s}_proj.pkl'.format(cloud_name)) 113 | with open(proj_file, 'rb') as f: 114 | proj_idx, labels = pickle.load(f) 115 | self.val_proj += [proj_idx] 116 | self.val_labels += [labels] 117 | print('{:s} done in {:.1f}s'.format(cloud_name, time.time() - t0)) 118 | 119 | # Generate the input data flow 120 | def get_batch_gen(self, split): 121 | if split == 'training': 122 | num_per_epoch = cfg.train_steps * cfg.batch_size 123 | elif split == 'validation': 124 | num_per_epoch = cfg.val_steps * cfg.val_batch_size 125 | 126 | self.possibility[split] = [] 127 | self.min_possibility[split] = [] 128 | # Random initialize 129 | for i, tree in enumerate(self.input_colors[split]): 130 | self.possibility[split] += [np.random.rand(tree.data.shape[0]) * 1e-3] 131 | self.min_possibility[split] += [float(np.min(self.possibility[split][-1]))] 132 | 133 | def spatially_regular_gen(): 134 | # Generator loop 135 | for i in range(num_per_epoch): 136 | 137 | # Choose the cloud with the lowest probability 138 | cloud_idx = int(np.argmin(self.min_possibility[split])) 139 | 140 | # choose the point with the minimum of possibility in the cloud as query point 141 | point_ind = np.argmin(self.possibility[split][cloud_idx]) 142 | 143 | # Get all points within the cloud from tree structure 144 | points = np.array(self.input_trees[split][cloud_idx].data, copy=False) 145 | 146 | # Center point of input region 147 | center_point = points[point_ind, :].reshape(1, -1) 148 | 149 | # Add noise to the center point 150 | noise = np.random.normal(scale=cfg.noise_init / 10, size=center_point.shape) 151 | pick_point = center_point + noise.astype(center_point.dtype) 152 | 153 | # Check if the number of points in the selected cloud is less than the predefined num_points 154 | if len(points) < cfg.num_points: 155 | # Query all points within the cloud 156 | queried_idx = self.input_trees[split][cloud_idx].query(pick_point, k=len(points))[1][0] 157 | else: 158 | # Query the predefined number of points 159 | queried_idx = self.input_trees[split][cloud_idx].query(pick_point, k=cfg.num_points)[1][0] 160 | 161 | # Shuffle index 162 | queried_idx = DP.shuffle_idx(queried_idx) 163 | # Get corresponding points and colors based on the index 164 | queried_pc_xyz = points[queried_idx] 165 | queried_pc_xyz = queried_pc_xyz - pick_point 166 | queried_pc_colors = self.input_colors[split][cloud_idx][queried_idx] 167 | queried_pc_labels = self.input_labels[split][cloud_idx][queried_idx] 168 | 169 | # Update the possibility of the selected points 170 | dists = np.sum(np.square((points[queried_idx] - pick_point).astype(np.float32)), axis=1) 171 | delta = np.square(1 - dists / np.max(dists)) 172 | self.possibility[split][cloud_idx][queried_idx] += delta 173 | self.min_possibility[split][cloud_idx] = float(np.min(self.possibility[split][cloud_idx])) 174 | 175 | # up_sampled with replacement 176 | if len(points) < cfg.num_points: 177 | queried_pc_xyz, queried_pc_colors, queried_idx, queried_pc_labels = \ 178 | DP.data_aug(queried_pc_xyz, queried_pc_colors, queried_pc_labels, queried_idx, cfg.num_points) 179 | 180 | if True: 181 | yield (queried_pc_xyz.astype(np.float32), 182 | queried_pc_colors.astype(np.float32), 183 | queried_pc_labels, 184 | queried_idx.astype(np.int32), 185 | np.array([cloud_idx], dtype=np.int32)) 186 | 187 | gen_func = spatially_regular_gen 188 | gen_types = (tf.float32, tf.float32, tf.int32, tf.int32, tf.int32) 189 | gen_shapes = ([None, 3], [None, 3], [None], [None], [None]) 190 | return gen_func, gen_types, gen_shapes 191 | 192 | @staticmethod 193 | def get_tf_mapping2(): 194 | # Collect flat inputs 195 | def tf_map(batch_xyz, batch_features, batch_labels, batch_pc_idx, batch_cloud_idx): 196 | batch_features = tf.concat([batch_xyz, batch_features], axis=-1) 197 | 198 | input_points = [] 199 | input_neighbors = [] 200 | input_pools = [] 201 | input_up_samples = [] 202 | 203 | for i in range(cfg.num_layers): 204 | # KNN 205 | neighbour_idx = tf.py_func(DP.knn_search, [batch_xyz, batch_xyz, cfg.k_n], tf.int32) 206 | 207 | sub_points = batch_xyz[:, :tf.shape(batch_xyz)[1] // cfg.sub_sampling_ratio[i], :] 208 | pool_i = neighbour_idx[:, :tf.shape(batch_xyz)[1] // cfg.sub_sampling_ratio[i], :] 209 | up_i = tf.py_func(DP.knn_search, [sub_points, batch_xyz, 1], tf.int32) 210 | input_points.append(batch_xyz) 211 | input_neighbors.append(neighbour_idx) 212 | input_pools.append(pool_i) 213 | input_up_samples.append(up_i) 214 | batch_xyz = sub_points 215 | 216 | input_list = input_points + input_neighbors + input_pools + input_up_samples 217 | input_list += [batch_features, batch_labels, batch_pc_idx, batch_cloud_idx] 218 | 219 | return input_list 220 | 221 | return tf_map 222 | 223 | def init_input_pipeline(self): 224 | print('Initiating input pipelines') 225 | cfg.ignored_label_inds = [self.label_to_idx[ign_label] for ign_label in self.ignored_labels] 226 | gen_function, gen_types, gen_shapes = self.get_batch_gen('training') 227 | gen_function_val, _, _ = self.get_batch_gen('validation') 228 | self.train_data = tf.data.Dataset.from_generator(gen_function, gen_types, gen_shapes) 229 | self.val_data = tf.data.Dataset.from_generator(gen_function_val, gen_types, gen_shapes) 230 | 231 | self.batch_train_data = self.train_data.batch(cfg.batch_size) 232 | self.batch_val_data = self.val_data.batch(cfg.val_batch_size) 233 | map_func = self.get_tf_mapping2() 234 | 235 | self.batch_train_data = self.batch_train_data.map(map_func=map_func) 236 | self.batch_val_data = self.batch_val_data.map(map_func=map_func) 237 | 238 | self.batch_train_data = self.batch_train_data.prefetch(cfg.batch_size) 239 | self.batch_val_data = self.batch_val_data.prefetch(cfg.val_batch_size) 240 | 241 | iter = tf.data.Iterator.from_structure(self.batch_train_data.output_types, self.batch_train_data.output_shapes) 242 | self.flat_inputs = iter.get_next() 243 | self.train_init_op = iter.make_initializer(self.batch_train_data) 244 | self.val_init_op = iter.make_initializer(self.batch_val_data) 245 | 246 | 247 | if __name__ == '__main__': 248 | parser = argparse.ArgumentParser() 249 | parser.add_argument('--gpu', type=int, default=0, help='the number of GPUs to use [default: 0]') 250 | parser.add_argument('--test_area', type=int, default=5, help='Which area to use for test, option: 1-6 [default: 5]') 251 | parser.add_argument('--mode', type=str, default='train', help='options: train, test, vis') 252 | parser.add_argument('--model_path', type=str, default='None', help='pretrained model path') 253 | FLAGS = parser.parse_args() 254 | 255 | os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" 256 | os.environ['CUDA_VISIBLE_DEVICES'] = str(FLAGS.gpu) 257 | os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' 258 | Mode = FLAGS.mode 259 | 260 | test_area = FLAGS.test_area 261 | dataset = S3DIS(test_area) 262 | dataset.init_input_pipeline() 263 | 264 | if Mode == 'train': 265 | model = Network(dataset, cfg) 266 | model.train(dataset) 267 | elif Mode == 'test': 268 | cfg.saving = False 269 | model = Network(dataset, cfg) 270 | if FLAGS.model_path is not 'None': 271 | chosen_snap = FLAGS.model_path 272 | else: 273 | chosen_snapshot = -1 274 | logs = np.sort([os.path.join('results', f) for f in os.listdir('results') if f.startswith('Log')]) 275 | chosen_folder = logs[-1] 276 | snap_path = join(chosen_folder, 'snapshots') 277 | snap_steps = [int(f[:-5].split('-')[-1]) for f in os.listdir(snap_path) if f[-5:] == '.meta'] 278 | chosen_step = np.sort(snap_steps)[-1] 279 | chosen_snap = os.path.join(snap_path, 'snap-{:d}'.format(chosen_step)) 280 | tester = ModelTester(model, dataset, restore_snap=chosen_snap) 281 | tester.test(model, dataset) 282 | -------------------------------------------------------------------------------- /s3dis_test.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import numpy as np 3 | import time 4 | from os import makedirs 5 | from os.path import exists, join 6 | from sklearn.metrics import confusion_matrix 7 | 8 | from helper_ply import write_ply 9 | from helper_dp import DataProcessing as DP 10 | 11 | 12 | def log_out(out_str, log_f_out): 13 | log_f_out.write(out_str + '\n') 14 | log_f_out.flush() 15 | print(out_str) 16 | 17 | 18 | class ModelTester: 19 | def __init__(self, model, dataset, restore_snap=None): 20 | my_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES) 21 | self.saver = tf.train.Saver(my_vars, max_to_keep=100) 22 | self.Log_file = open('log_test_' + str(dataset.val_split) + '.txt', 'a') 23 | 24 | # Create a session for running Ops on the Graph. 25 | on_cpu = False 26 | if on_cpu: 27 | c_proto = tf.ConfigProto(device_count={'GPU': 0}) 28 | else: 29 | c_proto = tf.ConfigProto() 30 | c_proto.gpu_options.allow_growth = True 31 | self.sess = tf.Session(config=c_proto) 32 | self.sess.run(tf.global_variables_initializer()) 33 | 34 | # Load trained model 35 | if restore_snap is not None: 36 | self.saver.restore(self.sess, restore_snap) 37 | print("Model restored from " + restore_snap) 38 | 39 | self.prob_logits = tf.nn.softmax(model.logits) 40 | 41 | # Initiate global prediction over all test clouds 42 | self.test_probs = [np.zeros(shape=[l.shape[0], model.config.num_classes], dtype=np.float32) 43 | for l in dataset.input_labels['validation']] 44 | 45 | def test(self, model, dataset, num_votes=100): 46 | 47 | # Smoothing parameter for votes 48 | test_smooth = 0.95 49 | 50 | # Initialise iterator with validation/test data 51 | self.sess.run(dataset.val_init_op) 52 | 53 | # Number of points per class in validation set 54 | val_proportions = np.zeros(model.config.num_classes, dtype=np.float32) 55 | i = 0 56 | for label_val in dataset.label_values: 57 | if label_val not in dataset.ignored_labels: 58 | val_proportions[i] = np.sum([np.sum(labels == label_val) for labels in dataset.val_labels]) 59 | i += 1 60 | 61 | # Test saving path 62 | saving_path = time.strftime('results/Log_%Y-%m-%d_%H-%M-%S', time.gmtime()) 63 | test_path = join('test') # Save all results in the same path for the convinience of 6fold evaluation 64 | makedirs(test_path) if not exists(test_path) else None 65 | makedirs(join(test_path, 's3dis_preds')) if not exists(join(test_path, 's3dis_preds')) else None 66 | 67 | step_id = 0 68 | epoch_id = 0 69 | last_min = -0.5 70 | 71 | while last_min < num_votes: 72 | try: 73 | ops = (self.prob_logits, 74 | model.labels, 75 | model.inputs['input_inds'], 76 | model.inputs['cloud_inds'], 77 | ) 78 | 79 | stacked_probs, stacked_labels, point_idx, cloud_idx = self.sess.run(ops, {model.is_training: False}) 80 | correct = np.sum(np.argmax(stacked_probs, axis=1) == stacked_labels) 81 | acc = correct / float(np.prod(np.shape(stacked_labels))) 82 | print('step' + str(step_id) + ' acc:' + str(acc)) 83 | stacked_probs = np.reshape(stacked_probs, [model.config.val_batch_size, model.config.num_points, 84 | model.config.num_classes]) 85 | 86 | for j in range(np.shape(stacked_probs)[0]): 87 | probs = stacked_probs[j, :, :] 88 | p_idx = point_idx[j, :] 89 | c_i = cloud_idx[j][0] 90 | self.test_probs[c_i][p_idx] = test_smooth * self.test_probs[c_i][p_idx] + (1 - test_smooth) * probs 91 | step_id += 1 92 | 93 | except tf.errors.OutOfRangeError: 94 | 95 | new_min = np.min(dataset.min_possibility['validation']) 96 | log_out('Epoch {:3d}, end. Min possibility = {:.1f}'.format(epoch_id, new_min), self.Log_file) 97 | 98 | if last_min + 1 < new_min: 99 | 100 | # Update last_min 101 | last_min += 1 102 | 103 | # Show vote results (On subcloud so it is not the good values here) 104 | log_out('\nConfusion on sub clouds', self.Log_file) 105 | confusion_list = [] 106 | 107 | num_val = len(dataset.input_labels['validation']) 108 | 109 | for i_test in range(num_val): 110 | probs = self.test_probs[i_test] 111 | preds = dataset.label_values[np.argmax(probs, axis=1)].astype(np.int32) 112 | labels = dataset.input_labels['validation'][i_test] 113 | 114 | # Confs 115 | confusion_list += [confusion_matrix(labels, preds, dataset.label_values)] 116 | 117 | # Regroup confusions 118 | C = np.sum(np.stack(confusion_list), axis=0).astype(np.float32) 119 | 120 | # Rescale with the right number of point per class 121 | C *= np.expand_dims(val_proportions / (np.sum(C, axis=1) + 1e-6), 1) 122 | 123 | # Compute IoUs 124 | IoUs = DP.IoU_from_confusions(C) 125 | m_IoU = np.mean(IoUs) 126 | s = '{:5.2f} | '.format(100 * m_IoU) 127 | for IoU in IoUs: 128 | s += '{:5.2f} '.format(100 * IoU) 129 | log_out(s + '\n', self.Log_file) 130 | 131 | if int(np.ceil(new_min)) % 1 == 0: 132 | 133 | # Project predictions 134 | log_out('\nReproject Vote #{:d}'.format(int(np.floor(new_min))), self.Log_file) 135 | proj_probs_list = [] 136 | 137 | for i_val in range(num_val): 138 | # Reproject probs back to the evaluations points 139 | proj_idx = dataset.val_proj[i_val] 140 | probs = self.test_probs[i_val][proj_idx, :] 141 | proj_probs_list += [probs] 142 | 143 | # Show vote results 144 | log_out('Confusion on full clouds', self.Log_file) 145 | confusion_list = [] 146 | for i_test in range(num_val): 147 | # Get the predicted labels 148 | preds = dataset.label_values[np.argmax(proj_probs_list[i_test], axis=1)].astype(np.uint8) 149 | 150 | # Confusion 151 | labels = dataset.val_labels[i_test] 152 | acc = np.sum(preds == labels) / len(labels) 153 | log_out(dataset.input_names['validation'][i_test] + ' Acc:' + str(acc), self.Log_file) 154 | 155 | confusion_list += [confusion_matrix(labels, preds, dataset.label_values)] 156 | name = dataset.input_names['validation'][i_test] + '.ply' 157 | write_ply(join(test_path, 's3dis_preds', name), [preds, labels], ['pred', 'label']) 158 | 159 | # Regroup confusions 160 | C = np.sum(np.stack(confusion_list), axis=0) 161 | 162 | IoUs = DP.IoU_from_confusions(C) 163 | m_IoU = np.mean(IoUs) 164 | s = '{:5.2f} | '.format(100 * m_IoU) 165 | for IoU in IoUs: 166 | s += '{:5.2f} '.format(100 * IoU) 167 | log_out('-' * len(s), self.Log_file) 168 | log_out(s, self.Log_file) 169 | log_out('-' * len(s) + '\n', self.Log_file) 170 | print('finished \n') 171 | self.sess.close() 172 | return 173 | 174 | self.sess.run(dataset.val_init_op) 175 | epoch_id += 1 176 | step_id = 0 177 | continue 178 | 179 | return 180 | -------------------------------------------------------------------------------- /semantic3d_main.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import numpy as np 3 | import pickle, argparse, os 4 | from os.path import join, exists 5 | 6 | from SCFNet import Network 7 | from semantic3d_test import ModelTester 8 | from helper_ply import read_ply 9 | from helper_dp import DataProcessing as DP 10 | 11 | class cfg: 12 | k_n = 16 # KNN 13 | num_layers = 5 # Number of layers 14 | num_points = 65536 # Number of input points 15 | num_classes = 8 # Number of valid classes 16 | sub_grid_size = 0.06 # preprocess_parameter 17 | 18 | batch_size = 3 # batch_size during training 19 | val_batch_size = 16 # batch_size during validation and test 20 | train_steps = 1000 # Number of steps per epochs 21 | val_steps = 100 # Number of validation steps per epoch 22 | 23 | sub_sampling_ratio = [4, 4, 4, 4, 2] # sampling ratio of random sampling at each layer 24 | d_out = [16, 64, 128, 256, 512] # feature dimension 25 | 26 | noise_init = 3.5 # noise initial parameter 27 | max_epoch = 100 # maximum epoch during training 28 | learning_rate = 1e-2 # initial learning rate 29 | lr_decays = {i: 0.95 for i in range(0, 500)} # decay rate of learning rate 30 | 31 | train_sum_dir = 'train_log' 32 | saving = True 33 | saving_path = None 34 | 35 | augment_scale_anisotropic = True 36 | augment_symmetries = [True, False, False] 37 | augment_rotation = 'vertical' 38 | augment_scale_min = 0.8 39 | augment_scale_max = 1.2 40 | augment_noise = 0.001 41 | augment_occlusion = 'none' 42 | augment_color = 0.8 43 | 44 | class Semantic3D: 45 | def __init__(self): 46 | self.name = 'Semantic3D' 47 | self.path = './data/semantic3d' 48 | self.label_to_names = {0: 'unlabeled', 49 | 1: 'man-made terrain', 50 | 2: 'natural terrain', 51 | 3: 'high vegetation', 52 | 4: 'low vegetation', 53 | 5: 'buildings', 54 | 6: 'hard scape', 55 | 7: 'scanning artefacts', 56 | 8: 'cars'} 57 | self.num_classes = len(self.label_to_names) 58 | self.label_values = np.sort([k for k, v in self.label_to_names.items()]) 59 | self.label_to_idx = {l: i for i, l in enumerate(self.label_values)} 60 | self.ignored_labels = np.sort([0]) 61 | 62 | self.original_folder = join(self.path, 'original_data') 63 | self.full_pc_folder = join(self.path, 'original_ply') 64 | self.sub_pc_folder = join(self.path, 'input_{:.3f}'.format(cfg.sub_grid_size)) 65 | 66 | # Following KPConv to do the train-validation split 67 | self.all_splits = [0, 1, 4, 5, 3, 4, 3, 0, 1, 2, 3, 4, 2, 0, 5] 68 | self.val_split = 1 69 | 70 | # Initial training-validation-testing files 71 | self.train_files = [] 72 | self.val_files = [] 73 | self.test_files = [] 74 | cloud_names = [file_name[:-4] for file_name in os.listdir(self.original_folder) if file_name[-4:] == '.txt'] 75 | for pc_name in cloud_names: 76 | if exists(join(self.original_folder, pc_name + '.labels')): 77 | self.train_files.append(join(self.sub_pc_folder, pc_name + '.ply')) 78 | else: 79 | self.test_files.append(join(self.full_pc_folder, pc_name + '.ply')) 80 | # elif '-reduced' in pc_name: 81 | # self.test_files.append(join(self.full_pc_folder, pc_name + '.ply')) 82 | 83 | self.train_files = np.sort(self.train_files) 84 | self.test_files = np.sort(self.test_files) 85 | 86 | for i, file_path in enumerate(self.train_files): 87 | if self.all_splits[i] == self.val_split: 88 | self.val_files.append(file_path) 89 | 90 | self.train_files = np.sort([x for x in self.train_files if x not in self.val_files]) 91 | 92 | # Initiate containers 93 | self.val_proj = [] 94 | self.val_labels = [] 95 | self.test_proj = [] 96 | self.test_labels = [] 97 | 98 | self.possibility = {} 99 | self.min_possibility = {} 100 | self.class_weight = {} 101 | self.input_trees = {'training': [], 'validation': [], 'test': []} 102 | self.input_colors = {'training': [], 'validation': [], 'test': []} 103 | self.input_labels = {'training': [], 'validation': []} 104 | 105 | # Ascii files dict for testing 106 | self.ascii_files = { 107 | 'MarketplaceFeldkirch_Station4_rgb_intensity-reduced.ply': 'marketsquarefeldkirch4-reduced.labels', 108 | 'sg27_station10_rgb_intensity-reduced.ply': 'sg27_10-reduced.labels', 109 | 'sg28_Station2_rgb_intensity-reduced.ply': 'sg28_2-reduced.labels', 110 | 'StGallenCathedral_station6_rgb_intensity-reduced.ply': 'stgallencathedral6-reduced.labels', 111 | 'birdfountain_station1_xyz_intensity_rgb.ply': 'birdfountain1.labels', 112 | 'castleblatten_station1_intensity_rgb.ply': 'castleblatten1.labels', 113 | 'castleblatten_station5_xyz_intensity_rgb.ply': 'castleblatten5.labels', 114 | 'marketplacefeldkirch_station1_intensity_rgb.ply': 'marketsquarefeldkirch1.labels', 115 | 'marketplacefeldkirch_station4_intensity_rgb.ply': 'marketsquarefeldkirch4.labels', 116 | 'marketplacefeldkirch_station7_intensity_rgb.ply': 'marketsquarefeldkirch7.labels', 117 | 'sg27_station10_intensity_rgb.ply': 'sg27_10.labels', 118 | 'sg27_station3_intensity_rgb.ply': 'sg27_3.labels', 119 | 'sg27_station6_intensity_rgb.ply': 'sg27_6.labels', 120 | 'sg27_station8_intensity_rgb.ply': 'sg27_8.labels', 121 | 'sg28_station2_intensity_rgb.ply': 'sg28_2.labels', 122 | 'sg28_station5_xyz_intensity_rgb.ply': 'sg28_5.labels', 123 | 'stgallencathedral_station1_intensity_rgb.ply': 'stgallencathedral1.labels', 124 | 'stgallencathedral_station3_intensity_rgb.ply': 'stgallencathedral3.labels', 125 | 'stgallencathedral_station6_intensity_rgb.ply': 'stgallencathedral6.labels'} 126 | 127 | self.load_sub_sampled_clouds(cfg.sub_grid_size) 128 | 129 | def load_sub_sampled_clouds(self, sub_grid_size): 130 | 131 | tree_path = join(self.path, 'input_{:.3f}'.format(sub_grid_size)) 132 | files = np.hstack((self.train_files, self.val_files, self.test_files)) 133 | 134 | for i, file_path in enumerate(files): 135 | cloud_name = file_path.split('/')[-1][:-4] 136 | print('Load_pc_' + str(i) + ': ' + cloud_name) 137 | if file_path in self.val_files: 138 | cloud_split = 'validation' 139 | elif file_path in self.train_files: 140 | cloud_split = 'training' 141 | else: 142 | cloud_split = 'test' 143 | # elif file_path in self.test_files: 144 | # cloud_split = 'test' 145 | 146 | # Name of the input files 147 | kd_tree_file = join(tree_path, '{:s}_KDTree.pkl'.format(cloud_name)) 148 | sub_ply_file = join(tree_path, '{:s}.ply'.format(cloud_name)) 149 | 150 | # read ply with data 151 | data = read_ply(sub_ply_file) 152 | sub_colors = np.vstack((data['red'], data['green'], data['blue'])).T 153 | if cloud_split == 'test': 154 | sub_labels = None 155 | else: 156 | sub_labels = data['class'] 157 | 158 | # Read pkl with search tree 159 | with open(kd_tree_file, 'rb') as f: 160 | search_tree = pickle.load(f) 161 | 162 | self.input_trees[cloud_split] += [search_tree] 163 | self.input_colors[cloud_split] += [sub_colors] 164 | if cloud_split in ['training', 'validation']: 165 | self.input_labels[cloud_split] += [sub_labels] 166 | 167 | # Get validation and test re_projection indices 168 | print('\nPreparing reprojection indices for validation and test') 169 | 170 | for i, file_path in enumerate(files): 171 | 172 | # get cloud name and split 173 | cloud_name = file_path.split('/')[-1][:-4] 174 | 175 | # Validation projection and labels 176 | if file_path in self.val_files: 177 | proj_file = join(tree_path, '{:s}_proj.pkl'.format(cloud_name)) 178 | with open(proj_file, 'rb') as f: 179 | proj_idx, labels = pickle.load(f) 180 | self.val_proj += [proj_idx] 181 | self.val_labels += [labels] 182 | 183 | # Test projection 184 | if file_path in self.test_files: 185 | proj_file = join(tree_path, '{:s}_proj.pkl'.format(cloud_name)) 186 | with open(proj_file, 'rb') as f: 187 | proj_idx, labels = pickle.load(f) 188 | self.test_proj += [proj_idx] 189 | self.test_labels += [labels] 190 | print('finished') 191 | return 192 | 193 | # Generate the input data flow 194 | def get_batch_gen(self, split): 195 | if split == 'training': 196 | num_per_epoch = cfg.train_steps * cfg.batch_size 197 | elif split == 'validation': 198 | num_per_epoch = cfg.val_steps * cfg.val_batch_size 199 | elif split == 'test': 200 | num_per_epoch = cfg.val_steps * cfg.val_batch_size 201 | 202 | # Reset possibility 203 | self.possibility[split] = [] 204 | self.min_possibility[split] = [] 205 | self.class_weight[split] = [] 206 | 207 | # Random initialize 208 | for i, tree in enumerate(self.input_trees[split]): 209 | self.possibility[split] += [np.random.rand(tree.data.shape[0]) * 1e-3] 210 | self.min_possibility[split] += [float(np.min(self.possibility[split][-1]))] 211 | 212 | if split != 'test': 213 | _, num_class_total = np.unique(np.hstack(self.input_labels[split]), return_counts=True) 214 | self.class_weight[split] += [np.squeeze([num_class_total / np.sum(num_class_total)], axis=0)] 215 | 216 | def spatially_regular_gen(): 217 | 218 | # Generator loop 219 | for i in range(num_per_epoch): # num_per_epoch 220 | 221 | # Choose the cloud with the lowest probability 222 | cloud_idx = int(np.argmin(self.min_possibility[split])) 223 | 224 | # choose the point with the minimum of possibility in the cloud as query point 225 | point_ind = np.argmin(self.possibility[split][cloud_idx]) 226 | 227 | # Get all points within the cloud from tree structure 228 | points = np.array(self.input_trees[split][cloud_idx].data, copy=False) 229 | 230 | # Center point of input region 231 | center_point = points[point_ind, :].reshape(1, -1) 232 | 233 | # Add noise to the center point 234 | noise = np.random.normal(scale=cfg.noise_init / 10, size=center_point.shape) 235 | pick_point = center_point + noise.astype(center_point.dtype) 236 | query_idx = self.input_trees[split][cloud_idx].query(pick_point, k=cfg.num_points)[1][0] 237 | 238 | # Shuffle index 239 | query_idx = DP.shuffle_idx(query_idx) 240 | 241 | # Get corresponding points and colors based on the index 242 | queried_pc_xyz = points[query_idx] 243 | queried_pc_xyz[:, 0:2] = queried_pc_xyz[:, 0:2] - pick_point[:, 0:2] 244 | queried_pc_colors = self.input_colors[split][cloud_idx][query_idx] 245 | if split == 'test': 246 | queried_pc_labels = np.zeros(queried_pc_xyz.shape[0]) 247 | queried_pt_weight = 1 248 | else: 249 | queried_pc_labels = self.input_labels[split][cloud_idx][query_idx] 250 | queried_pc_labels = np.array([self.label_to_idx[l] for l in queried_pc_labels]) 251 | queried_pt_weight = np.array([self.class_weight[split][0][n] for n in queried_pc_labels]) 252 | 253 | # Update the possibility of the selected points 254 | dists = np.sum(np.square((points[query_idx] - pick_point).astype(np.float32)), axis=1) 255 | delta = np.square(1 - dists / np.max(dists)) * queried_pt_weight 256 | self.possibility[split][cloud_idx][query_idx] += delta 257 | self.min_possibility[split][cloud_idx] = float(np.min(self.possibility[split][cloud_idx])) 258 | 259 | if True: 260 | yield (queried_pc_xyz, 261 | queried_pc_colors.astype(np.float32), 262 | queried_pc_labels, 263 | query_idx.astype(np.int32), 264 | np.array([cloud_idx], dtype=np.int32)) 265 | 266 | gen_func = spatially_regular_gen 267 | gen_types = (tf.float32, tf.float32, tf.int32, tf.int32, tf.int32) 268 | gen_shapes = ([None, 3], [None, 3], [None], [None], [None]) 269 | return gen_func, gen_types, gen_shapes 270 | 271 | def get_tf_mapping(self): 272 | # Collect flat inputs 273 | def tf_map(batch_xyz, batch_features, batch_labels, batch_pc_idx, batch_cloud_idx): 274 | batch_features = tf.map_fn(self.tf_augment_input, [batch_xyz, batch_features], dtype=tf.float32) 275 | input_points = [] 276 | input_neighbors = [] 277 | input_pools = [] 278 | input_up_samples = [] 279 | 280 | for i in range(cfg.num_layers): 281 | neigh_idx = tf.py_func(DP.knn_search, [batch_xyz, batch_xyz, cfg.k_n], tf.int32) 282 | sub_points = batch_xyz[:, :tf.shape(batch_xyz)[1] // cfg.sub_sampling_ratio[i], :] 283 | pool_i = neigh_idx[:, :tf.shape(batch_xyz)[1] // cfg.sub_sampling_ratio[i], :] 284 | up_i = tf.py_func(DP.knn_search, [sub_points, batch_xyz, 1], tf.int32) 285 | input_points.append(batch_xyz) 286 | input_neighbors.append(neigh_idx) 287 | input_pools.append(pool_i) 288 | input_up_samples.append(up_i) 289 | batch_xyz = sub_points 290 | 291 | input_list = input_points + input_neighbors + input_pools + input_up_samples 292 | input_list += [batch_features, batch_labels, batch_pc_idx, batch_cloud_idx] 293 | 294 | return input_list 295 | 296 | return tf_map 297 | 298 | # data augmentation 299 | @staticmethod 300 | def tf_augment_input(inputs): 301 | xyz = inputs[0] 302 | features = inputs[1] 303 | theta = tf.random_uniform((1,), minval=0, maxval=2 * np.pi) 304 | # Rotation matrices 305 | c, s = tf.cos(theta), tf.sin(theta) 306 | cs0 = tf.zeros_like(c) 307 | cs1 = tf.ones_like(c) 308 | R = tf.stack([c, -s, cs0, s, c, cs0, cs0, cs0, cs1], axis=1) 309 | stacked_rots = tf.reshape(R, (3, 3)) 310 | 311 | # Apply rotations 312 | transformed_xyz = tf.reshape(tf.matmul(xyz, stacked_rots), [-1, 3]) 313 | # Choose random scales for each example 314 | min_s = cfg.augment_scale_min 315 | max_s = cfg.augment_scale_max 316 | if cfg.augment_scale_anisotropic: 317 | s = tf.random_uniform((1, 3), minval=min_s, maxval=max_s) 318 | else: 319 | s = tf.random_uniform((1, 1), minval=min_s, maxval=max_s) 320 | 321 | symmetries = [] 322 | for i in range(3): 323 | if cfg.augment_symmetries[i]: 324 | symmetries.append(tf.round(tf.random_uniform((1, 1))) * 2 - 1) 325 | else: 326 | symmetries.append(tf.ones([1, 1], dtype=tf.float32)) 327 | s *= tf.concat(symmetries, 1) 328 | 329 | # Create N x 3 vector of scales to multiply with stacked_points 330 | stacked_scales = tf.tile(s, [tf.shape(transformed_xyz)[0], 1]) 331 | 332 | # Apply scales 333 | transformed_xyz = transformed_xyz * stacked_scales 334 | 335 | noise = tf.random_normal(tf.shape(transformed_xyz), stddev=cfg.augment_noise) 336 | transformed_xyz = transformed_xyz + noise 337 | rgb = features[:, :3] 338 | stacked_features = tf.concat([transformed_xyz, rgb], axis=-1) 339 | return stacked_features 340 | 341 | def init_input_pipeline(self): 342 | print('Initiating input pipelines') 343 | cfg.ignored_label_inds = [self.label_to_idx[ign_label] for ign_label in self.ignored_labels] 344 | gen_function, gen_types, gen_shapes = self.get_batch_gen('training') 345 | gen_function_val, _, _ = self.get_batch_gen('validation') 346 | gen_function_test, _, _ = self.get_batch_gen('test') 347 | self.train_data = tf.data.Dataset.from_generator(gen_function, gen_types, gen_shapes) 348 | self.val_data = tf.data.Dataset.from_generator(gen_function_val, gen_types, gen_shapes) 349 | self.test_data = tf.data.Dataset.from_generator(gen_function_test, gen_types, gen_shapes) 350 | 351 | self.batch_train_data = self.train_data.batch(cfg.batch_size) 352 | self.batch_val_data = self.val_data.batch(cfg.val_batch_size) 353 | self.batch_test_data = self.test_data.batch(cfg.val_batch_size) 354 | map_func = self.get_tf_mapping() 355 | 356 | self.batch_train_data = self.batch_train_data.map(map_func=map_func) 357 | self.batch_val_data = self.batch_val_data.map(map_func=map_func) 358 | self.batch_test_data = self.batch_test_data.map(map_func=map_func) 359 | 360 | self.batch_train_data = self.batch_train_data.prefetch(cfg.batch_size) 361 | self.batch_val_data = self.batch_val_data.prefetch(cfg.val_batch_size) 362 | self.batch_test_data = self.batch_test_data.prefetch(cfg.val_batch_size) 363 | 364 | iter = tf.data.Iterator.from_structure(self.batch_train_data.output_types, self.batch_train_data.output_shapes) 365 | self.flat_inputs = iter.get_next() 366 | self.train_init_op = iter.make_initializer(self.batch_train_data) 367 | self.val_init_op = iter.make_initializer(self.batch_val_data) 368 | self.test_init_op = iter.make_initializer(self.batch_test_data) 369 | 370 | 371 | if __name__ == '__main__': 372 | parser = argparse.ArgumentParser() 373 | parser.add_argument('--gpu', type=int, default=0, help='the number of GPUs to use [default: 0]') 374 | parser.add_argument('--mode', type=str, default='train', help='options: train, test, vis') 375 | parser.add_argument('--model_path', type=str, default='None', help='pretrained model path') 376 | FLAGS = parser.parse_args() 377 | 378 | GPU_ID = FLAGS.gpu 379 | os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" 380 | os.environ['CUDA_VISIBLE_DEVICES'] = str(GPU_ID) 381 | os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' 382 | 383 | Mode = FLAGS.mode 384 | dataset = Semantic3D() 385 | dataset.init_input_pipeline() 386 | 387 | if Mode == 'train': 388 | model = Network(dataset, cfg) 389 | model.train(dataset) 390 | elif Mode == 'test': 391 | cfg.saving = False 392 | model = Network(dataset, cfg) 393 | if FLAGS.model_path is not 'None': 394 | chosen_snap = FLAGS.model_path 395 | else: 396 | chosen_snapshot = -1 397 | logs = np.sort([os.path.join('results', f) for f in os.listdir('results') if f.startswith('Log')]) 398 | chosen_folder = logs[-1] 399 | snap_path = join(chosen_folder, 'snapshots') 400 | snap_steps = [int(f[:-5].split('-')[-1]) for f in os.listdir(snap_path) if f[-5:] == '.meta'] 401 | chosen_step = np.sort(snap_steps)[-1] 402 | chosen_snap = os.path.join(snap_path, 'snap-{:d}'.format(chosen_step)) 403 | tester = ModelTester(model, dataset, restore_snap=chosen_snap) 404 | tester.test(model, dataset) 405 | -------------------------------------------------------------------------------- /semantic3d_test.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import numpy as np 3 | import time 4 | from os import makedirs 5 | from os.path import exists, join 6 | 7 | from helper_ply import read_ply, write_ply 8 | 9 | 10 | def log_string(out_str, log_out): 11 | log_out.write(out_str + '\n') 12 | log_out.flush() 13 | print(out_str) 14 | 15 | 16 | class ModelTester: 17 | def __init__(self, model, dataset, restore_snap=None): 18 | # Tensorflow Saver definition 19 | my_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES) 20 | self.saver = tf.train.Saver(my_vars, max_to_keep=100) 21 | 22 | # Create a session for running Ops on the Graph. 23 | on_cpu = False 24 | if on_cpu: 25 | c_proto = tf.ConfigProto(device_count={'GPU': 0}) 26 | else: 27 | c_proto = tf.ConfigProto() 28 | c_proto.gpu_options.allow_growth = True 29 | self.sess = tf.Session(config=c_proto) 30 | self.sess.run(tf.global_variables_initializer()) 31 | 32 | if restore_snap is not None: 33 | self.saver.restore(self.sess, restore_snap) 34 | print("Model restored from " + restore_snap) 35 | 36 | # Add a softmax operation for predictions 37 | self.prob_logits = tf.nn.softmax(model.logits) 38 | self.test_probs = [np.zeros((l.data.shape[0], model.config.num_classes), dtype=np.float16) 39 | for l in dataset.input_trees['test']] 40 | 41 | self.log_out = open('log_test_' + dataset.name + '.txt', 'a') 42 | 43 | def test(self, model, dataset, num_votes=100): 44 | 45 | # Smoothing parameter for votes 46 | test_smooth = 0.98 47 | 48 | # Initialise iterator with train data 49 | self.sess.run(dataset.test_init_op) 50 | 51 | # Test saving path 52 | saving_path = time.strftime('results/Log_%Y-%m-%d_%H-%M-%S', time.gmtime()) 53 | test_path = join('test', saving_path.split('/')[-1]) 54 | makedirs(test_path) if not exists(test_path) else None 55 | makedirs(join(test_path, 'predictions')) if not exists(join(test_path, 'predictions')) else None 56 | makedirs(join(test_path, 'probs')) if not exists(join(test_path, 'probs')) else None 57 | 58 | ##################### 59 | # Network predictions 60 | ##################### 61 | 62 | step_id = 0 63 | epoch_id = 0 64 | last_min = -0.5 65 | 66 | t1 = time.time() 67 | 68 | while last_min < num_votes: 69 | 70 | try: 71 | ops = (self.prob_logits, 72 | model.labels, 73 | model.inputs['input_inds'], 74 | model.inputs['cloud_inds'],) 75 | 76 | stacked_probs, stacked_labels, point_idx, cloud_idx = self.sess.run(ops, {model.is_training: False}) 77 | stacked_probs = np.reshape(stacked_probs, [model.config.val_batch_size, model.config.num_points, 78 | model.config.num_classes]) 79 | 80 | for j in range(np.shape(stacked_probs)[0]): 81 | probs = stacked_probs[j, :, :] 82 | inds = point_idx[j, :] 83 | c_i = cloud_idx[j][0] 84 | self.test_probs[c_i][inds] = test_smooth * self.test_probs[c_i][inds] + (1 - test_smooth) * probs 85 | step_id += 1 86 | log_string('Epoch {:3d}, step {:3d}. min possibility = {:.1f}'.format(epoch_id, step_id, np.min( 87 | dataset.min_possibility['test'])), self.log_out) 88 | 89 | except tf.errors.OutOfRangeError: 90 | 91 | # Save predicted cloud 92 | new_min = np.min(dataset.min_possibility['test']) 93 | log_string('Epoch {:3d}, end. Min possibility = {:.1f}'.format(epoch_id, new_min), self.log_out) 94 | 95 | if last_min + 4 < new_min: 96 | t2 = time.time() 97 | print('Saving clouds') 98 | 99 | # Update last_min 100 | last_min = new_min 101 | 102 | # Project predictions 103 | print('\nReproject Vote #{:d}'.format(int(np.floor(new_min)))) 104 | 105 | files = dataset.test_files 106 | i_test = 0 107 | for i, file_path in enumerate(files): 108 | # Get file 109 | points = self.load_evaluation_points(file_path) 110 | points = points.astype(np.float16) 111 | 112 | # Reproject probs 113 | probs = np.zeros(shape=[np.shape(points)[0], 8], dtype=np.float16) 114 | proj_index = dataset.test_proj[i_test] 115 | 116 | probs = self.test_probs[i_test][proj_index, :] 117 | 118 | # Insert false columns for ignored labels 119 | probs2 = probs 120 | for l_ind, label_value in enumerate(dataset.label_values): 121 | if label_value in dataset.ignored_labels: 122 | probs2 = np.insert(probs2, l_ind, 0, axis=1) 123 | 124 | # Get the predicted labels 125 | preds = dataset.label_values[np.argmax(probs2, axis=1)].astype(np.uint8) 126 | 127 | # Save plys 128 | cloud_name = file_path.split('/')[-1] 129 | 130 | # Save ascii preds 131 | ascii_name = join(test_path, 'predictions', dataset.ascii_files[cloud_name]) 132 | np.savetxt(ascii_name, preds, fmt='%d') 133 | log_string(ascii_name + 'has saved', self.log_out) 134 | i_test += 1 135 | 136 | t3 = time.time() 137 | # print('Prediction: {:.1f} s Saving: {:.1f} s\n'.format(t2-t1, t3-t2)) 138 | self.sess.close() 139 | return 140 | 141 | self.sess.run(dataset.test_init_op) 142 | epoch_id += 1 143 | step_id = 0 144 | continue 145 | return 146 | 147 | @staticmethod 148 | def load_evaluation_points(file_path): 149 | data = read_ply(file_path) 150 | return np.vstack((data['x'], data['y'], data['z'])).T 151 | -------------------------------------------------------------------------------- /utils/cpp_wrappers/compile_wrappers.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # Compile cpp subsampling 4 | cd cpp_subsampling 5 | python3 setup.py build_ext --inplace 6 | cd .. 7 | 8 | -------------------------------------------------------------------------------- /utils/cpp_wrappers/cpp_subsampling/grid_subsampling/grid_subsampling.cpp: -------------------------------------------------------------------------------- 1 | 2 | #include "grid_subsampling.h" 3 | 4 | 5 | void grid_subsampling(vector& original_points, 6 | vector& subsampled_points, 7 | vector& original_features, 8 | vector& subsampled_features, 9 | vector& original_classes, 10 | vector& subsampled_classes, 11 | float sampleDl, 12 | int verbose) { 13 | 14 | // Initiate variables 15 | // ****************** 16 | 17 | // Number of points in the cloud 18 | size_t N = original_points.size(); 19 | 20 | // Dimension of the features 21 | size_t fdim = original_features.size() / N; 22 | size_t ldim = original_classes.size() / N; 23 | 24 | // Limits of the cloud 25 | PointXYZ minCorner = min_point(original_points); 26 | PointXYZ maxCorner = max_point(original_points); 27 | PointXYZ originCorner = floor(minCorner * (1/sampleDl)) * sampleDl; 28 | 29 | // Dimensions of the grid 30 | size_t sampleNX = (size_t)floor((maxCorner.x - originCorner.x) / sampleDl) + 1; 31 | size_t sampleNY = (size_t)floor((maxCorner.y - originCorner.y) / sampleDl) + 1; 32 | //size_t sampleNZ = (size_t)floor((maxCorner.z - originCorner.z) / sampleDl) + 1; 33 | 34 | // Check if features and classes need to be processed 35 | bool use_feature = original_features.size() > 0; 36 | bool use_classes = original_classes.size() > 0; 37 | 38 | 39 | // Create the sampled map 40 | // ********************** 41 | 42 | // Verbose parameters 43 | int i = 0; 44 | int nDisp = N / 100; 45 | 46 | // Initiate variables 47 | size_t iX, iY, iZ, mapIdx; 48 | unordered_map data; 49 | 50 | for (auto& p : original_points) 51 | { 52 | // Position of point in sample map 53 | iX = (size_t)floor((p.x - originCorner.x) / sampleDl); 54 | iY = (size_t)floor((p.y - originCorner.y) / sampleDl); 55 | iZ = (size_t)floor((p.z - originCorner.z) / sampleDl); 56 | mapIdx = iX + sampleNX*iY + sampleNX*sampleNY*iZ; 57 | 58 | // If not already created, create key 59 | if (data.count(mapIdx) < 1) 60 | data.emplace(mapIdx, SampledData(fdim, ldim)); 61 | 62 | // Fill the sample map 63 | if (use_feature && use_classes) 64 | data[mapIdx].update_all(p, original_features.begin() + i * fdim, original_classes.begin() + i * ldim); 65 | else if (use_feature) 66 | data[mapIdx].update_features(p, original_features.begin() + i * fdim); 67 | else if (use_classes) 68 | data[mapIdx].update_classes(p, original_classes.begin() + i * ldim); 69 | else 70 | data[mapIdx].update_points(p); 71 | 72 | // Display 73 | i++; 74 | if (verbose > 1 && i%nDisp == 0) 75 | std::cout << "\rSampled Map : " << std::setw(3) << i / nDisp << "%"; 76 | 77 | } 78 | 79 | // Divide for barycentre and transfer to a vector 80 | subsampled_points.reserve(data.size()); 81 | if (use_feature) 82 | subsampled_features.reserve(data.size() * fdim); 83 | if (use_classes) 84 | subsampled_classes.reserve(data.size() * ldim); 85 | for (auto& v : data) 86 | { 87 | subsampled_points.push_back(v.second.point * (1.0 / v.second.count)); 88 | if (use_feature) 89 | { 90 | float count = (float)v.second.count; 91 | transform(v.second.features.begin(), 92 | v.second.features.end(), 93 | v.second.features.begin(), 94 | [count](float f) { return f / count;}); 95 | subsampled_features.insert(subsampled_features.end(),v.second.features.begin(),v.second.features.end()); 96 | } 97 | if (use_classes) 98 | { 99 | for (int i = 0; i < ldim; i++) 100 | subsampled_classes.push_back(max_element(v.second.labels[i].begin(), v.second.labels[i].end(), 101 | [](const pair&a, const pair&b){return a.second < b.second;})->first); 102 | } 103 | } 104 | 105 | return; 106 | } 107 | -------------------------------------------------------------------------------- /utils/cpp_wrappers/cpp_subsampling/grid_subsampling/grid_subsampling.h: -------------------------------------------------------------------------------- 1 | 2 | 3 | #include "../../cpp_utils/cloud/cloud.h" 4 | 5 | #include 6 | #include 7 | 8 | using namespace std; 9 | 10 | class SampledData 11 | { 12 | public: 13 | 14 | // Elements 15 | // ******** 16 | 17 | int count; 18 | PointXYZ point; 19 | vector features; 20 | vector> labels; 21 | 22 | 23 | // Methods 24 | // ******* 25 | 26 | // Constructor 27 | SampledData() 28 | { 29 | count = 0; 30 | point = PointXYZ(); 31 | } 32 | 33 | SampledData(const size_t fdim, const size_t ldim) 34 | { 35 | count = 0; 36 | point = PointXYZ(); 37 | features = vector(fdim); 38 | labels = vector>(ldim); 39 | } 40 | 41 | // Method Update 42 | void update_all(const PointXYZ p, vector::iterator f_begin, vector::iterator l_begin) 43 | { 44 | count += 1; 45 | point += p; 46 | transform (features.begin(), features.end(), f_begin, features.begin(), plus()); 47 | int i = 0; 48 | for(vector::iterator it = l_begin; it != l_begin + labels.size(); ++it) 49 | { 50 | labels[i][*it] += 1; 51 | i++; 52 | } 53 | return; 54 | } 55 | void update_features(const PointXYZ p, vector::iterator f_begin) 56 | { 57 | count += 1; 58 | point += p; 59 | transform (features.begin(), features.end(), f_begin, features.begin(), plus()); 60 | return; 61 | } 62 | void update_classes(const PointXYZ p, vector::iterator l_begin) 63 | { 64 | count += 1; 65 | point += p; 66 | int i = 0; 67 | for(vector::iterator it = l_begin; it != l_begin + labels.size(); ++it) 68 | { 69 | labels[i][*it] += 1; 70 | i++; 71 | } 72 | return; 73 | } 74 | void update_points(const PointXYZ p) 75 | { 76 | count += 1; 77 | point += p; 78 | return; 79 | } 80 | }; 81 | 82 | 83 | 84 | void grid_subsampling(vector& original_points, 85 | vector& subsampled_points, 86 | vector& original_features, 87 | vector& subsampled_features, 88 | vector& original_classes, 89 | vector& subsampled_classes, 90 | float sampleDl, 91 | int verbose); 92 | 93 | -------------------------------------------------------------------------------- /utils/cpp_wrappers/cpp_subsampling/setup.py: -------------------------------------------------------------------------------- 1 | from distutils.core import setup, Extension 2 | import numpy.distutils.misc_util 3 | 4 | # Adding OpenCV to project 5 | # ************************ 6 | 7 | # Adding sources of the project 8 | # ***************************** 9 | 10 | m_name = "grid_subsampling" 11 | 12 | SOURCES = ["../cpp_utils/cloud/cloud.cpp", 13 | "grid_subsampling/grid_subsampling.cpp", 14 | "wrapper.cpp"] 15 | 16 | module = Extension(m_name, 17 | sources=SOURCES, 18 | extra_compile_args=['-std=c++11', 19 | '-D_GLIBCXX_USE_CXX11_ABI=0']) 20 | 21 | setup(ext_modules=[module], include_dirs=numpy.distutils.misc_util.get_numpy_include_dirs()) 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | -------------------------------------------------------------------------------- /utils/cpp_wrappers/cpp_subsampling/wrapper.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include "grid_subsampling/grid_subsampling.h" 4 | #include 5 | 6 | 7 | 8 | // docstrings for our module 9 | // ************************* 10 | 11 | static char module_docstring[] = "This module provides an interface for the subsampling of a pointcloud"; 12 | 13 | static char compute_docstring[] = "function subsampling a pointcloud"; 14 | 15 | 16 | // Declare the functions 17 | // ********************* 18 | 19 | static PyObject *grid_subsampling_compute(PyObject *self, PyObject *args, PyObject *keywds); 20 | 21 | 22 | // Specify the members of the module 23 | // ********************************* 24 | 25 | static PyMethodDef module_methods[] = 26 | { 27 | { "compute", (PyCFunction)grid_subsampling_compute, METH_VARARGS | METH_KEYWORDS, compute_docstring }, 28 | {NULL, NULL, 0, NULL} 29 | }; 30 | 31 | 32 | // Initialize the module 33 | // ********************* 34 | 35 | static struct PyModuleDef moduledef = 36 | { 37 | PyModuleDef_HEAD_INIT, 38 | "grid_subsampling", // m_name 39 | module_docstring, // m_doc 40 | -1, // m_size 41 | module_methods, // m_methods 42 | NULL, // m_reload 43 | NULL, // m_traverse 44 | NULL, // m_clear 45 | NULL, // m_free 46 | }; 47 | 48 | PyMODINIT_FUNC PyInit_grid_subsampling(void) 49 | { 50 | import_array(); 51 | return PyModule_Create(&moduledef); 52 | } 53 | 54 | 55 | // Actual wrapper 56 | // ************** 57 | 58 | static PyObject *grid_subsampling_compute(PyObject *self, PyObject *args, PyObject *keywds) 59 | { 60 | 61 | // Manage inputs 62 | // ************* 63 | 64 | // Args containers 65 | PyObject *points_obj = NULL; 66 | PyObject *features_obj = NULL; 67 | PyObject *classes_obj = NULL; 68 | 69 | // Keywords containers 70 | static char *kwlist[] = {"points", "features", "classes", "sampleDl", "method", "verbose", NULL }; 71 | float sampleDl = 0.1; 72 | const char *method_buffer = "barycenters"; 73 | int verbose = 0; 74 | 75 | // Parse the input 76 | if (!PyArg_ParseTupleAndKeywords(args, keywds, "O|$OOfsi", kwlist, &points_obj, &features_obj, &classes_obj, &sampleDl, &method_buffer, &verbose)) 77 | { 78 | PyErr_SetString(PyExc_RuntimeError, "Error parsing arguments"); 79 | return NULL; 80 | } 81 | 82 | // Get the method argument 83 | string method(method_buffer); 84 | 85 | // Interpret method 86 | if (method.compare("barycenters") && method.compare("voxelcenters")) 87 | { 88 | PyErr_SetString(PyExc_RuntimeError, "Error parsing method. Valid method names are \"barycenters\" and \"voxelcenters\" "); 89 | return NULL; 90 | } 91 | 92 | // Check if using features or classes 93 | bool use_feature = true, use_classes = true; 94 | if (features_obj == NULL) 95 | use_feature = false; 96 | if (classes_obj == NULL) 97 | use_classes = false; 98 | 99 | // Interpret the input objects as numpy arrays. 100 | PyObject *points_array = PyArray_FROM_OTF(points_obj, NPY_FLOAT, NPY_IN_ARRAY); 101 | PyObject *features_array = NULL; 102 | PyObject *classes_array = NULL; 103 | if (use_feature) 104 | features_array = PyArray_FROM_OTF(features_obj, NPY_FLOAT, NPY_IN_ARRAY); 105 | if (use_classes) 106 | classes_array = PyArray_FROM_OTF(classes_obj, NPY_INT, NPY_IN_ARRAY); 107 | 108 | // Verify data was load correctly. 109 | if (points_array == NULL) 110 | { 111 | Py_XDECREF(points_array); 112 | Py_XDECREF(classes_array); 113 | Py_XDECREF(features_array); 114 | PyErr_SetString(PyExc_RuntimeError, "Error converting input points to numpy arrays of type float32"); 115 | return NULL; 116 | } 117 | if (use_feature && features_array == NULL) 118 | { 119 | Py_XDECREF(points_array); 120 | Py_XDECREF(classes_array); 121 | Py_XDECREF(features_array); 122 | PyErr_SetString(PyExc_RuntimeError, "Error converting input features to numpy arrays of type float32"); 123 | return NULL; 124 | } 125 | if (use_classes && classes_array == NULL) 126 | { 127 | Py_XDECREF(points_array); 128 | Py_XDECREF(classes_array); 129 | Py_XDECREF(features_array); 130 | PyErr_SetString(PyExc_RuntimeError, "Error converting input classes to numpy arrays of type int32"); 131 | return NULL; 132 | } 133 | 134 | // Check that the input array respect the dims 135 | if ((int)PyArray_NDIM(points_array) != 2 || (int)PyArray_DIM(points_array, 1) != 3) 136 | { 137 | Py_XDECREF(points_array); 138 | Py_XDECREF(classes_array); 139 | Py_XDECREF(features_array); 140 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : points.shape is not (N, 3)"); 141 | return NULL; 142 | } 143 | if (use_feature && ((int)PyArray_NDIM(features_array) != 2)) 144 | { 145 | Py_XDECREF(points_array); 146 | Py_XDECREF(classes_array); 147 | Py_XDECREF(features_array); 148 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : features.shape is not (N, d)"); 149 | return NULL; 150 | } 151 | 152 | if (use_classes && (int)PyArray_NDIM(classes_array) > 2) 153 | { 154 | Py_XDECREF(points_array); 155 | Py_XDECREF(classes_array); 156 | Py_XDECREF(features_array); 157 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : classes.shape is not (N,) or (N, d)"); 158 | return NULL; 159 | } 160 | 161 | // Number of points 162 | int N = (int)PyArray_DIM(points_array, 0); 163 | 164 | // Dimension of the features 165 | int fdim = 0; 166 | if (use_feature) 167 | fdim = (int)PyArray_DIM(features_array, 1); 168 | 169 | //Dimension of labels 170 | int ldim = 1; 171 | if (use_classes && (int)PyArray_NDIM(classes_array) == 2) 172 | ldim = (int)PyArray_DIM(classes_array, 1); 173 | 174 | // Check that the input array respect the number of points 175 | if (use_feature && (int)PyArray_DIM(features_array, 0) != N) 176 | { 177 | Py_XDECREF(points_array); 178 | Py_XDECREF(classes_array); 179 | Py_XDECREF(features_array); 180 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : features.shape is not (N, d)"); 181 | return NULL; 182 | } 183 | if (use_classes && (int)PyArray_DIM(classes_array, 0) != N) 184 | { 185 | Py_XDECREF(points_array); 186 | Py_XDECREF(classes_array); 187 | Py_XDECREF(features_array); 188 | PyErr_SetString(PyExc_RuntimeError, "Wrong dimensions : classes.shape is not (N,) or (N, d)"); 189 | return NULL; 190 | } 191 | 192 | 193 | // Call the C++ function 194 | // ********************* 195 | 196 | // Create pyramid 197 | if (verbose > 0) 198 | cout << "Computing cloud pyramid with support points: " << endl; 199 | 200 | 201 | // Convert PyArray to Cloud C++ class 202 | vector original_points; 203 | vector original_features; 204 | vector original_classes; 205 | original_points = vector((PointXYZ*)PyArray_DATA(points_array), (PointXYZ*)PyArray_DATA(points_array) + N); 206 | if (use_feature) 207 | original_features = vector((float*)PyArray_DATA(features_array), (float*)PyArray_DATA(features_array) + N*fdim); 208 | if (use_classes) 209 | original_classes = vector((int*)PyArray_DATA(classes_array), (int*)PyArray_DATA(classes_array) + N*ldim); 210 | 211 | // Subsample 212 | vector subsampled_points; 213 | vector subsampled_features; 214 | vector subsampled_classes; 215 | grid_subsampling(original_points, 216 | subsampled_points, 217 | original_features, 218 | subsampled_features, 219 | original_classes, 220 | subsampled_classes, 221 | sampleDl, 222 | verbose); 223 | 224 | // Check result 225 | if (subsampled_points.size() < 1) 226 | { 227 | PyErr_SetString(PyExc_RuntimeError, "Error"); 228 | return NULL; 229 | } 230 | 231 | // Manage outputs 232 | // ************** 233 | 234 | // Dimension of input containers 235 | npy_intp* point_dims = new npy_intp[2]; 236 | point_dims[0] = subsampled_points.size(); 237 | point_dims[1] = 3; 238 | npy_intp* feature_dims = new npy_intp[2]; 239 | feature_dims[0] = subsampled_points.size(); 240 | feature_dims[1] = fdim; 241 | npy_intp* classes_dims = new npy_intp[2]; 242 | classes_dims[0] = subsampled_points.size(); 243 | classes_dims[1] = ldim; 244 | 245 | // Create output array 246 | PyObject *res_points_obj = PyArray_SimpleNew(2, point_dims, NPY_FLOAT); 247 | PyObject *res_features_obj = NULL; 248 | PyObject *res_classes_obj = NULL; 249 | PyObject *ret = NULL; 250 | 251 | // Fill output array with values 252 | size_t size_in_bytes = subsampled_points.size() * 3 * sizeof(float); 253 | memcpy(PyArray_DATA(res_points_obj), subsampled_points.data(), size_in_bytes); 254 | if (use_feature) 255 | { 256 | size_in_bytes = subsampled_points.size() * fdim * sizeof(float); 257 | res_features_obj = PyArray_SimpleNew(2, feature_dims, NPY_FLOAT); 258 | memcpy(PyArray_DATA(res_features_obj), subsampled_features.data(), size_in_bytes); 259 | } 260 | if (use_classes) 261 | { 262 | size_in_bytes = subsampled_points.size() * ldim * sizeof(int); 263 | res_classes_obj = PyArray_SimpleNew(2, classes_dims, NPY_INT); 264 | memcpy(PyArray_DATA(res_classes_obj), subsampled_classes.data(), size_in_bytes); 265 | } 266 | 267 | 268 | // Merge results 269 | if (use_feature && use_classes) 270 | ret = Py_BuildValue("NNN", res_points_obj, res_features_obj, res_classes_obj); 271 | else if (use_feature) 272 | ret = Py_BuildValue("NN", res_points_obj, res_features_obj); 273 | else if (use_classes) 274 | ret = Py_BuildValue("NN", res_points_obj, res_classes_obj); 275 | else 276 | ret = Py_BuildValue("N", res_points_obj); 277 | 278 | // Clean up 279 | // ******** 280 | 281 | Py_DECREF(points_array); 282 | Py_XDECREF(features_array); 283 | Py_XDECREF(classes_array); 284 | 285 | return ret; 286 | } -------------------------------------------------------------------------------- /utils/cpp_wrappers/cpp_utils/cloud/cloud.cpp: -------------------------------------------------------------------------------- 1 | // 2 | // 3 | // 0==========================0 4 | // | Local feature test | 5 | // 0==========================0 6 | // 7 | // version 1.0 : 8 | // > 9 | // 10 | //--------------------------------------------------- 11 | // 12 | // Cloud source : 13 | // Define usefull Functions/Methods 14 | // 15 | //---------------------------------------------------- 16 | // 17 | // Hugues THOMAS - 10/02/2017 18 | // 19 | 20 | 21 | #include "cloud.h" 22 | 23 | 24 | // Getters 25 | // ******* 26 | 27 | PointXYZ max_point(std::vector points) 28 | { 29 | // Initiate limits 30 | PointXYZ maxP(points[0]); 31 | 32 | // Loop over all points 33 | for (auto p : points) 34 | { 35 | if (p.x > maxP.x) 36 | maxP.x = p.x; 37 | 38 | if (p.y > maxP.y) 39 | maxP.y = p.y; 40 | 41 | if (p.z > maxP.z) 42 | maxP.z = p.z; 43 | } 44 | 45 | return maxP; 46 | } 47 | 48 | PointXYZ min_point(std::vector points) 49 | { 50 | // Initiate limits 51 | PointXYZ minP(points[0]); 52 | 53 | // Loop over all points 54 | for (auto p : points) 55 | { 56 | if (p.x < minP.x) 57 | minP.x = p.x; 58 | 59 | if (p.y < minP.y) 60 | minP.y = p.y; 61 | 62 | if (p.z < minP.z) 63 | minP.z = p.z; 64 | } 65 | 66 | return minP; 67 | } -------------------------------------------------------------------------------- /utils/cpp_wrappers/cpp_utils/cloud/cloud.h: -------------------------------------------------------------------------------- 1 | // 2 | // 3 | // 0==========================0 4 | // | Local feature test | 5 | // 0==========================0 6 | // 7 | // version 1.0 : 8 | // > 9 | // 10 | //--------------------------------------------------- 11 | // 12 | // Cloud header 13 | // 14 | //---------------------------------------------------- 15 | // 16 | // Hugues THOMAS - 10/02/2017 17 | // 18 | 19 | 20 | # pragma once 21 | 22 | #include 23 | #include 24 | #include 25 | #include 26 | #include 27 | #include 28 | #include 29 | #include 30 | 31 | #include 32 | 33 | 34 | 35 | 36 | // Point class 37 | // *********** 38 | 39 | 40 | class PointXYZ 41 | { 42 | public: 43 | 44 | // Elements 45 | // ******** 46 | 47 | float x, y, z; 48 | 49 | 50 | // Methods 51 | // ******* 52 | 53 | // Constructor 54 | PointXYZ() { x = 0; y = 0; z = 0; } 55 | PointXYZ(float x0, float y0, float z0) { x = x0; y = y0; z = z0; } 56 | 57 | // array type accessor 58 | float operator [] (int i) const 59 | { 60 | if (i == 0) return x; 61 | else if (i == 1) return y; 62 | else return z; 63 | } 64 | 65 | // opperations 66 | float dot(const PointXYZ P) const 67 | { 68 | return x * P.x + y * P.y + z * P.z; 69 | } 70 | 71 | float sq_norm() 72 | { 73 | return x*x + y*y + z*z; 74 | } 75 | 76 | PointXYZ cross(const PointXYZ P) const 77 | { 78 | return PointXYZ(y*P.z - z*P.y, z*P.x - x*P.z, x*P.y - y*P.x); 79 | } 80 | 81 | PointXYZ& operator+=(const PointXYZ& P) 82 | { 83 | x += P.x; 84 | y += P.y; 85 | z += P.z; 86 | return *this; 87 | } 88 | 89 | PointXYZ& operator-=(const PointXYZ& P) 90 | { 91 | x -= P.x; 92 | y -= P.y; 93 | z -= P.z; 94 | return *this; 95 | } 96 | 97 | PointXYZ& operator*=(const float& a) 98 | { 99 | x *= a; 100 | y *= a; 101 | z *= a; 102 | return *this; 103 | } 104 | }; 105 | 106 | 107 | // Point Opperations 108 | // ***************** 109 | 110 | inline PointXYZ operator + (const PointXYZ A, const PointXYZ B) 111 | { 112 | return PointXYZ(A.x + B.x, A.y + B.y, A.z + B.z); 113 | } 114 | 115 | inline PointXYZ operator - (const PointXYZ A, const PointXYZ B) 116 | { 117 | return PointXYZ(A.x - B.x, A.y - B.y, A.z - B.z); 118 | } 119 | 120 | inline PointXYZ operator * (const PointXYZ P, const float a) 121 | { 122 | return PointXYZ(P.x * a, P.y * a, P.z * a); 123 | } 124 | 125 | inline PointXYZ operator * (const float a, const PointXYZ P) 126 | { 127 | return PointXYZ(P.x * a, P.y * a, P.z * a); 128 | } 129 | 130 | inline std::ostream& operator << (std::ostream& os, const PointXYZ P) 131 | { 132 | return os << "[" << P.x << ", " << P.y << ", " << P.z << "]"; 133 | } 134 | 135 | inline bool operator == (const PointXYZ A, const PointXYZ B) 136 | { 137 | return A.x == B.x && A.y == B.y && A.z == B.z; 138 | } 139 | 140 | inline PointXYZ floor(const PointXYZ P) 141 | { 142 | return PointXYZ(std::floor(P.x), std::floor(P.y), std::floor(P.z)); 143 | } 144 | 145 | 146 | PointXYZ max_point(std::vector points); 147 | PointXYZ min_point(std::vector points); 148 | 149 | 150 | 151 | 152 | 153 | 154 | 155 | 156 | 157 | 158 | 159 | -------------------------------------------------------------------------------- /utils/download_semantic3d.sh: -------------------------------------------------------------------------------- 1 | BASE_DIR=${1-./data/semantic3d/original_data} 2 | 3 | # Training data 4 | wget -c -N http://semantic3d.net/data/point-clouds/training1/bildstein_station1_xyz_intensity_rgb.7z -P $BASE_DIR 5 | wget -c -N http://semantic3d.net/data/point-clouds/training1/bildstein_station3_xyz_intensity_rgb.7z -P $BASE_DIR 6 | wget -c -N http://semantic3d.net/data/point-clouds/training1/bildstein_station5_xyz_intensity_rgb.7z -P $BASE_DIR 7 | wget -c -N http://semantic3d.net/data/point-clouds/training1/domfountain_station1_xyz_intensity_rgb.7z -P $BASE_DIR 8 | wget -c -N http://semantic3d.net/data/point-clouds/training1/domfountain_station2_xyz_intensity_rgb.7z -P $BASE_DIR/ 9 | wget -c -N http://semantic3d.net/data/point-clouds/training1/domfountain_station3_xyz_intensity_rgb.7z -P $BASE_DIR 10 | wget -c -N http://semantic3d.net/data/point-clouds/training1/neugasse_station1_xyz_intensity_rgb.7z -P $BASE_DIR 11 | wget -c -N http://semantic3d.net/data/point-clouds/training1/sg27_station1_intensity_rgb.7z -P $BASE_DIR 12 | wget -c -N http://semantic3d.net/data/point-clouds/training1/sg27_station2_intensity_rgb.7z -P $BASE_DIR 13 | wget -c -N http://semantic3d.net/data/point-clouds/training1/sg27_station4_intensity_rgb.7z -P $BASE_DIR/ 14 | wget -c -N http://semantic3d.net/data/point-clouds/training1/sg27_station5_intensity_rgb.7z -P $BASE_DIR 15 | wget -c -N http://semantic3d.net/data/point-clouds/training1/sg27_station9_intensity_rgb.7z -P $BASE_DIR 16 | wget -c -N http://semantic3d.net/data/point-clouds/training1/sg28_station4_intensity_rgb.7z -P $BASE_DIR 17 | wget -c -N http://semantic3d.net/data/point-clouds/training1/untermaederbrunnen_station1_xyz_intensity_rgb.7z -P $BASE_DIR 18 | wget -c -N http://semantic3d.net/data/point-clouds/training1/untermaederbrunnen_station3_xyz_intensity_rgb.7z -P $BASE_DIR/ 19 | wget -c -N http://semantic3d.net/data/sem8_labels_training.7z -P $BASE_DIR 20 | 21 | 22 | # Test data 23 | wget -c -N http://semantic3d.net/data/point-clouds/testing1/birdfountain_station1_xyz_intensity_rgb.7z -P $BASE_DIR 24 | wget -c -N http://semantic3d.net/data/point-clouds/testing1/castleblatten_station1_intensity_rgb.7z -P $BASE_DIR 25 | wget -c -N http://semantic3d.net/data/point-clouds/testing1/castleblatten_station5_xyz_intensity_rgb.7z -P $BASE_DIR 26 | wget -c -N http://semantic3d.net/data/point-clouds/testing1/marketplacefeldkirch_station1_intensity_rgb.7z -P $BASE_DIR 27 | wget -c -N http://semantic3d.net/data/point-clouds/testing1/marketplacefeldkirch_station4_intensity_rgb.7z -P $BASE_DIR 28 | wget -c -N http://semantic3d.net/data/point-clouds/testing1/marketplacefeldkirch_station7_intensity_rgb.7z -P $BASE_DIR 29 | wget -c -N http://semantic3d.net/data/point-clouds/testing1/sg27_station10_intensity_rgb.7z -P $BASE_DIR 30 | wget -c -N http://semantic3d.net/data/point-clouds/testing1/sg27_station3_intensity_rgb.7z -P $BASE_DIR 31 | wget -c -N http://semantic3d.net/data/point-clouds/testing1/sg27_station6_intensity_rgb.7z -P $BASE_DIR 32 | wget -c -N http://semantic3d.net/data/point-clouds/testing1/sg27_station8_intensity_rgb.7z -P $BASE_DIR 33 | wget -c -N http://semantic3d.net/data/point-clouds/testing1/sg28_station2_intensity_rgb.7z -P $BASE_DIR 34 | wget -c -N http://semantic3d.net/data/point-clouds/testing1/sg28_station5_xyz_intensity_rgb.7z -P $BASE_DIR 35 | wget -c -N http://semantic3d.net/data/point-clouds/testing1/stgallencathedral_station1_intensity_rgb.7z -P $BASE_DIR 36 | wget -c -N http://semantic3d.net/data/point-clouds/testing1/stgallencathedral_station3_intensity_rgb.7z -P $BASE_DIR 37 | wget -c -N http://semantic3d.net/data/point-clouds/testing1/stgallencathedral_station6_intensity_rgb.7z -P $BASE_DIR 38 | 39 | # reduced-8 40 | wget -c -N http://semantic3d.net/data/point-clouds/testing2/MarketplaceFeldkirch_Station4_rgb_intensity-reduced.txt.7z -P $BASE_DIR 41 | wget -c -N http://semantic3d.net/data/point-clouds/testing2/StGallenCathedral_station6_rgb_intensity-reduced.txt.7z -P $BASE_DIR 42 | wget -c -N http://semantic3d.net/data/point-clouds/testing2/sg27_station10_rgb_intensity-reduced.txt.7z -P $BASE_DIR 43 | wget -c -N http://semantic3d.net/data/point-clouds/testing2/sg28_Station2_rgb_intensity-reduced.txt.7z -P $BASE_DIR 44 | 45 | 46 | 47 | for entry in "$BASE_DIR"/* 48 | do 49 | 7z x "$entry" -o$(dirname "$entry") -y 50 | done 51 | 52 | mv $BASE_DIR/station1_xyz_intensity_rgb.txt $BASE_DIR/#neugasse_station1_xyz_intensity_rgb.txt 53 | 54 | for entry in "$BASE_DIR"/*.7z 55 | do 56 | rm "$entry" 57 | done 58 | -------------------------------------------------------------------------------- /utils/meta/anno_paths.txt: -------------------------------------------------------------------------------- 1 | Area_1/conferenceRoom_1/Annotations 2 | Area_1/conferenceRoom_2/Annotations 3 | Area_1/copyRoom_1/Annotations 4 | Area_1/hallway_1/Annotations 5 | Area_1/hallway_2/Annotations 6 | Area_1/hallway_3/Annotations 7 | Area_1/hallway_4/Annotations 8 | Area_1/hallway_5/Annotations 9 | Area_1/hallway_6/Annotations 10 | Area_1/hallway_7/Annotations 11 | Area_1/hallway_8/Annotations 12 | Area_1/office_10/Annotations 13 | Area_1/office_11/Annotations 14 | Area_1/office_12/Annotations 15 | Area_1/office_13/Annotations 16 | Area_1/office_14/Annotations 17 | Area_1/office_15/Annotations 18 | Area_1/office_16/Annotations 19 | Area_1/office_17/Annotations 20 | Area_1/office_18/Annotations 21 | Area_1/office_19/Annotations 22 | Area_1/office_1/Annotations 23 | Area_1/office_20/Annotations 24 | Area_1/office_21/Annotations 25 | Area_1/office_22/Annotations 26 | Area_1/office_23/Annotations 27 | Area_1/office_24/Annotations 28 | Area_1/office_25/Annotations 29 | Area_1/office_26/Annotations 30 | Area_1/office_27/Annotations 31 | Area_1/office_28/Annotations 32 | Area_1/office_29/Annotations 33 | Area_1/office_2/Annotations 34 | Area_1/office_30/Annotations 35 | Area_1/office_31/Annotations 36 | Area_1/office_3/Annotations 37 | Area_1/office_4/Annotations 38 | Area_1/office_5/Annotations 39 | Area_1/office_6/Annotations 40 | Area_1/office_7/Annotations 41 | Area_1/office_8/Annotations 42 | Area_1/office_9/Annotations 43 | Area_1/pantry_1/Annotations 44 | Area_1/WC_1/Annotations 45 | Area_2/auditorium_1/Annotations 46 | Area_2/auditorium_2/Annotations 47 | Area_2/conferenceRoom_1/Annotations 48 | Area_2/hallway_10/Annotations 49 | Area_2/hallway_11/Annotations 50 | Area_2/hallway_12/Annotations 51 | Area_2/hallway_1/Annotations 52 | Area_2/hallway_2/Annotations 53 | Area_2/hallway_3/Annotations 54 | Area_2/hallway_4/Annotations 55 | Area_2/hallway_5/Annotations 56 | Area_2/hallway_6/Annotations 57 | Area_2/hallway_7/Annotations 58 | Area_2/hallway_8/Annotations 59 | Area_2/hallway_9/Annotations 60 | Area_2/office_10/Annotations 61 | Area_2/office_11/Annotations 62 | Area_2/office_12/Annotations 63 | Area_2/office_13/Annotations 64 | Area_2/office_14/Annotations 65 | Area_2/office_1/Annotations 66 | Area_2/office_2/Annotations 67 | Area_2/office_3/Annotations 68 | Area_2/office_4/Annotations 69 | Area_2/office_5/Annotations 70 | Area_2/office_6/Annotations 71 | Area_2/office_7/Annotations 72 | Area_2/office_8/Annotations 73 | Area_2/office_9/Annotations 74 | Area_2/storage_1/Annotations 75 | Area_2/storage_2/Annotations 76 | Area_2/storage_3/Annotations 77 | Area_2/storage_4/Annotations 78 | Area_2/storage_5/Annotations 79 | Area_2/storage_6/Annotations 80 | Area_2/storage_7/Annotations 81 | Area_2/storage_8/Annotations 82 | Area_2/storage_9/Annotations 83 | Area_2/WC_1/Annotations 84 | Area_2/WC_2/Annotations 85 | Area_3/conferenceRoom_1/Annotations 86 | Area_3/hallway_1/Annotations 87 | Area_3/hallway_2/Annotations 88 | Area_3/hallway_3/Annotations 89 | Area_3/hallway_4/Annotations 90 | Area_3/hallway_5/Annotations 91 | Area_3/hallway_6/Annotations 92 | Area_3/lounge_1/Annotations 93 | Area_3/lounge_2/Annotations 94 | Area_3/office_10/Annotations 95 | Area_3/office_1/Annotations 96 | Area_3/office_2/Annotations 97 | Area_3/office_3/Annotations 98 | Area_3/office_4/Annotations 99 | Area_3/office_5/Annotations 100 | Area_3/office_6/Annotations 101 | Area_3/office_7/Annotations 102 | Area_3/office_8/Annotations 103 | Area_3/office_9/Annotations 104 | Area_3/storage_1/Annotations 105 | Area_3/storage_2/Annotations 106 | Area_3/WC_1/Annotations 107 | Area_3/WC_2/Annotations 108 | Area_4/conferenceRoom_1/Annotations 109 | Area_4/conferenceRoom_2/Annotations 110 | Area_4/conferenceRoom_3/Annotations 111 | Area_4/hallway_10/Annotations 112 | Area_4/hallway_11/Annotations 113 | Area_4/hallway_12/Annotations 114 | Area_4/hallway_13/Annotations 115 | Area_4/hallway_14/Annotations 116 | Area_4/hallway_1/Annotations 117 | Area_4/hallway_2/Annotations 118 | Area_4/hallway_3/Annotations 119 | Area_4/hallway_4/Annotations 120 | Area_4/hallway_5/Annotations 121 | Area_4/hallway_6/Annotations 122 | Area_4/hallway_7/Annotations 123 | Area_4/hallway_8/Annotations 124 | Area_4/hallway_9/Annotations 125 | Area_4/lobby_1/Annotations 126 | Area_4/lobby_2/Annotations 127 | Area_4/office_10/Annotations 128 | Area_4/office_11/Annotations 129 | Area_4/office_12/Annotations 130 | Area_4/office_13/Annotations 131 | Area_4/office_14/Annotations 132 | Area_4/office_15/Annotations 133 | Area_4/office_16/Annotations 134 | Area_4/office_17/Annotations 135 | Area_4/office_18/Annotations 136 | Area_4/office_19/Annotations 137 | Area_4/office_1/Annotations 138 | Area_4/office_20/Annotations 139 | Area_4/office_21/Annotations 140 | Area_4/office_22/Annotations 141 | Area_4/office_2/Annotations 142 | Area_4/office_3/Annotations 143 | Area_4/office_4/Annotations 144 | Area_4/office_5/Annotations 145 | Area_4/office_6/Annotations 146 | Area_4/office_7/Annotations 147 | Area_4/office_8/Annotations 148 | Area_4/office_9/Annotations 149 | Area_4/storage_1/Annotations 150 | Area_4/storage_2/Annotations 151 | Area_4/storage_3/Annotations 152 | Area_4/storage_4/Annotations 153 | Area_4/WC_1/Annotations 154 | Area_4/WC_2/Annotations 155 | Area_4/WC_3/Annotations 156 | Area_4/WC_4/Annotations 157 | Area_5/conferenceRoom_1/Annotations 158 | Area_5/conferenceRoom_2/Annotations 159 | Area_5/conferenceRoom_3/Annotations 160 | Area_5/hallway_10/Annotations 161 | Area_5/hallway_11/Annotations 162 | Area_5/hallway_12/Annotations 163 | Area_5/hallway_13/Annotations 164 | Area_5/hallway_14/Annotations 165 | Area_5/hallway_15/Annotations 166 | Area_5/hallway_1/Annotations 167 | Area_5/hallway_2/Annotations 168 | Area_5/hallway_3/Annotations 169 | Area_5/hallway_4/Annotations 170 | Area_5/hallway_5/Annotations 171 | Area_5/hallway_6/Annotations 172 | Area_5/hallway_7/Annotations 173 | Area_5/hallway_8/Annotations 174 | Area_5/hallway_9/Annotations 175 | Area_5/lobby_1/Annotations 176 | Area_5/office_10/Annotations 177 | Area_5/office_11/Annotations 178 | Area_5/office_12/Annotations 179 | Area_5/office_13/Annotations 180 | Area_5/office_14/Annotations 181 | Area_5/office_15/Annotations 182 | Area_5/office_16/Annotations 183 | Area_5/office_17/Annotations 184 | Area_5/office_18/Annotations 185 | Area_5/office_19/Annotations 186 | Area_5/office_1/Annotations 187 | Area_5/office_20/Annotations 188 | Area_5/office_21/Annotations 189 | Area_5/office_22/Annotations 190 | Area_5/office_23/Annotations 191 | Area_5/office_24/Annotations 192 | Area_5/office_25/Annotations 193 | Area_5/office_26/Annotations 194 | Area_5/office_27/Annotations 195 | Area_5/office_28/Annotations 196 | Area_5/office_29/Annotations 197 | Area_5/office_2/Annotations 198 | Area_5/office_30/Annotations 199 | Area_5/office_31/Annotations 200 | Area_5/office_32/Annotations 201 | Area_5/office_33/Annotations 202 | Area_5/office_34/Annotations 203 | Area_5/office_35/Annotations 204 | Area_5/office_36/Annotations 205 | Area_5/office_37/Annotations 206 | Area_5/office_38/Annotations 207 | Area_5/office_39/Annotations 208 | Area_5/office_3/Annotations 209 | Area_5/office_40/Annotations 210 | Area_5/office_41/Annotations 211 | Area_5/office_42/Annotations 212 | Area_5/office_4/Annotations 213 | Area_5/office_5/Annotations 214 | Area_5/office_6/Annotations 215 | Area_5/office_7/Annotations 216 | Area_5/office_8/Annotations 217 | Area_5/office_9/Annotations 218 | Area_5/pantry_1/Annotations 219 | Area_5/storage_1/Annotations 220 | Area_5/storage_2/Annotations 221 | Area_5/storage_3/Annotations 222 | Area_5/storage_4/Annotations 223 | Area_5/WC_1/Annotations 224 | Area_5/WC_2/Annotations 225 | Area_6/conferenceRoom_1/Annotations 226 | Area_6/copyRoom_1/Annotations 227 | Area_6/hallway_1/Annotations 228 | Area_6/hallway_2/Annotations 229 | Area_6/hallway_3/Annotations 230 | Area_6/hallway_4/Annotations 231 | Area_6/hallway_5/Annotations 232 | Area_6/hallway_6/Annotations 233 | Area_6/lounge_1/Annotations 234 | Area_6/office_10/Annotations 235 | Area_6/office_11/Annotations 236 | Area_6/office_12/Annotations 237 | Area_6/office_13/Annotations 238 | Area_6/office_14/Annotations 239 | Area_6/office_15/Annotations 240 | Area_6/office_16/Annotations 241 | Area_6/office_17/Annotations 242 | Area_6/office_18/Annotations 243 | Area_6/office_19/Annotations 244 | Area_6/office_1/Annotations 245 | Area_6/office_20/Annotations 246 | Area_6/office_21/Annotations 247 | Area_6/office_22/Annotations 248 | Area_6/office_23/Annotations 249 | Area_6/office_24/Annotations 250 | Area_6/office_25/Annotations 251 | Area_6/office_26/Annotations 252 | Area_6/office_27/Annotations 253 | Area_6/office_28/Annotations 254 | Area_6/office_29/Annotations 255 | Area_6/office_2/Annotations 256 | Area_6/office_30/Annotations 257 | Area_6/office_31/Annotations 258 | Area_6/office_32/Annotations 259 | Area_6/office_33/Annotations 260 | Area_6/office_34/Annotations 261 | Area_6/office_35/Annotations 262 | Area_6/office_36/Annotations 263 | Area_6/office_37/Annotations 264 | Area_6/office_3/Annotations 265 | Area_6/office_4/Annotations 266 | Area_6/office_5/Annotations 267 | Area_6/office_6/Annotations 268 | Area_6/office_7/Annotations 269 | Area_6/office_8/Annotations 270 | Area_6/office_9/Annotations 271 | Area_6/openspace_1/Annotations 272 | Area_6/pantry_1/Annotations 273 | -------------------------------------------------------------------------------- /utils/meta/class_names.txt: -------------------------------------------------------------------------------- 1 | ceiling 2 | floor 3 | wall 4 | beam 5 | column 6 | window 7 | door 8 | table 9 | chair 10 | sofa 11 | bookcase 12 | board 13 | clutter 14 | -------------------------------------------------------------------------------- /utils/nearest_neighbors/KDTreeTableAdaptor.h: -------------------------------------------------------------------------------- 1 | /*********************************************************************** 2 | * Software License Agreement (BSD License) 3 | * 4 | * Copyright 2011-16 Jose Luis Blanco (joseluisblancoc@gmail.com). 5 | * All rights reserved. 6 | * 7 | * Redistribution and use in source and binary forms, with or without 8 | * modification, are permitted provided that the following conditions 9 | * are met: 10 | * 11 | * 1. Redistributions of source code must retain the above copyright 12 | * notice, this list of conditions and the following disclaimer. 13 | * 2. Redistributions in binary form must reproduce the above copyright 14 | * notice, this list of conditions and the following disclaimer in the 15 | * documentation and/or other materials provided with the distribution. 16 | * 17 | * THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR 18 | * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES 19 | * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. 20 | * IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT, 21 | * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT 22 | * NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, 23 | * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY 24 | * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 25 | * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF 26 | * THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 27 | *************************************************************************/ 28 | 29 | #pragma once 30 | 31 | #include "nanoflann.hpp" 32 | 33 | // #include 34 | 35 | // ===== This example shows how to use nanoflann with these types of containers: ======= 36 | //typedef std::vector > my_vector_of_vectors_t; 37 | //typedef std::vector my_vector_of_vectors_t; // This requires #include 38 | // ===================================================================================== 39 | 40 | 41 | /** A simple vector-of-vectors adaptor for nanoflann, without duplicating the storage. 42 | * The i'th vector represents a point in the state space. 43 | * 44 | * \tparam DIM If set to >0, it specifies a compile-time fixed dimensionality for the points in the data set, allowing more compiler optimizations. 45 | * \tparam num_t The type of the point coordinates (typically, double or float). 46 | * \tparam Distance The distance metric to use: nanoflann::metric_L1, nanoflann::metric_L2, nanoflann::metric_L2_Simple, etc. 47 | * \tparam IndexType The type for indices in the KD-tree index (typically, size_t of int) 48 | */ 49 | // template 50 | // struct KDTreeVectorAdaptor 51 | // { 52 | // typedef KDTreeVectorAdaptor self_t; 53 | // typedef typename Distance::template traits::distance_t metric_t; 54 | // typedef nanoflann::KDTreeSingleIndexAdaptor< metric_t,self_t,DIM,IndexType> index_t; 55 | 56 | // index_t* index; //! The kd-tree index for the user to call its methods as usual with any other FLANN index. 57 | // size_t dims; 58 | 59 | // /// Constructor: takes a const ref to the vector of vectors object with the data points 60 | // KDTreeVectorAdaptor(const size_t dims /* dimensionality */, const VectorType &mat, const int leaf_max_size = 10) : m_data(mat) 61 | // { 62 | // assert(mat.size() != 0); 63 | // this->dims= dims; 64 | // index = new index_t( static_cast(dims), *this /* adaptor */, nanoflann::KDTreeSingleIndexAdaptorParams(leaf_max_size ) ); 65 | // index->buildIndex(); 66 | // } 67 | 68 | // ~KDTreeVectorAdaptor() { 69 | // delete index; 70 | // } 71 | 72 | // const VectorType &m_data; 73 | 74 | // /** Query for the \a num_closest closest points to a given point (entered as query_point[0:dim-1]). 75 | // * Note that this is a short-cut method for index->findNeighbors(). 76 | // * The user can also call index->... methods as desired. 77 | // * \note nChecks_IGNORED is ignored but kept for compatibility with the original FLANN interface. 78 | // */ 79 | // inline void query(const num_t *query_point, const size_t num_closest, IndexType *out_indices, num_t *out_distances_sq, const int nChecks_IGNORED = 10) const 80 | // { 81 | // nanoflann::KNNResultSet resultSet(num_closest); 82 | // resultSet.init(out_indices, out_distances_sq); 83 | // index->findNeighbors(resultSet, query_point, nanoflann::SearchParams()); 84 | // } 85 | 86 | // /** @name Interface expected by KDTreeSingleIndexAdaptor 87 | // * @{ */ 88 | 89 | // const self_t & derived() const { 90 | // return *this; 91 | // } 92 | // self_t & derived() { 93 | // return *this; 94 | // } 95 | 96 | // // Must return the number of data points 97 | // inline size_t kdtree_get_point_count() const { 98 | // return m_data.size()/this->dims; 99 | // } 100 | 101 | // // Returns the dim'th component of the idx'th point in the class: 102 | // inline num_t kdtree_get_pt(const size_t idx, const size_t dim) const { 103 | // return m_data[idx*this->dims + dim]; 104 | // } 105 | 106 | // // Optional bounding-box computation: return false to default to a standard bbox computation loop. 107 | // // Return true if the BBOX was already computed by the class and returned in "bb" so it can be avoided to redo it again. 108 | // // Look at bb.size() to find out the expected dimensionality (e.g. 2 or 3 for point clouds) 109 | // template 110 | // bool kdtree_get_bbox(BBOX & /*bb*/) const { 111 | // return false; 112 | // } 113 | 114 | // /** @} */ 115 | 116 | // }; // end of KDTreeVectorOfVectorsAdaptor 117 | 118 | 119 | 120 | 121 | template 122 | struct KDTreeTableAdaptor 123 | { 124 | typedef KDTreeTableAdaptor self_t; 125 | typedef typename Distance::template traits::distance_t metric_t; 126 | typedef nanoflann::KDTreeSingleIndexAdaptor< metric_t,self_t,DIM,IndexType> index_t; 127 | 128 | index_t* index; //! The kd-tree index for the user to call its methods as usual with any other FLANN index. 129 | size_t dim; 130 | size_t npts; 131 | const TableType* m_data; 132 | 133 | /// Constructor: takes a const ref to the vector of vectors object with the data points 134 | KDTreeTableAdaptor(const size_t npts, const size_t dim, const TableType* mat, const int leaf_max_size = 10) : m_data(mat), dim(dim), npts(npts) 135 | { 136 | assert(npts != 0); 137 | index = new index_t( static_cast(dim), *this /* adaptor */, nanoflann::KDTreeSingleIndexAdaptorParams(leaf_max_size ) ); 138 | index->buildIndex(); 139 | } 140 | 141 | ~KDTreeTableAdaptor() { 142 | delete index; 143 | } 144 | 145 | 146 | /** Query for the \a num_closest closest points to a given point (entered as query_point[0:dim-1]). 147 | * Note that this is a short-cut method for index->findNeighbors(). 148 | * The user can also call index->... methods as desired. 149 | * \note nChecks_IGNORED is ignored but kept for compatibility with the original FLANN interface. 150 | */ 151 | inline void query(const num_t *query_point, const size_t num_closest, IndexType *out_indices, num_t *out_distances_sq, const int nChecks_IGNORED = 10) const 152 | { 153 | nanoflann::KNNResultSet resultSet(num_closest); 154 | resultSet.init(out_indices, out_distances_sq); 155 | index->findNeighbors(resultSet, query_point, nanoflann::SearchParams()); 156 | } 157 | 158 | /** @name Interface expected by KDTreeSingleIndexAdaptor 159 | * @{ */ 160 | 161 | const self_t & derived() const { 162 | return *this; 163 | } 164 | self_t & derived() { 165 | return *this; 166 | } 167 | 168 | // Must return the number of data points 169 | inline size_t kdtree_get_point_count() const { 170 | return this->npts; 171 | } 172 | 173 | // Returns the dim'th component of the idx'th point in the class: 174 | inline num_t kdtree_get_pt(const size_t pts_id, const size_t coord_id) const { 175 | return m_data[pts_id*this->dim + coord_id]; 176 | } 177 | 178 | // Optional bounding-box computation: return false to default to a standard bbox computation loop. 179 | // Return true if the BBOX was already computed by the class and returned in "bb" so it can be avoided to redo it again. 180 | // Look at bb.size() to find out the expected dimensionality (e.g. 2 or 3 for point clouds) 181 | template 182 | bool kdtree_get_bbox(BBOX & /*bb*/) const { 183 | return false; 184 | } 185 | 186 | /** @} */ 187 | 188 | }; // end of KDTreeVectorOfVectorsAdaptor 189 | 190 | -------------------------------------------------------------------------------- /utils/nearest_neighbors/knn.pyx: -------------------------------------------------------------------------------- 1 | # distutils: language = c++ 2 | # distutils: sources = knn.cxx 3 | 4 | import numpy as np 5 | cimport numpy as np 6 | import cython 7 | 8 | cdef extern from "knn_.h": 9 | void cpp_knn(const float* points, const size_t npts, const size_t dim, 10 | const float* queries, const size_t nqueries, 11 | const size_t K, long* indices) 12 | 13 | void cpp_knn_omp(const float* points, const size_t npts, const size_t dim, 14 | const float* queries, const size_t nqueries, 15 | const size_t K, long* indices) 16 | 17 | void cpp_knn_batch(const float* batch_data, const size_t batch_size, const size_t npts, const size_t dim, 18 | const float* queries, const size_t nqueries, 19 | const size_t K, long* batch_indices) 20 | 21 | void cpp_knn_batch_omp(const float* batch_data, const size_t batch_size, const size_t npts, const size_t dim, 22 | const float* queries, const size_t nqueries, 23 | const size_t K, long* batch_indices) 24 | 25 | void cpp_knn_batch_distance_pick(const float* batch_data, const size_t batch_size, const size_t npts, const size_t dim, 26 | float* queries, const size_t nqueries, 27 | const size_t K, long* batch_indices) 28 | 29 | void cpp_knn_batch_distance_pick_omp(const float* batch_data, const size_t batch_size, const size_t npts, const size_t dim, 30 | float* batch_queries, const size_t nqueries, 31 | const size_t K, long* batch_indices) 32 | 33 | def knn(pts, queries, K, omp=False): 34 | 35 | # define shape parameters 36 | cdef int npts 37 | cdef int dim 38 | cdef int K_cpp 39 | cdef int nqueries 40 | 41 | # define tables 42 | cdef np.ndarray[np.float32_t, ndim=2] pts_cpp 43 | cdef np.ndarray[np.float32_t, ndim=2] queries_cpp 44 | cdef np.ndarray[np.int64_t, ndim=2] indices_cpp 45 | 46 | # set shape values 47 | npts = pts.shape[0] 48 | nqueries = queries.shape[0] 49 | dim = pts.shape[1] 50 | K_cpp = K 51 | 52 | # create indices tensor 53 | indices = np.zeros((queries.shape[0], K), dtype=np.int64) 54 | 55 | pts_cpp = np.ascontiguousarray(pts, dtype=np.float32) 56 | queries_cpp = np.ascontiguousarray(queries, dtype=np.float32) 57 | indices_cpp = indices 58 | 59 | # normal estimation 60 | if omp: 61 | cpp_knn_omp( pts_cpp.data, npts, dim, 62 | queries_cpp.data, nqueries, 63 | K_cpp, indices_cpp.data) 64 | else: 65 | cpp_knn( pts_cpp.data, npts, dim, 66 | queries_cpp.data, nqueries, 67 | K_cpp, indices_cpp.data) 68 | 69 | return indices 70 | 71 | def knn_batch(pts, queries, K, omp=False): 72 | 73 | # define shape parameters 74 | cdef int batch_size 75 | cdef int npts 76 | cdef int nqueries 77 | cdef int K_cpp 78 | cdef int dim 79 | 80 | # define tables 81 | cdef np.ndarray[np.float32_t, ndim=3] pts_cpp 82 | cdef np.ndarray[np.float32_t, ndim=3] queries_cpp 83 | cdef np.ndarray[np.int64_t, ndim=3] indices_cpp 84 | 85 | # set shape values 86 | batch_size = pts.shape[0] 87 | npts = pts.shape[1] 88 | dim = pts.shape[2] 89 | nqueries = queries.shape[1] 90 | K_cpp = K 91 | 92 | # create indices tensor 93 | indices = np.zeros((pts.shape[0], queries.shape[1], K), dtype=np.int64) 94 | 95 | pts_cpp = np.ascontiguousarray(pts, dtype=np.float32) 96 | queries_cpp = np.ascontiguousarray(queries, dtype=np.float32) 97 | indices_cpp = indices 98 | 99 | # normal estimation 100 | if omp: 101 | cpp_knn_batch_omp( pts_cpp.data, batch_size, npts, dim, 102 | queries_cpp.data, nqueries, 103 | K_cpp, indices_cpp.data) 104 | else: 105 | cpp_knn_batch( pts_cpp.data, batch_size, npts, dim, 106 | queries_cpp.data, nqueries, 107 | K_cpp, indices_cpp.data) 108 | 109 | return indices 110 | 111 | def knn_batch_distance_pick(pts, nqueries, K, omp=False): 112 | 113 | # define shape parameters 114 | cdef int batch_size 115 | cdef int npts 116 | cdef int nqueries_cpp 117 | cdef int K_cpp 118 | cdef int dim 119 | 120 | # define tables 121 | cdef np.ndarray[np.float32_t, ndim=3] pts_cpp 122 | cdef np.ndarray[np.float32_t, ndim=3] queries_cpp 123 | cdef np.ndarray[np.int64_t, ndim=3] indices_cpp 124 | 125 | # set shape values 126 | batch_size = pts.shape[0] 127 | npts = pts.shape[1] 128 | dim = pts.shape[2] 129 | nqueries_cpp = nqueries 130 | K_cpp = K 131 | 132 | # create indices tensor 133 | indices = np.zeros((pts.shape[0], nqueries, K), dtype=np.long) 134 | queries = np.zeros((pts.shape[0], nqueries, dim), dtype=np.float32) 135 | 136 | pts_cpp = np.ascontiguousarray(pts, dtype=np.float32) 137 | queries_cpp = np.ascontiguousarray(queries, dtype=np.float32) 138 | indices_cpp = indices 139 | 140 | if omp: 141 | cpp_knn_batch_distance_pick_omp( pts_cpp.data, batch_size, npts, dim, 142 | queries_cpp.data, nqueries, 143 | K_cpp, indices_cpp.data) 144 | else: 145 | cpp_knn_batch_distance_pick( pts_cpp.data, batch_size, npts, dim, 146 | queries_cpp.data, nqueries, 147 | K_cpp, indices_cpp.data) 148 | 149 | return indices, queries -------------------------------------------------------------------------------- /utils/nearest_neighbors/knn_.cxx: -------------------------------------------------------------------------------- 1 | 2 | #include "knn_.h" 3 | #include "nanoflann.hpp" 4 | using namespace nanoflann; 5 | 6 | #include "KDTreeTableAdaptor.h" 7 | 8 | #include 9 | #include 10 | #include 11 | #include 12 | 13 | #include 14 | #include 15 | #include 16 | #include 17 | 18 | using namespace std; 19 | 20 | 21 | 22 | void cpp_knn(const float* points, const size_t npts, const size_t dim, 23 | const float* queries, const size_t nqueries, 24 | const size_t K, long* indices){ 25 | 26 | // create the kdtree 27 | typedef KDTreeTableAdaptor< float, float> KDTree; 28 | KDTree mat_index(npts, dim, points, 10); 29 | mat_index.index->buildIndex(); 30 | 31 | std::vector out_dists_sqr(K); 32 | std::vector out_ids(K); 33 | 34 | // iterate over the points 35 | for(size_t i=0; i resultSet(K); 38 | resultSet.init(&out_ids[0], &out_dists_sqr[0] ); 39 | mat_index.index->findNeighbors(resultSet, &queries[i*dim], nanoflann::SearchParams(10)); 40 | for(size_t j=0; j KDTree; 52 | KDTree mat_index(npts, dim, points, 10); 53 | mat_index.index->buildIndex(); 54 | 55 | 56 | // iterate over the points 57 | # pragma omp parallel for 58 | for(size_t i=0; i out_ids(K); 60 | std::vector out_dists_sqr(K); 61 | 62 | nanoflann::KNNResultSet resultSet(K); 63 | resultSet.init(&out_ids[0], &out_dists_sqr[0] ); 64 | mat_index.index->findNeighbors(resultSet, &queries[i*dim], nanoflann::SearchParams(10)); 65 | for(size_t j=0; j KDTree; 83 | KDTree mat_index(npts, dim, points, 10); 84 | 85 | mat_index.index->buildIndex(); 86 | 87 | std::vector out_dists_sqr(K); 88 | std::vector out_ids(K); 89 | 90 | // iterate over the points 91 | for(size_t i=0; i resultSet(K); 93 | resultSet.init(&out_ids[0], &out_dists_sqr[0] ); 94 | mat_index.index->findNeighbors(resultSet, &queries[bid*nqueries*dim + i*dim], nanoflann::SearchParams(10)); 95 | for(size_t j=0; j KDTree; 116 | KDTree mat_index(npts, dim, points, 10); 117 | 118 | mat_index.index->buildIndex(); 119 | 120 | std::vector out_dists_sqr(K); 121 | std::vector out_ids(K); 122 | 123 | // iterate over the points 124 | for(size_t i=0; i resultSet(K); 126 | resultSet.init(&out_ids[0], &out_dists_sqr[0] ); 127 | mat_index.index->findNeighbors(resultSet, &queries[bid*nqueries*dim + i*dim], nanoflann::SearchParams(10)); 128 | for(size_t j=0; j KDTree; 153 | KDTree tree(npts, dim, points, 10); 154 | tree.index->buildIndex(); 155 | 156 | vector used(npts, 0); 157 | int current_id = 0; 158 | for(size_t ptid=0; ptid possible_ids; 162 | while(possible_ids.size() == 0){ 163 | for(size_t i=0; i query(3); 178 | for(size_t i=0; i dists(K); 183 | std::vector ids(K); 184 | nanoflann::KNNResultSet resultSet(K); 185 | resultSet.init(&ids[0], &dists[0] ); 186 | tree.index->findNeighbors(resultSet, &query[0], nanoflann::SearchParams(10)); 187 | 188 | for(size_t i=0; i KDTree; 221 | KDTree tree(npts, dim, points, 10); 222 | tree.index->buildIndex(); 223 | 224 | vector used(npts, 0); 225 | int current_id = 0; 226 | for(size_t ptid=0; ptid possible_ids; 230 | while(possible_ids.size() == 0){ 231 | for(size_t i=0; i query(3); 246 | for(size_t i=0; i dists(K); 251 | std::vector ids(K); 252 | nanoflann::KNNResultSet resultSet(K); 253 | resultSet.init(&ids[0], &dists[0] ); 254 | tree.index->findNeighbors(resultSet, &query[0], nanoflann::SearchParams(10)); 255 | 256 | for(size_t i=0; i 4 | void cpp_knn(const float* points, const size_t npts, const size_t dim, 5 | const float* queries, const size_t nqueries, 6 | const size_t K, long* indices); 7 | 8 | void cpp_knn_omp(const float* points, const size_t npts, const size_t dim, 9 | const float* queries, const size_t nqueries, 10 | const size_t K, long* indices); 11 | 12 | 13 | void cpp_knn_batch(const float* batch_data, const size_t batch_size, const size_t npts, const size_t dim, 14 | const float* queries, const size_t nqueries, 15 | const size_t K, long* batch_indices); 16 | 17 | void cpp_knn_batch_omp(const float* batch_data, const size_t batch_size, const size_t npts, const size_t dim, 18 | const float* queries, const size_t nqueries, 19 | const size_t K, long* batch_indices); 20 | 21 | void cpp_knn_batch_distance_pick(const float* batch_data, const size_t batch_size, const size_t npts, const size_t dim, 22 | float* queries, const size_t nqueries, 23 | const size_t K, long* batch_indices); 24 | 25 | void cpp_knn_batch_distance_pick_omp(const float* batch_data, const size_t batch_size, const size_t npts, const size_t dim, 26 | float* batch_queries, const size_t nqueries, 27 | const size_t K, long* batch_indices); -------------------------------------------------------------------------------- /utils/nearest_neighbors/setup.py: -------------------------------------------------------------------------------- 1 | from distutils.core import setup 2 | from distutils.extension import Extension 3 | from Cython.Distutils import build_ext 4 | import numpy 5 | 6 | 7 | 8 | ext_modules = [Extension( 9 | "nearest_neighbors", 10 | sources=["knn.pyx", "knn_.cxx",], # source file(s) 11 | include_dirs=["./", numpy.get_include()], 12 | language="c++", 13 | extra_compile_args = [ "-std=c++11", "-fopenmp",], 14 | extra_link_args=["-std=c++11", '-fopenmp'], 15 | )] 16 | 17 | setup( 18 | name = "KNN NanoFLANN", 19 | ext_modules = ext_modules, 20 | cmdclass = {'build_ext': build_ext}, 21 | ) 22 | -------------------------------------------------------------------------------- /utils/nearest_neighbors/test.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import lib.python.nearest_neighbors as nearest_neighbors 3 | import time 4 | 5 | batch_size = 16 6 | num_points = 81920 7 | K = 16 8 | pc = np.random.rand(batch_size, num_points, 3).astype(np.float32) 9 | 10 | # nearest neighbours 11 | start = time.time() 12 | neigh_idx = nearest_neighbors.knn_batch(pc, pc, K, omp=True) 13 | print(time.time() - start) 14 | 15 | 16 | -------------------------------------------------------------------------------- /utils/s3dis_dp.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import pandas as pd 3 | import os, sys, glob, pickle 4 | from os.path import join, exists, dirname, abspath 5 | from sklearn.neighbors import KDTree 6 | 7 | BASE_DIR = dirname(abspath(__file__)) 8 | ROOT_DIR = dirname(BASE_DIR) 9 | sys.path.append(BASE_DIR) 10 | sys.path.append(ROOT_DIR) 11 | from helper_ply import write_ply 12 | from helper_dp import DataProcessing as DP 13 | 14 | dataset_path = './data/S3DIS/Stanford3dDataset_v1.2_Aligned_Version' 15 | anno_paths = [line.rstrip() for line in open(join(BASE_DIR, 'meta/anno_paths.txt'))] 16 | anno_paths = [join(dataset_path, p) for p in anno_paths] 17 | 18 | gt_class = [x.rstrip() for x in open(join(BASE_DIR, 'meta/class_names.txt'))] 19 | gt_class2label = {cls: i for i, cls in enumerate(gt_class)} 20 | 21 | sub_grid_size = 0.04 22 | original_pc_folder = join(dirname(dataset_path), 'original_ply') 23 | sub_pc_folder = join(dirname(dataset_path), 'input_{:.3f}'.format(sub_grid_size)) 24 | os.mkdir(original_pc_folder) if not exists(original_pc_folder) else None 25 | os.mkdir(sub_pc_folder) if not exists(sub_pc_folder) else None 26 | out_format = '.ply' 27 | 28 | 29 | def convert_pc2ply(anno_path, save_path): 30 | """ 31 | Convert original dataset files to ply file (each line is XYZRGBL). 32 | We aggregated all the points from each instance in the room. 33 | :param anno_path: path to annotations. e.g. Area_1/office_2/Annotations/ 34 | :param save_path: path to save original point clouds (each line is XYZRGBL) 35 | :return: None 36 | """ 37 | data_list = [] 38 | 39 | for f in glob.glob(join(anno_path, '*.txt')): 40 | class_name = os.path.basename(f).split('_')[0] 41 | if class_name not in gt_class: # note: in some room there is 'staris' class.. 42 | class_name = 'clutter' 43 | pc = pd.read_csv(f, header=None, delim_whitespace=True).values 44 | labels = np.ones((pc.shape[0], 1)) * gt_class2label[class_name] 45 | data_list.append(np.concatenate([pc, labels], 1)) # Nx7 46 | 47 | pc_label = np.concatenate(data_list, 0) 48 | xyz_min = np.amin(pc_label, axis=0)[0:3] 49 | pc_label[:, 0:3] -= xyz_min 50 | 51 | xyz = pc_label[:, :3].astype(np.float32) 52 | colors = pc_label[:, 3:6].astype(np.uint8) 53 | labels = pc_label[:, 6].astype(np.uint8) 54 | write_ply(save_path, (xyz, colors, labels), ['x', 'y', 'z', 'red', 'green', 'blue', 'class']) 55 | 56 | # save sub_cloud and KDTree file 57 | sub_xyz, sub_colors, sub_labels = DP.grid_sub_sampling(xyz, colors, labels, sub_grid_size) 58 | sub_colors = sub_colors / 255.0 59 | sub_ply_file = join(sub_pc_folder, save_path.split('/')[-1][:-4] + '.ply') 60 | write_ply(sub_ply_file, [sub_xyz, sub_colors, sub_labels], ['x', 'y', 'z', 'red', 'green', 'blue', 'class']) 61 | 62 | search_tree = KDTree(sub_xyz) 63 | kd_tree_file = join(sub_pc_folder, str(save_path.split('/')[-1][:-4]) + '_KDTree.pkl') 64 | with open(kd_tree_file, 'wb') as f: 65 | pickle.dump(search_tree, f) 66 | 67 | proj_idx = np.squeeze(search_tree.query(xyz, return_distance=False)) 68 | proj_idx = proj_idx.astype(np.int32) 69 | proj_save = join(sub_pc_folder, str(save_path.split('/')[-1][:-4]) + '_proj.pkl') 70 | with open(proj_save, 'wb') as f: 71 | pickle.dump([proj_idx, labels], f) 72 | 73 | 74 | if __name__ == '__main__': 75 | # Note: there is an extra character in the v1.2 data in Area_5/hallway_6. It's fixed manually. 76 | for annotation_path in anno_paths: 77 | print(annotation_path) 78 | elements = str(annotation_path).split('/') 79 | out_file_name = elements[-3] + '_' + elements[-2] + out_format 80 | convert_pc2ply(annotation_path, join(original_pc_folder, out_file_name)) 81 | -------------------------------------------------------------------------------- /utils/semantic3d_dp.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import sys 3 | import os, glob, pickle 4 | from os.path import join, exists, dirname, abspath 5 | from sklearn.neighbors import KDTree 6 | 7 | BASE_DIR = dirname(abspath(__file__)) 8 | ROOT_DIR = dirname(BASE_DIR) 9 | sys.path.append(BASE_DIR) 10 | sys.path.append(ROOT_DIR) 11 | from helper_ply import write_ply 12 | from helper_dp import DataProcessing as DP 13 | 14 | 15 | grid_size = 0.06 16 | dataset_path = './data/semantic3d/original_data' 17 | original_pc_folder = join(dirname(dataset_path), 'original_ply') 18 | sub_pc_folder = join(dirname(dataset_path), 'input_{:.3f}'.format(grid_size)) 19 | os.mkdir(original_pc_folder) if not exists(original_pc_folder) else None 20 | os.mkdir(sub_pc_folder) if not exists(sub_pc_folder) else None 21 | 22 | for pc_path in glob.glob(join(dataset_path, '*.txt')): 23 | print(pc_path) 24 | file_name = pc_path.split('/')[-1][:-4] 25 | 26 | # check if it has already calculated 27 | if exists(join(sub_pc_folder, file_name + '_KDTree.pkl')): 28 | continue 29 | 30 | pc = DP.load_pc_semantic3d(pc_path) 31 | # check if label exists 32 | label_path = pc_path[:-4] + '.labels' 33 | if exists(label_path): 34 | labels = DP.load_label_semantic3d(label_path) 35 | full_ply_path = join(original_pc_folder, file_name + '.ply') 36 | 37 | #  Subsample to save space 38 | sub_points, sub_colors, sub_labels = DP.grid_sub_sampling(pc[:, :3].astype(np.float32), 39 | pc[:, 4:7].astype(np.uint8), labels, 0.01) 40 | sub_labels = np.squeeze(sub_labels) 41 | 42 | write_ply(full_ply_path, (sub_points, sub_colors, sub_labels), ['x', 'y', 'z', 'red', 'green', 'blue', 'class']) 43 | 44 | # save sub_cloud and KDTree file 45 | sub_xyz, sub_colors, sub_labels = DP.grid_sub_sampling(sub_points, sub_colors, sub_labels, grid_size) 46 | sub_colors = sub_colors / 255.0 47 | sub_labels = np.squeeze(sub_labels) 48 | sub_ply_file = join(sub_pc_folder, file_name + '.ply') 49 | write_ply(sub_ply_file, [sub_xyz, sub_colors, sub_labels], ['x', 'y', 'z', 'red', 'green', 'blue', 'class']) 50 | 51 | search_tree = KDTree(sub_xyz, leaf_size=50) 52 | kd_tree_file = join(sub_pc_folder, file_name + '_KDTree.pkl') 53 | with open(kd_tree_file, 'wb') as f: 54 | pickle.dump(search_tree, f) 55 | 56 | proj_idx = np.squeeze(search_tree.query(sub_points, return_distance=False)) 57 | proj_idx = proj_idx.astype(np.int32) 58 | proj_save = join(sub_pc_folder, file_name + '_proj.pkl') 59 | with open(proj_save, 'wb') as f: 60 | pickle.dump([proj_idx, labels], f) 61 | 62 | else: 63 | full_ply_path = join(original_pc_folder, file_name + '.ply') 64 | write_ply(full_ply_path, (pc[:, :3].astype(np.float32), pc[:, 4:7].astype(np.uint8)), 65 | ['x', 'y', 'z', 'red', 'green', 'blue']) 66 | 67 | # save sub_cloud and KDTree file 68 | sub_xyz, sub_colors = DP.grid_sub_sampling(pc[:, :3].astype(np.float32), pc[:, 4:7].astype(np.uint8), 69 | grid_size=grid_size) 70 | sub_colors = sub_colors / 255.0 71 | sub_ply_file = join(sub_pc_folder, file_name + '.ply') 72 | write_ply(sub_ply_file, [sub_xyz, sub_colors], ['x', 'y', 'z', 'red', 'green', 'blue']) 73 | labels = np.zeros(pc.shape[0], dtype=np.uint8) 74 | 75 | search_tree = KDTree(sub_xyz, leaf_size=50) 76 | kd_tree_file = join(sub_pc_folder, file_name + '_KDTree.pkl') 77 | with open(kd_tree_file, 'wb') as f: 78 | pickle.dump(search_tree, f) 79 | 80 | proj_idx = np.squeeze(search_tree.query(pc[:, :3].astype(np.float32), return_distance=False)) 81 | proj_idx = proj_idx.astype(np.int32) 82 | proj_save = join(sub_pc_folder, file_name + '_proj.pkl') 83 | with open(proj_save, 'wb') as f: 84 | pickle.dump([proj_idx, labels], f) 85 | -------------------------------------------------------------------------------- /visualization.py: -------------------------------------------------------------------------------- 1 | # open3d == 0.10.0.1 2 | import open3d 3 | import argparse 4 | import numpy as np 5 | from helper_ply import read_ply 6 | 7 | s3dis_color = [[255,248,220], [220,220,220], [139,71,38], [238,197,145], [70,130,180], [179,238,58], [110,139,61], [105,105,105], [0,0,128], [205,92,92], [244,164,96], [147,112,219], [255,228,225]] 8 | semantic3d_color = [[220,220,220], [154,205,50], [0,100,0], [238,220,130], [139,115,85], [70,130,180], [255,231,186], [255,99,71]] 9 | 10 | def draw_pc(pc_xyzrgb): 11 | """ 12 | Draw Point Cloud 13 | """ 14 | pc = open3d.geometry.PointCloud() 15 | pc.points = open3d.utility.Vector3dVector(pc_xyzrgb[:, 0:3]) 16 | if pc_xyzrgb.shape[1] == 3: 17 | open3d.draw_geometries([pc]) 18 | return 0 19 | if np.max(pc_xyzrgb[:, 3:6]) > 20: ## 0-255 20 | pc.colors = open3d.utility.Vector3dVector(pc_xyzrgb[:, 3:6] / 255.) 21 | else: 22 | pc.colors = open3d.utility.Vector3dVector(pc_xyzrgb[:, 3:6]) 23 | open3d.visualization.draw_geometries([pc]) 24 | return 0 25 | 26 | def s3dis(origin_file_name, label_file_name): 27 | """ 28 | Visualization for S3DIS 29 | """ 30 | origin_data = read_ply(origin_file_name) 31 | label_data = read_ply(label_file_name) 32 | 33 | pred = label_data["pred"] 34 | label = label_data["label"] 35 | 36 | x = origin_data["x"] 37 | y = origin_data["y"] 38 | z = origin_data["z"] 39 | 40 | r = origin_data["red"] 41 | g = origin_data["green"] 42 | b = origin_data["blue"] 43 | 44 | label_rgb = np.array([s3dis_color[i] for i in label]) 45 | 46 | pred_rgb = np.array([s3dis_color[i] for i in pred]) 47 | 48 | xyzrgb = np.vstack((x,y,z,r,g,b)).T 49 | draw_pc(xyzrgb) 50 | 51 | pred_xyzrgb = np.vstack(([x,y,z],pred_rgb.T)).T 52 | draw_pc(pred_xyzrgb) 53 | 54 | label_xyzrgb = np.vstack(([x,y,z],label_rgb.T)).T 55 | draw_pc(label_xyzrgb) 56 | 57 | def semantic3d(origin_file_name, label_file_name): 58 | """ 59 | Visualization for Semantic3D 60 | """ 61 | origin_data = read_ply(origin_file_name) 62 | 63 | with open(label_file_name, 'r') as f: pred = f.readlines() 64 | 65 | x = origin_data["x"] 66 | y = origin_data["y"] 67 | z = origin_data["z"] 68 | 69 | r = origin_data["red"] 70 | g = origin_data["green"] 71 | b = origin_data["blue"] 72 | 73 | pred_rgb = np.array([semantic3d_color[int(i)-1] for i in pred]) 74 | 75 | xyzrgb = np.vstack((x,y,z,r,g,b)).T 76 | draw_pc(xyzrgb) 77 | 78 | pred_xyzrgb = np.vstack(([x,y,z],pred_rgb.T)).T 79 | draw_pc(pred_xyzrgb) 80 | 81 | 82 | 83 | if __name__ == "__main__": 84 | 85 | parser = argparse.ArgumentParser() 86 | parser.add_argument('--dataset', type=str, default='s3dis', help='s3dis or semantic3d [default: s3dis]') 87 | parser.add_argument('--ply_path', type=str, help='origin ply path') 88 | parser.add_argument('--label_path', type=str, help='label path') 89 | args = parser.parse_args() 90 | 91 | if args.dataset == 's3dis': 92 | s3dis(args.ply_path, args.label_path) 93 | elif args.dataset == 'semantic3d': 94 | semantic3d(args.ply_path, args.label_path) 95 | else: 96 | print("Wrong dataset. S3DIS or Semantic3D") --------------------------------------------------------------------------------