├── README.md ├── evaluate.py ├── modelnet_dataset.py ├── models ├── acnn_cls_rings.py └── acnn_segm_rings.py ├── part_segm ├── evaluate.py ├── evaluate_job.sh ├── part_dataset_all_normal.py └── train.py ├── pics └── teaser.png ├── provider.py ├── scannet ├── README.md └── normal_estimation │ ├── 1_read_and_save_as_files.py │ ├── 2_compute_normals_and_curvature.cpp │ ├── 3_pickle_normals.py │ ├── CMakeLists.txt │ └── run.sh ├── tf_ops ├── 3d_interpolation │ ├── .gitignore │ ├── interpolate.cpp │ ├── tf_interpolate.cpp │ ├── tf_interpolate.py │ └── tf_interpolate_compile.sh ├── grouping_ring │ ├── .gitignore │ ├── tf_grouping_ring.cpp │ ├── tf_grouping_ring.py │ ├── tf_grouping_ring_compile.sh │ └── tf_grouping_ring_g.cu ├── ordering │ ├── .gitignore │ ├── tf_ordering.cpp │ ├── tf_ordering.py │ ├── tf_ordering_compile.sh │ └── tf_ordering_g.cu └── sampling │ ├── .gitignore │ ├── tf_sampling.cpp │ ├── tf_sampling.py │ ├── tf_sampling_compile.sh │ └── tf_sampling_g.cu ├── train.py └── utils ├── pointnet_util.py └── tf_util.py /README.md: -------------------------------------------------------------------------------- 1 | ## A-CNN: *Annularly Convolutional Neural Networks on Point Clouds* 2 | Created by Artem Komarichev, Zichun Zhong, Jing Hua from Department of Computer Science, Wayne State University. 3 | 4 | ![teaser image](https://github.com/artemkomarichev/a-cnn/blob/master/pics/teaser.png) 5 | 6 | ### Introduction 7 | 8 | Our paper (arXiV) proposes a new approach to define and compute convolution directly on 3D point clouds by the proposed annular convolution. 9 | 10 | To appear, Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), June 2019 11 | 12 | ### A-CNN usage 13 | 14 | We provide the code of A-CNN model that was tested with Tensorflow 1.3.0, CUDA 8.0, and python 3.6 on Ubuntu 16.04. We run all our experiments on a single NVIDIA Titan Xp GPU with 12GB GDDR5X. 15 | 16 | * Classification Task 17 | 18 | Download *ModelNet-40* dataset first. Point clouds are sampled from meshes with 10K points (XYZ + normals) per shape and provided by PointNet++. 19 | 20 | To train a classification A-CNN model on *ModelNet-40* dataset type the following command: 21 | 22 | python train.py 23 | 24 | To evaluate a trained model run the following script: 25 | 26 | python evaluate.py 27 | 28 | * Part Segmentation Task 29 | 30 | Download *ShapeNet-part* dataset first. Each point cloud represented by 2K points (XYZ + normals) and provided by PointNet++. 31 | 32 | To train a part segmentation A-CNN model on *ShapeNet-part* dataset type the following commands: 33 | 34 | cd part_segm 35 | python train.py 36 | 37 | To evaluate a trained segmentation model run the following script: 38 | 39 | ./evaluate_job.sh 40 | 41 | * Semantic Segmentation Task 42 | 43 | Download *S3DIS* and *ScanNet* datasets provided by PointNet/PointNet++. *S3DIS* contains XYZ + RGB information. *ScanNet* only has geometry information (*XYZ* only), no color. 44 | 45 | To estimate normals we used PCL library. The script to estimate normals for ScanNet data could be found here: 46 | 47 | cd scannet/normal_estimation 48 | ./run.sh 49 | 50 | ### Citation 51 | If you find our work useful in your research, please cite our work: 52 | 53 | @InProceedings{komarichev2019acnn, 54 | title={A-CNN: Annularly Convolutional Neural Networks on Point Clouds}, 55 | author={Komarichev, Artem and Zhong, Zichun and Hua, Jing}, 56 | booktitle={Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition}, 57 | year={2019} 58 | } -------------------------------------------------------------------------------- /evaluate.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import numpy as np 3 | import argparse 4 | import socket 5 | import importlib 6 | import time 7 | import os 8 | # import scipy.misc 9 | import sys 10 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 11 | ROOT_DIR = os.path.dirname(BASE_DIR) 12 | DATA_DIR = os.path.dirname(ROOT_DIR) 13 | sys.path.append(BASE_DIR) 14 | sys.path.append(os.path.join(BASE_DIR, 'models')) 15 | sys.path.append(os.path.join(BASE_DIR, 'utils')) 16 | import provider 17 | import modelnet_dataset 18 | 19 | parser = argparse.ArgumentParser() 20 | parser.add_argument('--gpu', type=int, default=1, help='GPU to use [default: GPU 0]') 21 | parser.add_argument('--model', default='acnn_cls_rings', help='Model name. [default: pointnet2_cls_ssg]') 22 | parser.add_argument('--batch_size', type=int, default=16, help='Batch Size during training [default: 16]') 23 | parser.add_argument('--num_point', type=int, default=1024, help='Point Number [256/512/1024/2048] [default: 1024]') 24 | parser.add_argument('--model_path', default='log/model_best_acc.ckpt', help='model checkpoint file path [default: log/model.ckpt]') 25 | parser.add_argument('--dump_dir', default='dump', help='dump folder path [dump]') 26 | parser.add_argument('--normal', action='store_true', default='store_true', help='Whether to use normal information') 27 | parser.add_argument('--num_votes', type=int, default=12, help='Aggregate classification scores from multiple rotations [default: 1]') 28 | FLAGS = parser.parse_args() 29 | 30 | 31 | BATCH_SIZE = FLAGS.batch_size 32 | NUM_POINT = FLAGS.num_point 33 | MODEL_PATH = FLAGS.model_path 34 | GPU_INDEX = FLAGS.gpu 35 | MODEL = importlib.import_module(FLAGS.model) # import network module 36 | DUMP_DIR = FLAGS.dump_dir 37 | if not os.path.exists(DUMP_DIR): os.mkdir(DUMP_DIR) 38 | LOG_FOUT = open(os.path.join(DUMP_DIR, 'log_evaluate.txt'), 'w') 39 | LOG_FOUT.write(str(FLAGS)+'\n') 40 | 41 | NUM_CLASSES = 40 42 | SHAPE_NAMES = [line.rstrip() for line in \ 43 | open(os.path.join(DATA_DIR, 'data_10K/modelnet40_normal_resampled/shape_names.txt'))] 44 | 45 | HOSTNAME = socket.gethostname() 46 | 47 | # Shapenet official train/test split 48 | if FLAGS.normal: 49 | assert(NUM_POINT<=10000) 50 | DATA_PATH = os.path.join(DATA_DIR, 'data_10K/modelnet40_normal_resampled') 51 | TRAIN_DATASET = modelnet_dataset.ModelNetDataset(root=DATA_PATH, npoints=NUM_POINT, split='train', normal_channel=FLAGS.normal, batch_size=BATCH_SIZE) 52 | TEST_DATASET = modelnet_dataset.ModelNetDataset(root=DATA_PATH, npoints=NUM_POINT, split='test', normal_channel=FLAGS.normal, batch_size=BATCH_SIZE) 53 | 54 | def log_string(out_str): 55 | LOG_FOUT.write(out_str+'\n') 56 | LOG_FOUT.flush() 57 | print(out_str) 58 | 59 | def evaluate(num_votes): 60 | is_training = False 61 | 62 | with tf.device('/gpu:'+str(GPU_INDEX)): 63 | pointclouds_pl, labels_pl, normals_pl = MODEL.placeholder_inputs(BATCH_SIZE, NUM_POINT) 64 | is_training_pl = tf.placeholder(tf.bool, shape=()) 65 | 66 | # simple model 67 | pred, end_points = MODEL.get_model(pointclouds_pl, normals_pl, is_training_pl) 68 | MODEL.get_loss(pred, labels_pl, end_points) 69 | losses = tf.get_collection('losses') 70 | total_loss = tf.add_n(losses, name='total_loss') 71 | 72 | # Add ops to save and restore all the variables. 73 | saver = tf.train.Saver() 74 | 75 | # Create a session 76 | config = tf.ConfigProto() 77 | config.gpu_options.allow_growth = True 78 | config.allow_soft_placement = True 79 | config.log_device_placement = False 80 | sess = tf.Session(config=config) 81 | 82 | # Restore variables from disk. 83 | saver.restore(sess, MODEL_PATH) 84 | log_string("Model restored.") 85 | 86 | ops = {'pointclouds_pl': pointclouds_pl, 87 | 'normals_pl': normals_pl, 88 | 'labels_pl': labels_pl, 89 | 'is_training_pl': is_training_pl, 90 | 'pred': pred, 91 | 'loss': total_loss} 92 | 93 | eval_one_epoch(sess, ops, num_votes) 94 | 95 | def eval_one_epoch(sess, ops, num_votes=1, topk=1): 96 | is_training = False 97 | 98 | # Make sure batch data is of same size 99 | cur_batch_data = np.zeros((BATCH_SIZE,NUM_POINT,3)) 100 | cur_batch_normals = np.zeros((BATCH_SIZE,NUM_POINT,3)) 101 | cur_batch_label = np.zeros((BATCH_SIZE), dtype=np.int32) 102 | 103 | total_correct = 0 104 | total_seen = 0 105 | loss_sum = 0 106 | batch_idx = 0 107 | shape_ious = [] 108 | total_seen_class = [0 for _ in range(NUM_CLASSES)] 109 | total_correct_class = [0 for _ in range(NUM_CLASSES)] 110 | 111 | while TEST_DATASET.has_next_batch(): 112 | batch_data, batch_label = TEST_DATASET.next_batch(augment=False) 113 | bsize = batch_data.shape[0] 114 | 115 | print('Batch: %03d, batch size: %d'%(batch_idx, bsize)) 116 | 117 | batch_pred_sum = np.zeros((BATCH_SIZE, NUM_CLASSES)) # score for classes 118 | cur_batch_label[0:bsize] = batch_label 119 | 120 | for vote_idx in range(num_votes): 121 | original_data = np.copy(batch_data) 122 | jittered_data = provider.random_scale_point_cloud(original_data[:,:,:3]) 123 | original_data[:,:,:3] = jittered_data 124 | shuffled_data = provider.shuffle_points(original_data) 125 | 126 | cur_batch_data[0:bsize,...] = shuffled_data[:,:,:3] 127 | cur_batch_normals[0:bsize,...] = shuffled_data[:,:,3:] 128 | 129 | feed_dict = {ops['pointclouds_pl']: cur_batch_data, 130 | ops['labels_pl']: cur_batch_label, 131 | ops['normals_pl']: cur_batch_normals, 132 | ops['is_training_pl']: is_training} 133 | loss_val, pred_val = sess.run([ops['loss'], ops['pred']], feed_dict=feed_dict) 134 | batch_pred_sum += pred_val 135 | pred_val = np.argmax(batch_pred_sum, 1) 136 | correct = np.sum(pred_val[0:bsize] == batch_label[0:bsize]) 137 | total_correct += correct 138 | total_seen += bsize 139 | loss_sum += loss_val 140 | batch_idx += 1 141 | for i in range(bsize): 142 | l = batch_label[i] 143 | total_seen_class[l] += 1 144 | total_correct_class[l] += (pred_val[i] == l) 145 | 146 | log_string('eval mean loss: %f' % (loss_sum / float(batch_idx))) 147 | log_string('eval accuracy: %f'% (total_correct / float(total_seen))) 148 | log_string('eval avg class acc: %f' % (np.mean(np.array(total_correct_class)/np.array(total_seen_class,dtype=np.float)))) 149 | 150 | class_accuracies = np.array(total_correct_class)/np.array(total_seen_class,dtype=np.float) 151 | for i, name in enumerate(SHAPE_NAMES): 152 | log_string('%10s:\t%0.3f' % (name, class_accuracies[i])) 153 | 154 | 155 | if __name__=='__main__': 156 | with tf.Graph().as_default(): 157 | evaluate(num_votes=FLAGS.num_votes) 158 | LOG_FOUT.close() 159 | -------------------------------------------------------------------------------- /modelnet_dataset.py: -------------------------------------------------------------------------------- 1 | ''' 2 | ModelNet dataset. Support ModelNet40, ModelNet10, XYZ and normal channels. Up to 10000 points. 3 | ''' 4 | 5 | import os 6 | import os.path 7 | import json 8 | import numpy as np 9 | import sys 10 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 11 | sys.path.append(os.path.join(BASE_DIR, 'utils')) 12 | import provider 13 | 14 | def pc_normalize(pc): 15 | l = pc.shape[0] 16 | centroid = np.mean(pc, axis=0) 17 | pc = pc - centroid 18 | m = np.max(np.sqrt(np.sum(pc**2, axis=1))) 19 | pc = pc / m 20 | return pc 21 | 22 | class ModelNetDataset(): 23 | def __init__(self, root, batch_size = 32, npoints = 1024, split='train', normalize=True, normal_channel=False, modelnet10=False, cache_size=15000, shuffle=None): 24 | self.root = root 25 | self.batch_size = batch_size 26 | self.npoints = npoints 27 | self.normalize = normalize 28 | if modelnet10: 29 | self.catfile = os.path.join(self.root, 'modelnet10_shape_names.txt') 30 | else: 31 | self.catfile = os.path.join(self.root, 'shape_names.txt') 32 | self.cat = [line.rstrip() for line in open(self.catfile)] 33 | self.classes = dict(zip(self.cat, range(len(self.cat)))) 34 | self.normal_channel = normal_channel 35 | 36 | shape_ids = {} 37 | if modelnet10: 38 | shape_ids['train'] = [line.rstrip() for line in open(os.path.join(self.root, 'modelnet10_train.txt'))] 39 | shape_ids['test']= [line.rstrip() for line in open(os.path.join(self.root, 'modelnet10_test.txt'))] 40 | else: 41 | shape_ids['train'] = [line.rstrip() for line in open(os.path.join(self.root, 'modelnet40_train.txt'))] 42 | shape_ids['test']= [line.rstrip() for line in open(os.path.join(self.root, 'modelnet40_test.txt'))] 43 | assert(split=='train' or split=='test') 44 | shape_names = ['_'.join(x.split('_')[0:-1]) for x in shape_ids[split]] 45 | # list of (shape_name, shape_txt_file_path) tuple 46 | self.datapath = [(shape_names[i], os.path.join(self.root, shape_names[i], shape_ids[split][i])+'.txt') for i in range(len(shape_ids[split]))] 47 | 48 | self.cache_size = cache_size # how many data points to cache in memory 49 | self.cache = {} # from index to (point_set, cls) tuple 50 | 51 | if shuffle is None: 52 | if split == 'train': self.shuffle = True 53 | else: self.shuffle = False 54 | else: 55 | self.shuffle = shuffle 56 | 57 | self.reset() 58 | 59 | def _augment_batch_data(self, batch_data): 60 | jittered_data = provider.random_scale_point_cloud(batch_data[:,:,0:3]) 61 | jittered_data = provider.shift_point_cloud(jittered_data) 62 | jittered_data = provider.jitter_point_cloud(jittered_data) 63 | batch_data[:,:,0:3] = jittered_data 64 | return provider.shuffle_points(batch_data) 65 | 66 | def _get_item(self, index): 67 | if index in self.cache: 68 | point_set, cls = self.cache[index] 69 | else: 70 | fn = self.datapath[index] 71 | cls = self.classes[self.datapath[index][0]] 72 | cls = np.array([cls]).astype(np.int32) 73 | point_set = np.loadtxt(fn[1],delimiter=',').astype(np.float32) 74 | # Take the first npoints 75 | point_set = point_set[0:self.npoints,:] 76 | if self.normalize: 77 | point_set[:,0:3] = pc_normalize(point_set[:,0:3]) 78 | if not self.normal_channel: 79 | point_set = point_set[:,0:3] 80 | if len(self.cache) < self.cache_size: 81 | self.cache[index] = (point_set, cls) 82 | return point_set, cls 83 | 84 | def __getitem__(self, index): 85 | return self._get_item(index) 86 | 87 | def __len__(self): 88 | return len(self.datapath) 89 | 90 | def num_channel(self): 91 | if self.normal_channel: 92 | return 6 93 | else: 94 | return 3 95 | 96 | def reset(self): 97 | self.idxs = np.arange(0, len(self.datapath)) 98 | if self.shuffle: 99 | np.random.shuffle(self.idxs) 100 | self.num_batches = (len(self.datapath)+self.batch_size-1) // self.batch_size 101 | self.batch_idx = 0 102 | 103 | def has_next_batch(self): 104 | return self.batch_idx < self.num_batches 105 | 106 | def next_batch(self, augment=False): 107 | ''' returned dimension may be smaller than self.batch_size ''' 108 | start_idx = self.batch_idx * self.batch_size 109 | end_idx = min((self.batch_idx+1) * self.batch_size, len(self.datapath)) 110 | bsize = end_idx - start_idx 111 | batch_data = np.zeros((bsize, self.npoints, self.num_channel())) 112 | batch_label = np.zeros((bsize), dtype=np.int32) 113 | for i in range(bsize): 114 | ps,cls = self._get_item(self.idxs[i+start_idx]) 115 | batch_data[i] = ps 116 | batch_label[i] = cls 117 | self.batch_idx += 1 118 | if augment: batch_data = self._augment_batch_data(batch_data) 119 | return batch_data, batch_label 120 | 121 | if __name__ == '__main__': 122 | d = ModelNetDataset(root = '../data/modelnet40_normal_resampled', split='test') 123 | print(d.shuffle) 124 | print(len(d)) 125 | import time 126 | tic = time.time() 127 | for i in range(10): 128 | ps, cls = d[i] 129 | print(time.time() - tic) 130 | print(ps.shape, type(ps), cls) 131 | 132 | print(d.has_next_batch()) 133 | ps_batch, cls_batch = d.next_batch(True) 134 | print(ps_batch.shape) 135 | print(cls_batch.shape) 136 | -------------------------------------------------------------------------------- /models/acnn_cls_rings.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | BASE_DIR = os.path.dirname(__file__) 4 | sys.path.append(BASE_DIR) 5 | sys.path.append(os.path.join(BASE_DIR, '../utils')) 6 | import tensorflow as tf 7 | import numpy as np 8 | import tf_util 9 | from pointnet_util import pointnet_sa_module, acnn_module_rings 10 | 11 | def placeholder_inputs(batch_size, num_point): 12 | pointclouds_pl = tf.placeholder(tf.float32, shape=(batch_size, num_point, 3)) 13 | labels_pl = tf.placeholder(tf.int32, shape=(batch_size)) 14 | normals_pl = tf.placeholder(tf.float32, shape=(batch_size, num_point, 3)) 15 | return pointclouds_pl, labels_pl, normals_pl 16 | 17 | def get_model(point_cloud, normals, is_training, bn_decay=None): 18 | """ Classification A-CNN, input is points BxNx3 and normals BxNx3, output Bx40 """ 19 | batch_size = point_cloud.get_shape()[0].value 20 | num_point = point_cloud.get_shape()[1].value 21 | end_points = {} 22 | 23 | l0_xyz = point_cloud 24 | l0_normals = normals 25 | l0_points = None 26 | 27 | # Abstraction layers 28 | l1_xyz, l1_points, l1_normals = acnn_module_rings(l0_xyz, l0_points, l0_normals, 512, [[0.0, 0.1], [0.1, 0.2]], [16,48], [[32,32,64], [64,64,128]], is_training, bn_decay, scope='layer1') 29 | l2_xyz, l2_points, l2_normals = acnn_module_rings(l1_xyz, l1_points, l1_normals, 128, [[0.1, 0.2], [0.3, 0.4]], [16,48], [[64,64,128], [128,128,256]], is_training, bn_decay, scope='layer2') 30 | _, l3_points, _ = pointnet_sa_module(l2_xyz, l2_points, npoint=None, radius=None, nsample=None, mlp=[256,512,1024], mlp2=None, group_all=True, is_training=is_training, bn_decay=bn_decay, scope='layer3') 31 | 32 | # Fully connected layers 33 | net = tf.reshape(l3_points, [batch_size, -1]) 34 | net = tf_util.fully_connected(net, 512, bn=True, is_training=is_training, scope='fc1', bn_decay=bn_decay) 35 | net = tf_util.dropout(net, keep_prob=0.4, is_training=is_training, scope='dp1') 36 | net = tf_util.fully_connected(net, 256, bn=True, is_training=is_training, scope='fc2', bn_decay=bn_decay) 37 | net = tf_util.dropout(net, keep_prob=0.4, is_training=is_training, scope='dp2') 38 | net = tf_util.fully_connected(net, 40, activation_fn=None, scope='fc3') 39 | 40 | return net, end_points 41 | 42 | 43 | def get_loss(pred, label, end_points): 44 | """ pred: B*NUM_CLASSES, 45 | label: B, """ 46 | loss = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=pred, labels=label) 47 | classify_loss = tf.reduce_mean(loss) 48 | tf.summary.scalar('classify loss', classify_loss) 49 | tf.add_to_collection('losses', classify_loss) 50 | return classify_loss 51 | 52 | 53 | if __name__=='__main__': 54 | with tf.Graph().as_default(): 55 | inputs = tf.zeros((32,1024,3)) 56 | net, _ = get_model(inputs, tf.constant(True)) 57 | print(net) 58 | -------------------------------------------------------------------------------- /models/acnn_segm_rings.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | BASE_DIR = os.path.dirname(__file__) 4 | sys.path.append(BASE_DIR) 5 | sys.path.append(os.path.join(BASE_DIR, '../utils')) 6 | import tensorflow as tf 7 | import numpy as np 8 | import tf_util 9 | from pointnet_util import pointnet_sa_module, pointnet_fp_module, acnn_module_rings 10 | 11 | def placeholder_inputs(batch_size, num_point): 12 | pointclouds_pl = tf.placeholder(tf.float32, shape=(batch_size, num_point, 3)) 13 | labels_pl = tf.placeholder(tf.int32, shape=(batch_size, num_point)) 14 | normals_pl = tf.placeholder(tf.float32, shape=(batch_size, num_point, 3)) 15 | cls_labels_pl = tf.placeholder(tf.int32, shape=(batch_size)) 16 | return pointclouds_pl, labels_pl, normals_pl, cls_labels_pl 17 | 18 | NUM_CATEGORIES = 16 19 | 20 | def get_model(point_cloud, cls_label, normals, is_training, bn_decay=None): 21 | """ Part segmentation A-CNN, input is points BxNx3 and normals BxNx3, output Bx50 """ 22 | batch_size = point_cloud.get_shape()[0].value 23 | num_point = point_cloud.get_shape()[1].value 24 | end_points = {} 25 | 26 | l0_xyz = point_cloud 27 | l0_normals = normals 28 | l0_points = normals 29 | 30 | # Set Abstraction layers 31 | l1_xyz, l1_points, l1_normals = acnn_module_rings(l0_xyz, l0_points, l0_normals, 512, [[0.0, 0.1], [0.1, 0.2]], [16,48], [[32,32,64], [64,64,128]], is_training, bn_decay, scope='layer1') 32 | l2_xyz, l2_points, _ = acnn_module_rings(l1_xyz, l1_points, l1_normals, 128, [[0.1, 0.2], [0.3, 0.4]], [16,48], [[64,64,128], [128,128,256]], is_training, bn_decay, scope='layer2') 33 | l3_xyz, l3_points, _ = pointnet_sa_module(l2_xyz, l2_points, npoint=None, radius=None, nsample=None, mlp=[256,512,1024], mlp2=None, group_all=True, is_training=is_training, bn_decay=bn_decay, scope='layer3') 34 | 35 | cls_label_one_hot = tf.one_hot(cls_label, depth=NUM_CATEGORIES, on_value=1.0, off_value=0.0) 36 | cls_label_one_hot = tf.reshape(cls_label_one_hot, [batch_size, 1, NUM_CATEGORIES]) 37 | cls_label_one_hot = tf.tile(cls_label_one_hot, [1,num_point,1]) 38 | 39 | # Feature Propagation layers 40 | up_l3_points = pointnet_fp_module(l0_xyz, l3_xyz, None, l3_points, [64], is_training, bn_decay, scope='fa_layer1_up') 41 | up_l2_points = pointnet_fp_module(l0_xyz, l2_xyz, None, l2_points, [64], is_training, bn_decay, scope='fa_layer2_up') 42 | up_l1_points = pointnet_fp_module(l0_xyz, l1_xyz, l0_points, l1_points, [128,128,128], is_training, bn_decay, scope='fa_layer3_up') 43 | 44 | concat = tf.concat(axis=-1, values=[ 45 | up_l3_points, 46 | up_l2_points, 47 | up_l1_points, 48 | cls_label_one_hot, 49 | l0_xyz 50 | ]) 51 | 52 | net = tf_util.conv1d(concat, 256, 1, padding='VALID', bn=True, is_training=is_training, scope='fc1', bn_decay=bn_decay) 53 | net = tf_util.dropout(net, keep_prob=0.6, is_training=is_training, scope='dp1') 54 | net = tf_util.conv1d(net, 256, 1, padding='VALID', bn=True, is_training=is_training, scope='fc2', bn_decay=bn_decay) 55 | net = tf_util.dropout(net, keep_prob=0.6, is_training=is_training, scope='dp2') 56 | net = tf_util.conv1d(net, 128, 1, padding='VALID', bn=True, is_training=is_training, scope='fc3', bn_decay=bn_decay) 57 | net = tf_util.conv1d(net, 50, 1, padding='VALID', activation_fn=None, scope='fc4') 58 | 59 | return net, end_points 60 | 61 | 62 | def get_loss(pred, label): 63 | """ pred: BxNxC, 64 | label: BxN, """ 65 | loss = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=pred, labels=label) 66 | classify_loss = tf.reduce_mean(loss) 67 | tf.summary.scalar('classify loss', classify_loss) 68 | tf.add_to_collection('losses', classify_loss) 69 | return classify_loss 70 | 71 | if __name__=='__main__': 72 | with tf.Graph().as_default(): 73 | inputs = tf.zeros((32,2048,6)) 74 | net, _ = get_model(inputs, tf.constant(True)) 75 | print(net) 76 | -------------------------------------------------------------------------------- /part_segm/evaluate.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import math 3 | from datetime import datetime 4 | import h5py 5 | import json 6 | import numpy as np 7 | import tensorflow as tf 8 | import socket 9 | import importlib 10 | import os 11 | import sys 12 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 13 | ROOT_DIR = os.path.dirname(BASE_DIR) 14 | DATA_DIR = os.path.dirname(os.path.dirname(ROOT_DIR)) 15 | sys.path.append(ROOT_DIR) 16 | sys.path.append(os.path.join(ROOT_DIR, 'models')) 17 | sys.path.append(os.path.join(ROOT_DIR, 'utils')) 18 | import provider 19 | import tf_util 20 | import part_dataset_all_normal 21 | 22 | parser = argparse.ArgumentParser() 23 | parser.add_argument('--gpu', type=int, default=0, help='GPU to use [default: GPU 0]') 24 | parser.add_argument('--model', default='acnn_segm_rings', help='Model name [default: pointnet2_part_seg]') 25 | parser.add_argument('--model_path', default='log/model.ckpt', help='model checkpoint file path [default: log/model.ckpt]') 26 | parser.add_argument('--log_dir', default='log/log_eval', help='Log dir [default: log_eval]') 27 | parser.add_argument('--num_point', type=int, default=2048, help='Point Number [default: 2048]') 28 | parser.add_argument('--batch_size', type=int, default=16, help='Batch Size during training [default: 32]') 29 | FLAGS = parser.parse_args() 30 | 31 | VOTE_NUM = 12 32 | EPOCH_CNT = 0 33 | 34 | BATCH_SIZE = FLAGS.batch_size 35 | NUM_POINT = FLAGS.num_point 36 | GPU_INDEX = FLAGS.gpu 37 | 38 | MODEL_PATH = FLAGS.model_path 39 | MODEL = importlib.import_module(FLAGS.model) # import network module 40 | MODEL_FILE = os.path.join(ROOT_DIR, 'models', FLAGS.model+'.py') 41 | LOG_DIR = FLAGS.log_dir 42 | if not os.path.exists(LOG_DIR): os.mkdir(LOG_DIR) 43 | os.system('cp %s %s' % (MODEL_FILE, LOG_DIR)) # bkp of model def 44 | os.system('cp train.py %s' % (LOG_DIR)) # bkp of train procedure 45 | LOG_FOUT = open(os.path.join(LOG_DIR, 'log_train.txt'), 'w') 46 | LOG_FOUT.write(str(FLAGS)+'\n') 47 | NUM_CLASSES = 50 48 | 49 | # Shapenet official train/test split 50 | DATA_PATH = os.path.join(DATA_DIR, 'data', 'shapenetcore_partanno_segmentation_benchmark_v0_normal') 51 | TEST_DATASET = part_dataset_all_normal.PartNormalDataset(root=DATA_PATH, npoints=NUM_POINT, classification=False, split='test', return_cls_label=True) 52 | 53 | def log_string(out_str): 54 | LOG_FOUT.write(out_str+'\n') 55 | LOG_FOUT.flush() 56 | print(out_str) 57 | 58 | def evaluate(): 59 | with tf.Graph().as_default(): 60 | with tf.device('/gpu:'+str(GPU_INDEX)): 61 | pointclouds_pl, labels_pl, normals_pl, cls_labels_pl = MODEL.placeholder_inputs(BATCH_SIZE, NUM_POINT) 62 | is_training_pl = tf.placeholder(tf.bool, shape=()) 63 | print(is_training_pl) 64 | 65 | print("--- Get model and loss") 66 | # pred, end_points = MODEL.get_model(pointclouds_pl, is_training_pl) 67 | pred, end_points = MODEL.get_model(pointclouds_pl, cls_labels_pl, normals_pl, is_training_pl) 68 | loss = MODEL.get_loss(pred, labels_pl) 69 | saver = tf.train.Saver() 70 | 71 | # Create a session 72 | config = tf.ConfigProto() 73 | config.gpu_options.allow_growth = True 74 | config.allow_soft_placement = True 75 | sess = tf.Session(config=config) 76 | # Restore variables from disk. 77 | saver.restore(sess, MODEL_PATH) 78 | ops = {'pointclouds_pl': pointclouds_pl, 79 | 'normals_pl': normals_pl, 80 | 'cls_labels_pl': cls_labels_pl, 81 | 'labels_pl': labels_pl, 82 | 'is_training_pl': is_training_pl, 83 | 'pred': pred, 84 | 'loss': loss} 85 | 86 | eval_one_epoch(sess, ops) 87 | 88 | def get_batch(dataset, idxs, start_idx, end_idx): 89 | bsize = end_idx-start_idx 90 | batch_data = np.zeros((bsize, NUM_POINT, 6)) 91 | batch_label = np.zeros((bsize, NUM_POINT), dtype=np.int32) 92 | batch_cls_label = np.zeros((bsize,), dtype=np.int32) 93 | for i in range(bsize): 94 | ps,normal,seg,cls = dataset[idxs[i+start_idx]] 95 | batch_data[i,:,0:3] = ps 96 | batch_data[i,:,3:6] = normal 97 | batch_label[i,:] = seg 98 | batch_cls_label[i] = cls 99 | return batch_data, batch_label, batch_cls_label 100 | 101 | def eval_one_epoch(sess, ops): 102 | """ ops: dict mapping from string to tf ops """ 103 | is_training = False 104 | test_idxs = np.arange(0, len(TEST_DATASET)) 105 | # Test on all data: last batch might be smaller than BATCH_SIZE 106 | num_batches = (len(TEST_DATASET)+BATCH_SIZE-1)//BATCH_SIZE 107 | 108 | cur_batch_data = np.zeros((BATCH_SIZE,NUM_POINT,3)) 109 | cur_batch_normals = np.zeros((BATCH_SIZE,NUM_POINT,3)) 110 | cur_batch_label = np.zeros((BATCH_SIZE,NUM_POINT), dtype=np.int32) 111 | cur_batch_cls_label = np.zeros((BATCH_SIZE), dtype=np.int32) 112 | 113 | total_correct = 0 114 | total_seen = 0 115 | loss_sum = 0 116 | total_seen_class = [0 for _ in range(NUM_CLASSES)] 117 | total_correct_class = [0 for _ in range(NUM_CLASSES)] 118 | 119 | seg_classes = TEST_DATASET.seg_classes 120 | shape_ious = {cat:[] for cat in seg_classes.keys()} 121 | seg_label_to_cat = {} # {0:Airplane, 1:Airplane, ...49:Table} 122 | for cat in seg_classes.keys(): 123 | for label in seg_classes[cat]: 124 | seg_label_to_cat[label] = cat 125 | 126 | log_string(str(datetime.now())) 127 | log_string('---- EPOCH %03d EVALUATION ----'%(EPOCH_CNT)) 128 | 129 | batch_data = np.zeros((BATCH_SIZE, NUM_POINT, 6)) 130 | batch_label = np.zeros((BATCH_SIZE, NUM_POINT)).astype(np.int32) 131 | batch_cls_label = np.zeros((BATCH_SIZE,)).astype(np.int32) 132 | for batch_idx in range(num_batches): 133 | if batch_idx %50==0: 134 | log_string('%03d/%03d'%(batch_idx, num_batches)) 135 | start_idx = batch_idx * BATCH_SIZE 136 | end_idx = min(len(TEST_DATASET), (batch_idx+1) * BATCH_SIZE) 137 | cur_batch_size = end_idx-start_idx 138 | batch_data, batch_label, batch_cls_label = get_batch(TEST_DATASET, test_idxs, start_idx, end_idx) 139 | 140 | cur_batch_data[0:cur_batch_size,...] = batch_data[:,:,:3] 141 | cur_batch_normals[0:cur_batch_size,...] = batch_data[:,:,3:] 142 | cur_batch_label[0:cur_batch_size,...] = batch_label 143 | cur_batch_cls_label[0:cur_batch_size] = batch_cls_label 144 | 145 | loss_val = 0 146 | pred_val = np.zeros((BATCH_SIZE, NUM_POINT, NUM_CLASSES)) 147 | for _ in range(VOTE_NUM): 148 | jittered_data = provider.jitter_point_cloud(cur_batch_data) 149 | feed_dict = {ops['pointclouds_pl']: jittered_data, 150 | ops['labels_pl']: cur_batch_label, 151 | ops['cls_labels_pl']: cur_batch_cls_label, 152 | ops['normals_pl']: cur_batch_normals, 153 | ops['is_training_pl']: is_training} 154 | temp_loss_val, temp_pred_val = sess.run([ops['loss'], ops['pred']], feed_dict=feed_dict) 155 | loss_val += temp_loss_val 156 | pred_val += temp_pred_val 157 | loss_val /= float(VOTE_NUM) 158 | # --------------------------------------------------------------------- 159 | 160 | # Select valid data 161 | cur_pred_val = pred_val[0:cur_batch_size] 162 | # Constrain pred to the groundtruth classes (selected by seg_classes[cat]) 163 | cur_pred_val_logits = cur_pred_val 164 | cur_pred_val = np.zeros((cur_batch_size, NUM_POINT)).astype(np.int32) 165 | for i in range(cur_batch_size): 166 | cat = seg_label_to_cat[cur_batch_label[i,0]] 167 | logits = cur_pred_val_logits[i,:,:] 168 | cur_pred_val[i,:] = np.argmax(logits[:,seg_classes[cat]], 1) + seg_classes[cat][0] 169 | correct = np.sum(cur_pred_val == cur_batch_label) 170 | total_correct += correct 171 | total_seen += (cur_batch_size*NUM_POINT) 172 | if cur_batch_size==BATCH_SIZE: 173 | loss_sum += loss_val 174 | for l in range(NUM_CLASSES): 175 | total_seen_class[l] += np.sum(cur_batch_label==l) 176 | total_correct_class[l] += (np.sum((cur_pred_val==l) & (batch_label==l))) 177 | 178 | for i in range(cur_batch_size): 179 | segp = cur_pred_val[i,:] 180 | segl = cur_batch_label[i,:] 181 | pts = batch_data[i,:,:3] 182 | cat = seg_label_to_cat[segl[0]] 183 | part_ious = [0.0 for _ in range(len(seg_classes[cat]))] 184 | for l in seg_classes[cat]: 185 | if (np.sum(segl==l) == 0) and (np.sum(segp==l) == 0): # part is not present, no prediction as well 186 | part_ious[l-seg_classes[cat][0]] = 1.0 187 | else: 188 | part_ious[l-seg_classes[cat][0]] = np.sum((segl==l) & (segp==l)) / float(np.sum((segl==l) | (segp==l))) 189 | shape_ious[cat].append(np.mean(part_ious)) 190 | 191 | all_shape_ious = [] 192 | for cat in shape_ious.keys(): 193 | for iou in shape_ious[cat]: 194 | all_shape_ious.append(iou) 195 | shape_ious[cat] = np.mean(shape_ious[cat]) 196 | print(len(all_shape_ious)) 197 | mean_shape_ious = np.mean(list(shape_ious.values())) 198 | log_string('eval mean loss: %f' % (loss_sum / float(len(TEST_DATASET)/BATCH_SIZE))) 199 | log_string('eval accuracy: %f'% (total_correct / float(total_seen))) 200 | log_string('eval avg class acc: %f' % (np.mean(np.array(total_correct_class)/np.array(total_seen_class,dtype=np.float)))) 201 | for cat in sorted(shape_ious.keys()): 202 | log_string('eval mIoU of %s:\t %f'%(cat, shape_ious[cat])) 203 | log_string('eval mean mIoU: %f' % (mean_shape_ious)) 204 | log_string('eval mean mIoU (all shapes): %f' % (np.mean(all_shape_ious))) 205 | 206 | if __name__ == "__main__": 207 | log_string('pid: %s'%(str(os.getpid()))) 208 | evaluate() 209 | LOG_FOUT.close() 210 | -------------------------------------------------------------------------------- /part_segm/evaluate_job.sh: -------------------------------------------------------------------------------- 1 | python evaluate.py --model_path log/model.ckpt --log_dir log/log_eval 2 | # python evaluate.py --model_path log/model_best_acc_inst.ckpt --log_dir log/log_eval_inst 3 | # python evaluate.py --model_path log/model_best_acc_cat.ckpt --log_dir log/log_eval_cat -------------------------------------------------------------------------------- /part_segm/part_dataset_all_normal.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Dataset for ShapeNetPart segmentation 3 | ''' 4 | 5 | import os 6 | import os.path 7 | import json 8 | import numpy as np 9 | import sys 10 | 11 | def pc_normalize(pc): 12 | l = pc.shape[0] 13 | centroid = np.mean(pc, axis=0) 14 | pc = pc - centroid 15 | m = np.max(np.sqrt(np.sum(pc**2, axis=1))) 16 | pc = pc / m 17 | return pc 18 | 19 | class PartNormalDataset(): 20 | def __init__(self, root, npoints = 2500, classification = False, split='train', normalize=True, return_cls_label = False): 21 | self.npoints = npoints 22 | self.root = root 23 | self.catfile = os.path.join(self.root, 'synsetoffset2category.txt') 24 | self.cat = {} 25 | 26 | self.classification = classification 27 | self.normalize = normalize 28 | self.return_cls_label = return_cls_label 29 | 30 | with open(self.catfile, 'r') as f: 31 | for line in f: 32 | ls = line.strip().split() 33 | self.cat[ls[0]] = ls[1] 34 | self.cat = {k:v for k,v in self.cat.items()} 35 | #print(self.cat) 36 | 37 | self.meta = {} 38 | with open(os.path.join(self.root, 'train_test_split', 'shuffled_train_file_list.json'), 'r') as f: 39 | train_ids = set([str(d.split('/')[2]) for d in json.load(f)]) 40 | with open(os.path.join(self.root, 'train_test_split', 'shuffled_val_file_list.json'), 'r') as f: 41 | val_ids = set([str(d.split('/')[2]) for d in json.load(f)]) 42 | with open(os.path.join(self.root, 'train_test_split', 'shuffled_test_file_list.json'), 'r') as f: 43 | test_ids = set([str(d.split('/')[2]) for d in json.load(f)]) 44 | for item in self.cat: 45 | #print('category', item) 46 | self.meta[item] = [] 47 | dir_point = os.path.join(self.root, self.cat[item]) 48 | fns = sorted(os.listdir(dir_point)) 49 | #print(fns[0][0:-4]) 50 | if split=='trainval': 51 | fns = [fn for fn in fns if ((fn[0:-4] in train_ids) or (fn[0:-4] in val_ids))] 52 | elif split=='train': 53 | fns = [fn for fn in fns if fn[0:-4] in train_ids] 54 | elif split=='val': 55 | fns = [fn for fn in fns if fn[0:-4] in val_ids] 56 | elif split=='test': 57 | fns = [fn for fn in fns if fn[0:-4] in test_ids] 58 | else: 59 | print('Unknown split: %s. Exiting..'%(split)) 60 | exit(-1) 61 | 62 | #print(os.path.basename(fns)) 63 | for fn in fns: 64 | token = (os.path.splitext(os.path.basename(fn))[0]) 65 | self.meta[item].append(os.path.join(dir_point, token + '.txt')) 66 | 67 | self.datapath = [] 68 | for item in self.cat: 69 | for fn in self.meta[item]: 70 | self.datapath.append((item, fn)) 71 | 72 | 73 | self.classes = dict(zip(self.cat, range(len(self.cat)))) 74 | # Mapping from category ('Chair') to a list of int [10,11,12,13] as segmentation labels 75 | self.seg_classes = {'Earphone': [16, 17, 18], 'Motorbike': [30, 31, 32, 33, 34, 35], 'Rocket': [41, 42, 43], 'Car': [8, 9, 10, 11], 'Laptop': [28, 29], 'Cap': [6, 7], 'Skateboard': [44, 45, 46], 'Mug': [36, 37], 'Guitar': [19, 20, 21], 'Bag': [4, 5], 'Lamp': [24, 25, 26, 27], 'Table': [47, 48, 49], 'Airplane': [0, 1, 2, 3], 'Pistol': [38, 39, 40], 'Chair': [12, 13, 14, 15], 'Knife': [22, 23]} 76 | 77 | for cat in sorted(self.seg_classes.keys()): 78 | print(cat, self.seg_classes[cat]) 79 | 80 | self.cache = {} # from index to (point_set, cls, seg) tuple 81 | self.cache_size = 20000 82 | 83 | def __getitem__(self, index): 84 | if index in self.cache: 85 | point_set, normal, seg, cls = self.cache[index] 86 | else: 87 | fn = self.datapath[index] 88 | cat = self.datapath[index][0] 89 | cls = self.classes[cat] 90 | cls = np.array([cls]).astype(np.int32) 91 | data = np.loadtxt(fn[1]).astype(np.float32) 92 | point_set = data[:,0:3] 93 | if self.normalize: 94 | point_set = pc_normalize(point_set) 95 | normal = data[:,3:6] 96 | seg = data[:,-1].astype(np.int32) 97 | if len(self.cache) < self.cache_size: 98 | self.cache[index] = (point_set, normal, seg, cls) 99 | 100 | 101 | choice = np.random.choice(len(seg), self.npoints, replace=True) 102 | #resample 103 | point_set = point_set[choice, :] 104 | seg = seg[choice] 105 | normal = normal[choice,:] 106 | if self.classification: 107 | return point_set, normal, cls 108 | else: 109 | if self.return_cls_label: 110 | return point_set, normal, seg, cls 111 | else: 112 | return point_set, normal, seg 113 | 114 | def __len__(self): 115 | return len(self.datapath) 116 | 117 | 118 | if __name__ == '__main__': 119 | d = PartNormalDataset(root = '../data/shapenetcore_partanno_segmentation_benchmark_v0_normal', split='trainval', npoints=3000) 120 | print(len(d)) 121 | 122 | i = 500 123 | ps, normal, seg = d[i] 124 | print(d.datapath[i]) 125 | print(np.max(seg), np.min(seg)) 126 | print(ps.shape, seg.shape, normal.shape) 127 | print(ps) 128 | print(normal) 129 | 130 | sys.path.append('../utils') 131 | import show3d_balls 132 | show3d_balls.showpoints(ps, normal+1, ballradius=8) 133 | 134 | d = PartNormalDataset(root = '../data/shapenetcore_partanno_segmentation_benchmark_v0_normal', classification = True) 135 | print(len(d)) 136 | ps, normal, cls = d[0] 137 | print(ps.shape, type(ps), cls.shape,type(cls)) 138 | -------------------------------------------------------------------------------- /part_segm/train.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import math 3 | from datetime import datetime 4 | import h5py 5 | import numpy as np 6 | import tensorflow as tf 7 | import socket 8 | import importlib 9 | import os 10 | import sys 11 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 12 | ROOT_DIR = os.path.dirname(BASE_DIR) 13 | DATA_DIR = os.path.dirname(os.path.dirname(ROOT_DIR)) 14 | sys.path.append(ROOT_DIR) 15 | sys.path.append(os.path.join(ROOT_DIR, 'models')) 16 | sys.path.append(os.path.join(ROOT_DIR, 'utils')) 17 | import provider 18 | import tf_util 19 | import part_dataset_all_normal 20 | 21 | parser = argparse.ArgumentParser() 22 | parser.add_argument('--gpu', type=int, default=0, help='GPU to use [default: GPU 0]') 23 | parser.add_argument('--model', default='acnn_segm_rings', help='Model name [default: model]') 24 | parser.add_argument('--log_dir', default='log', help='Log dir [default: log]') 25 | parser.add_argument('--num_point', type=int, default=2048, help='Point Number [default: 2048]') 26 | parser.add_argument('--max_epoch', type=int, default=201, help='Epoch to run [default: 201]') 27 | parser.add_argument('--batch_size', type=int, default=16, help='Batch Size during training [default: 32]') 28 | parser.add_argument('--learning_rate', type=float, default=0.001, help='Initial learning rate [default: 0.001]') 29 | parser.add_argument('--momentum', type=float, default=0.9, help='Initial learning rate [default: 0.9]') 30 | parser.add_argument('--optimizer', default='adam', help='adam or momentum [default: adam]') 31 | parser.add_argument('--decay_step', type=int, default=16881*20, help='Decay step for lr decay [default: 200000]') 32 | parser.add_argument('--decay_rate', type=float, default=0.5, help='Decay rate for lr decay [default: 0.7]') 33 | FLAGS = parser.parse_args() 34 | 35 | EPOCH_CNT = 0 36 | 37 | BATCH_SIZE = FLAGS.batch_size 38 | NUM_POINT = FLAGS.num_point 39 | MAX_EPOCH = FLAGS.max_epoch 40 | BASE_LEARNING_RATE = FLAGS.learning_rate 41 | GPU_INDEX = FLAGS.gpu 42 | MOMENTUM = FLAGS.momentum 43 | OPTIMIZER = FLAGS.optimizer 44 | DECAY_STEP = FLAGS.decay_step 45 | DECAY_RATE = FLAGS.decay_rate 46 | 47 | MODEL = importlib.import_module(FLAGS.model) # import network module 48 | MODEL_FILE = os.path.join(ROOT_DIR, 'models', FLAGS.model+'.py') 49 | LOG_DIR = FLAGS.log_dir 50 | if not os.path.exists(LOG_DIR): os.mkdir(LOG_DIR) 51 | os.system('cp %s %s' % (MODEL_FILE, LOG_DIR)) # bkp of model def 52 | os.system('cp train.py %s' % (LOG_DIR)) # bkp of train procedure 53 | LOG_FOUT = open(os.path.join(LOG_DIR, 'log_train.txt'), 'w') 54 | LOG_FOUT.write(str(FLAGS)+'\n') 55 | 56 | BN_INIT_DECAY = 0.5 57 | BN_DECAY_DECAY_RATE = 0.5 58 | BN_DECAY_DECAY_STEP = float(DECAY_STEP) 59 | BN_DECAY_CLIP = 0.99 60 | 61 | HOSTNAME = socket.gethostname() 62 | 63 | NUM_CLASSES = 50 64 | 65 | # Shapenet official train/test split 66 | DATA_PATH = os.path.join(DATA_DIR, 'data', 'shapenetcore_partanno_segmentation_benchmark_v0_normal') 67 | TRAIN_DATASET = part_dataset_all_normal.PartNormalDataset(root=DATA_PATH, npoints=NUM_POINT, classification=False, split='trainval', return_cls_label=True) 68 | TEST_DATASET = part_dataset_all_normal.PartNormalDataset(root=DATA_PATH, npoints=NUM_POINT, classification=False, split='test', return_cls_label=True) 69 | 70 | def log_string(out_str): 71 | LOG_FOUT.write(out_str+'\n') 72 | LOG_FOUT.flush() 73 | print(out_str) 74 | 75 | def get_learning_rate(batch): 76 | learning_rate = tf.train.exponential_decay( 77 | BASE_LEARNING_RATE, # Base learning rate. 78 | batch * BATCH_SIZE, # Current index into the dataset. 79 | DECAY_STEP, # Decay step. 80 | DECAY_RATE, # Decay rate. 81 | staircase=True) 82 | learning_rate = tf.maximum(learning_rate, 0.00001) # CLIP THE LEARNING RATE! 83 | return learning_rate 84 | 85 | def get_bn_decay(batch): 86 | bn_momentum = tf.train.exponential_decay( 87 | BN_INIT_DECAY, 88 | batch*BATCH_SIZE, 89 | BN_DECAY_DECAY_STEP, 90 | BN_DECAY_DECAY_RATE, 91 | staircase=True) 92 | bn_decay = tf.minimum(BN_DECAY_CLIP, 1 - bn_momentum) 93 | return bn_decay 94 | 95 | def train(): 96 | with tf.Graph().as_default(): 97 | with tf.device('/gpu:'+str(GPU_INDEX)): 98 | pointclouds_pl, labels_pl, normals_pl, cls_labels_pl = MODEL.placeholder_inputs(BATCH_SIZE, NUM_POINT) 99 | is_training_pl = tf.placeholder(tf.bool, shape=()) 100 | 101 | # Note the global_step=batch parameter to minimize. 102 | # That tells the optimizer to helpfully increment the 'batch' parameter for you every time it trains. 103 | batch = tf.Variable(0) 104 | bn_decay = get_bn_decay(batch) 105 | tf.summary.scalar('bn_decay', bn_decay) 106 | 107 | print("--- Get model and loss") 108 | # Get model and loss 109 | pred, end_points = MODEL.get_model(pointclouds_pl, cls_labels_pl, normals_pl, is_training_pl, bn_decay=bn_decay) 110 | loss = MODEL.get_loss(pred, labels_pl) 111 | tf.summary.scalar('loss', loss) 112 | 113 | correct = tf.equal(tf.argmax(pred, 2), tf.to_int64(labels_pl)) 114 | accuracy = tf.reduce_sum(tf.cast(correct, tf.float32)) / float(BATCH_SIZE*NUM_POINT) 115 | tf.summary.scalar('accuracy', accuracy) 116 | 117 | print("--- Get training operator") 118 | # Get training operator 119 | learning_rate = get_learning_rate(batch) 120 | tf.summary.scalar('learning_rate', learning_rate) 121 | if OPTIMIZER == 'momentum': 122 | optimizer = tf.train.MomentumOptimizer(learning_rate, momentum=MOMENTUM) 123 | elif OPTIMIZER == 'adam': 124 | optimizer = tf.train.AdamOptimizer(learning_rate) 125 | train_op = optimizer.minimize(loss, global_step=batch) 126 | 127 | # Add ops to save and restore all the variables. 128 | saver = tf.train.Saver() 129 | 130 | # Create a session 131 | config = tf.ConfigProto() 132 | config.gpu_options.allow_growth = True 133 | config.allow_soft_placement = True 134 | config.log_device_placement = False 135 | sess = tf.Session(config=config) 136 | 137 | # Add summary writers 138 | merged = tf.summary.merge_all() 139 | train_writer = tf.summary.FileWriter(os.path.join(LOG_DIR, 'train'), sess.graph) 140 | test_writer = tf.summary.FileWriter(os.path.join(LOG_DIR, 'test'), sess.graph) 141 | 142 | # Init variables 143 | init = tf.global_variables_initializer() 144 | sess.run(init) 145 | 146 | ops = {'pointclouds_pl': pointclouds_pl, 147 | 'normals_pl': normals_pl, 148 | 'labels_pl': labels_pl, 149 | 'cls_labels_pl': cls_labels_pl, 150 | 'is_training_pl': is_training_pl, 151 | 'pred': pred, 152 | 'loss': loss, 153 | 'train_op': train_op, 154 | 'merged': merged, 155 | 'step': batch, 156 | 'end_points': end_points} 157 | 158 | best_acc_cat = -1 159 | best_acc_inst = -1 160 | for epoch in range(MAX_EPOCH): 161 | log_string('**** EPOCH %03d ****' % (epoch)) 162 | sys.stdout.flush() 163 | 164 | train_one_epoch(sess, ops, train_writer) 165 | _, acc_cat, acc_inst = eval_one_epoch(sess, ops, test_writer) 166 | 167 | if acc_cat > best_acc_cat: 168 | best_acc_cat = acc_cat 169 | save_path = saver.save(sess, os.path.join(LOG_DIR, "model_best_acc_cat.ckpt")) 170 | log_string("Model saved in file: %s" % save_path) 171 | log_string("Best category accuracy: %f" % best_acc_cat) 172 | 173 | if acc_inst > best_acc_inst: 174 | best_acc_inst = acc_inst 175 | save_path = saver.save(sess, os.path.join(LOG_DIR, "model_best_acc_inst.ckpt")) 176 | log_string("Model saved in file: %s" % save_path) 177 | log_string("Best instance accuracy: %f" % best_acc_inst) 178 | 179 | # Save the variables to disk. 180 | if epoch % 10 == 0: 181 | save_path = saver.save(sess, os.path.join(LOG_DIR, "model.ckpt")) 182 | log_string("Model saved in file: %s" % save_path) 183 | 184 | def get_batch(dataset, idxs, start_idx, end_idx): 185 | bsize = end_idx-start_idx 186 | batch_data = np.zeros((bsize, NUM_POINT, 6)) 187 | batch_label = np.zeros((bsize, NUM_POINT), dtype=np.int32) 188 | batch_cls_label = np.zeros((bsize,), dtype=np.int32) 189 | for i in range(bsize): 190 | ps,normal,seg,cls = dataset[idxs[i+start_idx]] 191 | batch_data[i,:,0:3] = ps 192 | batch_data[i,:,3:6] = normal 193 | batch_label[i,:] = seg 194 | batch_cls_label[i] = cls 195 | return batch_data, batch_label, batch_cls_label 196 | 197 | def train_one_epoch(sess, ops, train_writer): 198 | """ ops: dict mapping from string to tf ops """ 199 | is_training = True 200 | 201 | # Shuffle train samples 202 | train_idxs = np.arange(0, len(TRAIN_DATASET)) 203 | np.random.shuffle(train_idxs) 204 | num_batches = len(TRAIN_DATASET) // BATCH_SIZE 205 | 206 | log_string(str(datetime.now())) 207 | 208 | cur_batch_data = np.zeros((BATCH_SIZE,NUM_POINT,3)) 209 | cur_batch_normals = np.zeros((BATCH_SIZE,NUM_POINT,3)) 210 | cur_batch_label = np.zeros((BATCH_SIZE,NUM_POINT), dtype=np.int32) 211 | 212 | total_correct = 0 213 | total_seen = 0 214 | loss_sum = 0 215 | 216 | for batch_idx in range(num_batches): 217 | start_idx = batch_idx * BATCH_SIZE 218 | end_idx = (batch_idx+1) * BATCH_SIZE 219 | batch_data, batch_label, batch_cls_label = get_batch(TRAIN_DATASET, train_idxs, start_idx, end_idx) 220 | 221 | # data augmentation step 222 | batch_data[:,:,0:3] = provider.jitter_point_cloud(batch_data[:,:,0:3]) 223 | batch_data, batch_label = provider.shuffle_points_with_labels(batch_data, batch_label) 224 | 225 | bsize = batch_data.shape[0] 226 | cur_batch_data[0:bsize,...] = batch_data[:,:,:3] 227 | cur_batch_normals[0:bsize,...] = batch_data[:,:,3:] 228 | cur_batch_label[0:bsize,...] = batch_label 229 | 230 | feed_dict = {ops['pointclouds_pl']: cur_batch_data, 231 | ops['labels_pl']: cur_batch_label, 232 | ops['normals_pl']: cur_batch_normals, 233 | ops['cls_labels_pl']: batch_cls_label, 234 | ops['is_training_pl']: is_training,} 235 | summary, step, _, loss_val, pred_val = sess.run([ops['merged'], ops['step'], 236 | ops['train_op'], ops['loss'], ops['pred']], feed_dict=feed_dict) 237 | train_writer.add_summary(summary, step) 238 | pred_val = np.argmax(pred_val, 2) 239 | correct = np.sum(pred_val == batch_label) 240 | total_correct += correct 241 | total_seen += (BATCH_SIZE*NUM_POINT) 242 | loss_sum += loss_val 243 | 244 | if (batch_idx+1)%50 == 0: 245 | log_string(' -- %03d / %03d --' % (batch_idx+1, num_batches)) 246 | log_string('mean loss: %f' % (loss_sum / 50)) 247 | log_string('accuracy: %f' % (total_correct / float(total_seen))) 248 | total_correct = 0 249 | total_seen = 0 250 | loss_sum = 0 251 | 252 | def eval_one_epoch(sess, ops, test_writer): 253 | """ ops: dict mapping from string to tf ops """ 254 | global EPOCH_CNT 255 | is_training = False 256 | test_idxs = np.arange(0, len(TEST_DATASET)) 257 | # Test on all data: last batch might be smaller than BATCH_SIZE 258 | num_batches = (len(TEST_DATASET)+BATCH_SIZE-1) // BATCH_SIZE 259 | 260 | cur_batch_data = np.zeros((BATCH_SIZE,NUM_POINT,3)) 261 | cur_batch_normals = np.zeros((BATCH_SIZE,NUM_POINT,3)) 262 | cur_batch_label = np.zeros((BATCH_SIZE,NUM_POINT), dtype=np.int32) 263 | cur_batch_cls_label = np.zeros((BATCH_SIZE), dtype=np.int32) 264 | 265 | total_correct = 0 266 | total_seen = 0 267 | loss_sum = 0 268 | total_seen_class = [0 for _ in range(NUM_CLASSES)] 269 | total_correct_class = [0 for _ in range(NUM_CLASSES)] 270 | 271 | seg_classes = TEST_DATASET.seg_classes 272 | shape_ious = {cat:[] for cat in seg_classes.keys()} 273 | seg_label_to_cat = {} # {0:Airplane, 1:Airplane, ...49:Table} 274 | for cat in seg_classes.keys(): 275 | for label in seg_classes[cat]: 276 | seg_label_to_cat[label] = cat 277 | 278 | log_string(str(datetime.now())) 279 | log_string('---- EPOCH %03d EVALUATION ----'%(EPOCH_CNT)) 280 | 281 | for batch_idx in range(num_batches): 282 | if batch_idx %50==0: 283 | log_string('%03d/%03d'%(batch_idx, num_batches)) 284 | start_idx = batch_idx * BATCH_SIZE 285 | end_idx = min(len(TEST_DATASET), (batch_idx+1) * BATCH_SIZE) 286 | cur_batch_size = end_idx-start_idx 287 | batch_data, batch_label, batch_cls_label = get_batch(TEST_DATASET, test_idxs, start_idx, end_idx) 288 | 289 | cur_batch_data[0:cur_batch_size,...] = batch_data[:,:,:3] 290 | cur_batch_normals[0:cur_batch_size,...] = batch_data[:,:,3:] 291 | cur_batch_label[0:cur_batch_size,...] = batch_label 292 | cur_batch_cls_label[0:cur_batch_size] = batch_cls_label 293 | 294 | # --------------------------------------------------------------------- 295 | feed_dict = {ops['pointclouds_pl']: cur_batch_data, 296 | ops['labels_pl']: cur_batch_label, 297 | ops['cls_labels_pl']: cur_batch_cls_label, 298 | ops['normals_pl']: cur_batch_normals, 299 | ops['is_training_pl']: is_training} 300 | summary, step, loss_val, pred_val = sess.run([ops['merged'], ops['step'], 301 | ops['loss'], ops['pred']], feed_dict=feed_dict) 302 | test_writer.add_summary(summary, step) 303 | # --------------------------------------------------------------------- 304 | 305 | # Select valid data 306 | cur_pred_val = pred_val[0:cur_batch_size] 307 | # Constrain pred to the groundtruth classes (selected by seg_classes[cat]) 308 | cur_pred_val_logits = cur_pred_val 309 | cur_pred_val = np.zeros((cur_batch_size, NUM_POINT)).astype(np.int32) 310 | for i in range(cur_batch_size): 311 | cat = seg_label_to_cat[cur_batch_label[i,0]] 312 | logits = cur_pred_val_logits[i,:,:] 313 | cur_pred_val[i,:] = np.argmax(logits[:,seg_classes[cat]], 1) + seg_classes[cat][0] 314 | correct = np.sum(cur_pred_val == cur_batch_label) 315 | total_correct += correct 316 | total_seen += (cur_batch_size*NUM_POINT) 317 | if cur_batch_size==BATCH_SIZE: 318 | loss_sum += loss_val 319 | for l in range(NUM_CLASSES): 320 | total_seen_class[l] += np.sum(cur_batch_label==l) 321 | total_correct_class[l] += (np.sum((cur_pred_val==l) & (batch_label==l))) 322 | 323 | for i in range(cur_batch_size): 324 | segp = cur_pred_val[i,:] 325 | segl = cur_batch_label[i,:] 326 | cat = seg_label_to_cat[segl[0]] 327 | part_ious = [0.0 for _ in range(len(seg_classes[cat]))] 328 | for l in seg_classes[cat]: 329 | if (np.sum(segl==l) == 0) and (np.sum(segp==l) == 0): # part is not present, no prediction as well 330 | part_ious[l-seg_classes[cat][0]] = 1.0 331 | else: 332 | part_ious[l-seg_classes[cat][0]] = np.sum((segl==l) & (segp==l)) / float(np.sum((segl==l) | (segp==l))) 333 | shape_ious[cat].append(np.mean(part_ious)) 334 | 335 | all_shape_ious = [] 336 | for cat in shape_ious.keys(): 337 | for iou in shape_ious[cat]: 338 | all_shape_ious.append(iou) 339 | shape_ious[cat] = np.mean(shape_ious[cat]) 340 | mean_shape_ious = np.mean(list(shape_ious.values())) 341 | log_string('eval mean loss: %f' % (loss_sum / float(len(TEST_DATASET)/BATCH_SIZE))) 342 | log_string('eval accuracy: %f'% (total_correct / float(total_seen))) 343 | log_string('eval avg class acc: %f' % (np.mean(np.array(total_correct_class)/np.array(total_seen_class,dtype=np.float)))) 344 | for cat in sorted(shape_ious.keys()): 345 | log_string('eval mIoU of %s:\t %f'%(cat, shape_ious[cat])) 346 | log_string('eval mean mIoU: %f' % (mean_shape_ious)) 347 | log_string('eval mean mIoU (all shapes): %f' % (np.mean(all_shape_ious))) 348 | 349 | EPOCH_CNT += 1 350 | return total_correct/float(total_seen), mean_shape_ious, np.mean(all_shape_ious) 351 | 352 | 353 | if __name__ == "__main__": 354 | log_string('pid: %s'%(str(os.getpid()))) 355 | train() 356 | LOG_FOUT.close() 357 | -------------------------------------------------------------------------------- /pics/teaser.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/akomarichev/a-cnn/0cce386f8f423cf886a5ecb45c4df8fe8cdbbf6b/pics/teaser.png -------------------------------------------------------------------------------- /provider.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | import numpy as np 4 | import h5py 5 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 6 | sys.path.append(BASE_DIR) 7 | 8 | def shuffle_data(data, labels): 9 | """ Shuffle data and labels. 10 | Input: 11 | data: B,N,... numpy array 12 | label: B,... numpy array 13 | Return: 14 | shuffled data, label and shuffle indices 15 | """ 16 | idx = np.arange(len(labels)) 17 | np.random.shuffle(idx) 18 | return data[idx, ...], labels[idx], idx 19 | 20 | def shuffle_points(batch_data): 21 | """ Shuffle orders of points in each point cloud -- changes FPS behavior. 22 | Use the same shuffling idx for the entire batch. 23 | Input: 24 | BxNxC array 25 | Output: 26 | BxNxC array 27 | """ 28 | idx = np.arange(batch_data.shape[1]) 29 | np.random.shuffle(idx) 30 | return batch_data[:,idx,:] 31 | 32 | def shuffle_points_with_labels(batch_data, labels): 33 | """ Shuffle orders of points in each point cloud -- changes FPS behavior. 34 | Use the same shuffling idx for the entire batch. 35 | Input: 36 | BxNxC array 37 | Output: 38 | BxNxC array 39 | """ 40 | idx = np.arange(batch_data.shape[1]) 41 | np.random.shuffle(idx) 42 | return batch_data[:,idx,:], labels[:,idx] 43 | 44 | def rotate_point_cloud(batch_data): 45 | """ Randomly rotate the point clouds to augument the dataset 46 | rotation is per shape based along up direction 47 | Input: 48 | BxNx3 array, original batch of point clouds 49 | Return: 50 | BxNx3 array, rotated batch of point clouds 51 | """ 52 | rotated_data = np.zeros(batch_data.shape, dtype=np.float32) 53 | for k in range(batch_data.shape[0]): 54 | rotation_angle = np.random.uniform() * 2 * np.pi 55 | cosval = np.cos(rotation_angle) 56 | sinval = np.sin(rotation_angle) 57 | rotation_matrix = np.array([[cosval, 0, sinval], 58 | [0, 1, 0], 59 | [-sinval, 0, cosval]]) 60 | shape_pc = batch_data[k, ...] 61 | rotated_data[k, ...] = np.dot(shape_pc.reshape((-1, 3)), rotation_matrix) 62 | return rotated_data 63 | 64 | 65 | def rotate_point_cloud_with_normal(batch_xyz_normal): 66 | ''' Randomly rotate XYZ, normal point cloud. 67 | Input: 68 | batch_xyz_normal: B,N,6, first three channels are XYZ, last 3 all normal 69 | Output: 70 | B,N,6, rotated XYZ, normal point cloud 71 | ''' 72 | for k in range(batch_xyz_normal.shape[0]): 73 | rotation_angle = np.random.uniform() * 2 * np.pi 74 | cosval = np.cos(rotation_angle) 75 | sinval = np.sin(rotation_angle) 76 | rotation_matrix = np.array([[cosval, 0, sinval], 77 | [0, 1, 0], 78 | [-sinval, 0, cosval]]) 79 | shape_pc = batch_xyz_normal[k,:,0:3] 80 | shape_normal = batch_xyz_normal[k,:,3:6] 81 | batch_xyz_normal[k,:,0:3] = np.dot(shape_pc.reshape((-1, 3)), rotation_matrix) 82 | batch_xyz_normal[k,:,3:6] = np.dot(shape_normal.reshape((-1, 3)), rotation_matrix) 83 | return batch_xyz_normal 84 | 85 | def rotate_perturbation_point_cloud_with_normal(batch_data, angle_sigma=0.06, angle_clip=0.18): 86 | """ Randomly perturb the point clouds by small rotations 87 | Input: 88 | BxNx6 array, original batch of point clouds and point normals 89 | Return: 90 | BxNx3 array, rotated batch of point clouds 91 | """ 92 | rotated_data = np.zeros(batch_data.shape, dtype=np.float32) 93 | for k in range(batch_data.shape[0]): 94 | angles = np.clip(angle_sigma*np.random.randn(3), -angle_clip, angle_clip) 95 | Rx = np.array([[1,0,0], 96 | [0,np.cos(angles[0]),-np.sin(angles[0])], 97 | [0,np.sin(angles[0]),np.cos(angles[0])]]) 98 | Ry = np.array([[np.cos(angles[1]),0,np.sin(angles[1])], 99 | [0,1,0], 100 | [-np.sin(angles[1]),0,np.cos(angles[1])]]) 101 | Rz = np.array([[np.cos(angles[2]),-np.sin(angles[2]),0], 102 | [np.sin(angles[2]),np.cos(angles[2]),0], 103 | [0,0,1]]) 104 | R = np.dot(Rz, np.dot(Ry,Rx)) 105 | shape_pc = batch_data[k,:,0:3] 106 | shape_normal = batch_data[k,:,3:6] 107 | rotated_data[k,:,0:3] = np.dot(shape_pc.reshape((-1, 3)), R) 108 | rotated_data[k,:,3:6] = np.dot(shape_normal.reshape((-1, 3)), R) 109 | return rotated_data 110 | 111 | 112 | def rotate_point_cloud_by_angle(batch_data, rotation_angle): 113 | """ Rotate the point cloud along up direction with certain angle. 114 | Input: 115 | BxNx3 array, original batch of point clouds 116 | Return: 117 | BxNx3 array, rotated batch of point clouds 118 | """ 119 | rotated_data = np.zeros(batch_data.shape, dtype=np.float32) 120 | for k in range(batch_data.shape[0]): 121 | #rotation_angle = np.random.uniform() * 2 * np.pi 122 | cosval = np.cos(rotation_angle) 123 | sinval = np.sin(rotation_angle) 124 | rotation_matrix = np.array([[cosval, 0, sinval], 125 | [0, 1, 0], 126 | [-sinval, 0, cosval]]) 127 | shape_pc = batch_data[k,:,0:3] 128 | rotated_data[k,:,0:3] = np.dot(shape_pc.reshape((-1, 3)), rotation_matrix) 129 | return rotated_data 130 | 131 | def rotate_point_cloud_by_angle_with_normal(batch_data, rotation_angle): 132 | """ Rotate the point cloud along up direction with certain angle. 133 | Input: 134 | BxNx3 array, original batch of point clouds 135 | Return: 136 | BxNx3 array, rotated batch of point clouds 137 | """ 138 | rotated_data = np.zeros(batch_data.shape, dtype=np.float32) 139 | for k in range(batch_data.shape[0]): 140 | #rotation_angle = np.random.uniform() * 2 * np.pi 141 | cosval = np.cos(rotation_angle) 142 | sinval = np.sin(rotation_angle) 143 | rotation_matrix = np.array([[cosval, 0, sinval], 144 | [0, 1, 0], 145 | [-sinval, 0, cosval]]) 146 | shape_pc = batch_data[k, ...] 147 | shape_normal = batch_data[k,:,3:6] 148 | rotated_data[k, ...] = np.dot(shape_pc.reshape((-1, 3)), rotation_matrix) 149 | rotated_data[k,:,3:6] = np.dot(shape_normal.reshape((-1,3)), rotation_matrix) 150 | return rotated_data 151 | 152 | 153 | 154 | def rotate_perturbation_point_cloud(batch_data, angle_sigma=0.06, angle_clip=0.18): 155 | """ Randomly perturb the point clouds by small rotations 156 | Input: 157 | BxNx3 array, original batch of point clouds 158 | Return: 159 | BxNx3 array, rotated batch of point clouds 160 | """ 161 | rotated_data = np.zeros(batch_data.shape, dtype=np.float32) 162 | for k in range(batch_data.shape[0]): 163 | angles = np.clip(angle_sigma*np.random.randn(3), -angle_clip, angle_clip) 164 | Rx = np.array([[1,0,0], 165 | [0,np.cos(angles[0]),-np.sin(angles[0])], 166 | [0,np.sin(angles[0]),np.cos(angles[0])]]) 167 | Ry = np.array([[np.cos(angles[1]),0,np.sin(angles[1])], 168 | [0,1,0], 169 | [-np.sin(angles[1]),0,np.cos(angles[1])]]) 170 | Rz = np.array([[np.cos(angles[2]),-np.sin(angles[2]),0], 171 | [np.sin(angles[2]),np.cos(angles[2]),0], 172 | [0,0,1]]) 173 | R = np.dot(Rz, np.dot(Ry,Rx)) 174 | shape_pc = batch_data[k, ...] 175 | rotated_data[k, ...] = np.dot(shape_pc.reshape((-1, 3)), R) 176 | return rotated_data 177 | 178 | 179 | def jitter_point_cloud(batch_data, sigma=0.01, clip=0.05): 180 | """ Randomly jitter points. jittering is per point. 181 | Input: 182 | BxNx3 array, original batch of point clouds 183 | Return: 184 | BxNx3 array, jittered batch of point clouds 185 | """ 186 | B, N, C = batch_data.shape 187 | assert(clip > 0) 188 | jittered_data = np.clip(sigma * np.random.randn(B, N, C), -1*clip, clip) 189 | jittered_data += batch_data 190 | return jittered_data 191 | 192 | def shift_point_cloud(batch_data, shift_range=0.1): 193 | """ Randomly shift point cloud. Shift is per point cloud. 194 | Input: 195 | BxNx3 array, original batch of point clouds 196 | Return: 197 | BxNx3 array, shifted batch of point clouds 198 | """ 199 | B, N, C = batch_data.shape 200 | shifts = np.random.uniform(-shift_range, shift_range, (B,3)) 201 | for batch_index in range(B): 202 | batch_data[batch_index,:,:] += shifts[batch_index,:] 203 | return batch_data 204 | 205 | 206 | def random_scale_point_cloud(batch_data, scale_low=0.8, scale_high=1.25): 207 | """ Randomly scale the point cloud. Scale is per point cloud. 208 | Input: 209 | BxNx3 array, original batch of point clouds 210 | Return: 211 | BxNx3 array, scaled batch of point clouds 212 | """ 213 | B, N, C = batch_data.shape 214 | scales = np.random.uniform(scale_low, scale_high, B) 215 | for batch_index in range(B): 216 | batch_data[batch_index,:,:] *= scales[batch_index] 217 | return batch_data 218 | 219 | def random_point_dropout(batch_pc, max_dropout_ratio=0.875): 220 | ''' batch_pc: BxNx3 ''' 221 | for b in range(batch_pc.shape[0]): 222 | dropout_ratio = np.random.random()*max_dropout_ratio # 0~0.875 223 | drop_idx = np.where(np.random.random((batch_pc.shape[1]))<=dropout_ratio)[0] 224 | if len(drop_idx)>0: 225 | batch_pc[b,drop_idx,:] = batch_pc[b,0,:] # set to the first point 226 | return batch_pc 227 | 228 | 229 | def getDataFiles(list_filename): 230 | return [line.rstrip() for line in open(list_filename)] 231 | 232 | def load_h5(h5_filename): 233 | f = h5py.File(h5_filename) 234 | data = f['data'][:] 235 | label = f['label'][:] 236 | return (data, label) 237 | 238 | def loadDataFile(filename): 239 | return load_h5(filename) 240 | -------------------------------------------------------------------------------- /scannet/README.md: -------------------------------------------------------------------------------- 1 | ### Normal estimation on ScanNet data 2 | 3 | Download *ScanNet* datasets provided by PointNet++. 4 | 5 | Install PCL library on you machine. 6 | 7 | Run the following code to estimate normals on ScanNet scans: 8 | 9 | ./run.sh 10 | -------------------------------------------------------------------------------- /scannet/normal_estimation/1_read_and_save_as_files.py: -------------------------------------------------------------------------------- 1 | import pickle 2 | import os 3 | import sys 4 | import numpy as np 5 | from py3d import * 6 | 7 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 8 | BASE_DIR = os.path.dirname(BASE_DIR) 9 | sys.path.append(BASE_DIR) 10 | from visualize_pc import visualize_pc, print_statistics 11 | 12 | splits = ['train', 'test'] 13 | for split in splits: 14 | print("split: ", split) 15 | data_filename = '../scannet_%s.pickle'%(split) 16 | with open(data_filename,'rb') as fp: 17 | scene_points_list = pickle.load(fp, encoding='bytes') 18 | semantic_labels_list = pickle.load(fp, encoding='bytes') 19 | 20 | for i in range(len(scene_points_list)): 21 | point_cloud = scene_points_list[i] 22 | 23 | centered_point_cloud = np.copy(point_cloud) 24 | xyz_min = np.amin(centered_point_cloud, axis=0)[0:3] 25 | centered_point_cloud[:, 0:3] -= xyz_min 26 | 27 | xyz_max = np.amax(centered_point_cloud, axis=0)[0:3] 28 | centered_point_cloud[:, 0:3] -= (xyz_max / 2.0) 29 | 30 | path = '../scannet_%s/'%(split) 31 | out_filename = path+str(i)+'_xyz.pcd' 32 | fout = open(out_filename, 'w') 33 | fout.write('# .PCD v0.7 - Point Cloud Data file format\n') 34 | fout.write('VERSION 0.7\n') 35 | fout.write('FIELDS x y z\n') 36 | fout.write('SIZE 4 4 4\n') 37 | fout.write('TYPE F F F \n') 38 | fout.write('COUNT 1 1 1\n') 39 | fout.write('WIDTH '+str(centered_point_cloud.shape[0])+'\n') 40 | fout.write('HEIGHT 1\n') 41 | fout.write('VIEWPOINT 0 0 0 1 0 0 0\n') 42 | fout.write('POINTS '+str(centered_point_cloud.shape[0])+'\n') 43 | fout.write('DATA ascii\n') 44 | for i in range(centered_point_cloud.shape[0]): 45 | fout.write('%f %f %f\n' % \ 46 | (centered_point_cloud[i,0], centered_point_cloud[i,1], centered_point_cloud[i,2])) 47 | fout.close() 48 | 49 | print("Point cloud "+out_filename+" processed.") -------------------------------------------------------------------------------- /scannet/normal_estimation/2_compute_normals_and_curvature.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | #include 6 | 7 | #include 8 | #include 9 | #include 10 | #include 11 | 12 | using namespace std; 13 | 14 | int main (int, char** argv) 15 | { 16 | DIR * dir; 17 | struct dirent *pdir; 18 | vector folders {"../../scannet_train/", "../../scannet_test/"}; 19 | 20 | for(const auto path : folders){ 21 | if((dir = opendir(path.c_str())) != NULL){ 22 | while((pdir = readdir(dir)) != NULL){ 23 | string fileName = pdir->d_name; 24 | if(fileName != "." && fileName != ".."){ 25 | cout << "filename: " << fileName << endl; 26 | 27 | pcl::PointCloud::Ptr cloud (new pcl::PointCloud); 28 | if (pcl::io::loadPCDFile (path+fileName, *cloud) == -1) //* load the file 29 | { 30 | PCL_ERROR ("Couldn't read file"); 31 | return (-1); 32 | } 33 | 34 | // Compute the normals 35 | pcl::PointCloud::Ptr cloud_with_normals (new pcl::PointCloud); 36 | 37 | pcl::NormalEstimation normal_estimation; 38 | normal_estimation.setInputCloud (cloud); 39 | 40 | pcl::search::KdTree::Ptr tree (new pcl::search::KdTree); 41 | normal_estimation.setSearchMethod (tree); 42 | 43 | normal_estimation.setKSearch(10); 44 | 45 | normal_estimation.compute (*cloud_with_normals); 46 | 47 | std::ofstream myfile; 48 | size_t pos = fileName.find("_xyz.pcd"); 49 | string outputFileName = path + fileName.substr(0, pos) + "_normals.txt"; 50 | cout << "outputFileName: " + outputFileName << endl; 51 | myfile.open(outputFileName); 52 | 53 | for (size_t i = 0; i < cloud->points.size (); ++i) 54 | myfile << cloud_with_normals->points[i].normal_x 55 | << " " << cloud_with_normals->points[i].normal_y 56 | << " " << cloud_with_normals->points[i].normal_z 57 | << "\n"; 58 | 59 | myfile.close(); 60 | 61 | cout << "File " << outputFileName << " saved." << endl; 62 | } 63 | } 64 | closedir (dir); 65 | } else { 66 | perror("Could not open director."); 67 | return EXIT_FAILURE; 68 | } 69 | } 70 | return 0; 71 | } 72 | -------------------------------------------------------------------------------- /scannet/normal_estimation/3_pickle_normals.py: -------------------------------------------------------------------------------- 1 | import pickle 2 | import os 3 | import sys 4 | import numpy as np 5 | from py3d import * 6 | 7 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 8 | BASE_DIR = os.path.dirname(BASE_DIR) 9 | sys.path.append(BASE_DIR) 10 | from visualize_pc import visualize_pc, print_statistics 11 | 12 | splits = ['train', 'test'] 13 | for split in splits: 14 | print("split: ", split) 15 | data_filename = '../scannet_%s.pickle'%(split) 16 | with open(data_filename,'rb') as fp: 17 | scene_points_list = pickle.load(fp, encoding='bytes') 18 | semantic_labels_list = pickle.load(fp, encoding='bytes') 19 | 20 | scene_points_with_normals_list = [] 21 | for i in range(len(scene_points_list)): 22 | point_cloud = scene_points_list[i] 23 | 24 | path = '../scannet_%s/'%(split) 25 | normals_filename = path+str(i)+'_normals.txt' 26 | normals = np.loadtxt(normals_filename) 27 | 28 | scene_points_with_normals_list.append(np.concatenate([point_cloud, normals], 1)) 29 | print("Point cloud "+normals_filename+" processed.") 30 | 31 | assert(len(scene_points_with_normals_list) == len(scene_points_list)) 32 | pickle_filename = '../scannet_with_normals_%s.pickle'%(split) 33 | pickle_out = open(pickle_filename,'wb') 34 | pickle.dump(scene_points_with_normals_list, pickle_out) 35 | pickle.dump(semantic_labels_list, pickle_out) 36 | pickle_out.close() 37 | -------------------------------------------------------------------------------- /scannet/normal_estimation/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 2.6 FATAL_ERROR) 2 | 3 | project(MY_GRAND_PROJECT) 4 | 5 | find_package(PCL 1.3 REQUIRED COMPONENTS common io) 6 | 7 | list(REMOVE_ITEM PCL_LIBRARIES "vtkproj4") 8 | 9 | include_directories(${PCL_INCLUDE_DIRS}) 10 | link_directories(${PCL_LIBRARY_DIRS}) 11 | add_definitions(${PCL_DEFINITIONS}) 12 | 13 | add_executable(2_compute_normals_and_curvature 2_compute_normals_and_curvature.cpp) 14 | target_link_libraries(2_compute_normals_and_curvature ${PCL_COMMON_LIBRARIES} ${PCL_IO_LIBRARIES} pcl_common pcl_kdtree pcl_search pcl_features pcl_io) 15 | -------------------------------------------------------------------------------- /scannet/normal_estimation/run.sh: -------------------------------------------------------------------------------- 1 | #/usr/bin/sh 2 | python 1_read_and_save_as_files.py 3 | rm -rf build/ 4 | mkdir build 5 | cd build 6 | cmake .. 7 | make 8 | ./2_compute_normals_and_curvature 9 | python 3_pickle_normals.py -------------------------------------------------------------------------------- /tf_ops/3d_interpolation/.gitignore: -------------------------------------------------------------------------------- 1 | *.o 2 | *.so 3 | __pycache__/ -------------------------------------------------------------------------------- /tf_ops/3d_interpolation/interpolate.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include // memset 4 | #include // rand, RAND_MAX 5 | #include // sqrtf 6 | #include 7 | #include 8 | using namespace std; 9 | float randomf(){ 10 | return (rand()+0.5)/(RAND_MAX+1.0); 11 | } 12 | static double get_time(){ 13 | timespec tp; 14 | clock_gettime(CLOCK_MONOTONIC,&tp); 15 | return tp.tv_sec+tp.tv_nsec*1e-9; 16 | } 17 | 18 | // Find three nearest neigbors with square distance 19 | // input: xyz1 (b,n,3), xyz2(b,m,3) 20 | // output: dist (b,n,3), idx (b,n,3) 21 | void threenn_cpu(int b, int n, int m, const float *xyz1, const float *xyz2, float *dist, int *idx) { 22 | for (int i=0;i 2 | #include 3 | #include // memset 4 | #include // rand, RAND_MAX 5 | #include // sqrtf 6 | #include "tensorflow/core/framework/op.h" 7 | #include "tensorflow/core/framework/op_kernel.h" 8 | #include "tensorflow/core/framework/shape_inference.h" 9 | #include "tensorflow/core/framework/common_shape_fns.h" 10 | using namespace tensorflow; 11 | 12 | REGISTER_OP("ThreeNN") 13 | .Input("xyz1: float32") 14 | .Input("xyz2: float32") 15 | .Output("dist: float32") 16 | .Output("idx: int32") 17 | .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) { 18 | c->set_output(0, c->input(0)); 19 | c->set_output(1, c->input(0)); 20 | return Status::OK(); 21 | }); 22 | REGISTER_OP("ThreeInterpolate") 23 | .Input("points: float32") 24 | .Input("idx: int32") 25 | .Input("weight: float32") 26 | .Output("out: float32") 27 | .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) { 28 | ::tensorflow::shape_inference::ShapeHandle dims1; // (b,m,c) 29 | c->WithRank(c->input(0), 3, &dims1); 30 | ::tensorflow::shape_inference::ShapeHandle dims2; // (b,n,3) 31 | c->WithRank(c->input(1), 3, &dims2); 32 | // (b,n,c) 33 | ::tensorflow::shape_inference::ShapeHandle output = c->MakeShape({c->Dim(dims1, 0), c->Dim(dims2, 1), c->Dim(dims1, 2)}); 34 | c->set_output(0, output); 35 | return Status::OK(); 36 | }); 37 | REGISTER_OP("ThreeInterpolateGrad") 38 | .Input("points: float32") 39 | .Input("idx: int32") 40 | .Input("weight: float32") 41 | .Input("grad_out: float32") 42 | .Output("grad_points: float32") 43 | .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) { 44 | c->set_output(0, c->input(0)); 45 | return Status::OK(); 46 | }); 47 | 48 | float randomf(){ 49 | return (rand()+0.5)/(RAND_MAX+1.0); 50 | } 51 | static double get_time(){ 52 | timespec tp; 53 | clock_gettime(CLOCK_MONOTONIC,&tp); 54 | return tp.tv_sec+tp.tv_nsec*1e-9; 55 | } 56 | 57 | // Find three nearest neigbors with square distance 58 | // input: xyz1 (b,n,3), xyz2(b,m,3) 59 | // output: dist (b,n,3), idx (b,n,3) 60 | void threenn_cpu(int b, int n, int m, const float *xyz1, const float *xyz2, float *dist, int *idx) { 61 | for (int i=0;iinput(0); 163 | OP_REQUIRES(context, xyz1_tensor.dims()==3 && xyz1_tensor.shape().dim_size(2)==3, errors::InvalidArgument("ThreeNN expects (b,n,3) xyz1 shape.")); 164 | int b = xyz1_tensor.shape().dim_size(0); 165 | int n = xyz1_tensor.shape().dim_size(1); 166 | 167 | const Tensor& xyz2_tensor = context->input(1); 168 | OP_REQUIRES(context, xyz2_tensor.dims()==3 && xyz2_tensor.shape().dim_size(2)==3, errors::InvalidArgument("ThreeNN expects (b,m,3) xyz2 shape.")); 169 | int m = xyz2_tensor.shape().dim_size(1); 170 | 171 | Tensor *dist_tensor = nullptr; 172 | OP_REQUIRES_OK(context, context->allocate_output(0, TensorShape{b,n,3}, &dist_tensor)); 173 | Tensor *idx_tensor = nullptr; 174 | OP_REQUIRES_OK(context, context->allocate_output(1, TensorShape{b,n,3}, &idx_tensor)); 175 | 176 | auto xyz1_flat = xyz1_tensor.flat(); 177 | const float *xyz1 = &(xyz1_flat(0)); 178 | auto xyz2_flat = xyz2_tensor.flat(); 179 | const float *xyz2 = &(xyz2_flat(0)); 180 | auto dist_flat = dist_tensor->flat(); 181 | float *dist = &(dist_flat(0)); 182 | auto idx_flat = idx_tensor->flat(); 183 | int *idx = &(idx_flat(0)); 184 | threenn_cpu(b,n,m,xyz1,xyz2,dist,idx); 185 | } 186 | }; 187 | REGISTER_KERNEL_BUILDER(Name("ThreeNN").Device(DEVICE_CPU), ThreeNNOp); 188 | 189 | 190 | 191 | class ThreeInterpolateOp: public OpKernel{ 192 | public: 193 | explicit ThreeInterpolateOp(OpKernelConstruction * context):OpKernel(context){} 194 | 195 | void Compute(OpKernelContext * context) override { 196 | const Tensor& points_tensor=context->input(0); 197 | OP_REQUIRES(context, points_tensor.dims()==3, errors::InvalidArgument("ThreeInterpolate expects (b,m,c) points shape")); 198 | int b = points_tensor.shape().dim_size(0); 199 | int m = points_tensor.shape().dim_size(1); 200 | int c = points_tensor.shape().dim_size(2); 201 | 202 | const Tensor& idx_tensor=context->input(1); 203 | OP_REQUIRES(context,idx_tensor.dims()==3 && idx_tensor.shape().dim_size(0)==b && idx_tensor.shape().dim_size(2)==3, errors::InvalidArgument("ThreeInterpolate expects (b,n,3) idx shape")); 204 | int n = idx_tensor.shape().dim_size(1); 205 | const Tensor& weight_tensor=context->input(2); 206 | OP_REQUIRES(context,weight_tensor.dims()==3 && weight_tensor.shape().dim_size(0)==b && weight_tensor.shape().dim_size(1)==n && weight_tensor.shape().dim_size(2)==3, errors::InvalidArgument("ThreeInterpolate expects (b,n,3) weight shape")); 207 | 208 | Tensor * out_tensor = nullptr; 209 | OP_REQUIRES_OK(context, context->allocate_output(0,TensorShape{b,n,c}, &out_tensor)); 210 | 211 | auto points_flat = points_tensor.flat(); 212 | const float *points = &(points_flat(0)); 213 | auto idx_flat = idx_tensor.flat(); 214 | const int *idx = &(idx_flat(0)); 215 | auto weight_flat = weight_tensor.flat(); 216 | const float *weight = &(weight_flat(0)); 217 | auto out_flat = out_tensor->flat(); 218 | float *out = &(out_flat(0)); 219 | threeinterpolate_cpu(b,m,c,n,points,idx,weight,out); 220 | } 221 | }; 222 | REGISTER_KERNEL_BUILDER(Name("ThreeInterpolate").Device(DEVICE_CPU),ThreeInterpolateOp); 223 | 224 | 225 | class ThreeInterpolateGradOp: public OpKernel{ 226 | public: 227 | explicit ThreeInterpolateGradOp(OpKernelConstruction * context):OpKernel(context){} 228 | 229 | void Compute(OpKernelContext * context) override { 230 | const Tensor& points_tensor=context->input(0); 231 | OP_REQUIRES(context, points_tensor.dims()==3, errors::InvalidArgument("ThreeInterpolateGrad expects (b,m,c) points shape")); 232 | int b = points_tensor.shape().dim_size(0); 233 | int m = points_tensor.shape().dim_size(1); 234 | int c = points_tensor.shape().dim_size(2); 235 | 236 | const Tensor& idx_tensor=context->input(1); 237 | OP_REQUIRES(context,idx_tensor.dims()==3 && idx_tensor.shape().dim_size(0)==b, errors::InvalidArgument("ThreeInterpolateGrad expects (b,n,3) idx shape")); 238 | int n = idx_tensor.shape().dim_size(1); 239 | const Tensor& weight_tensor=context->input(2); 240 | OP_REQUIRES(context,weight_tensor.dims()==3 && weight_tensor.shape().dim_size(0)==b && weight_tensor.shape().dim_size(1)==n && weight_tensor.shape().dim_size(2)==3, errors::InvalidArgument("ThreeInterpolateGrad expects (b,n,3) weight shape")); 241 | 242 | const Tensor& grad_out_tensor=context->input(3); 243 | OP_REQUIRES(context,grad_out_tensor.dims()==3 && grad_out_tensor.shape().dim_size(0)==b && grad_out_tensor.shape().dim_size(1)==n && grad_out_tensor.shape().dim_size(2)==c, errors::InvalidArgument("ThreeInterpolateGrad expects (b,n,c) grad_out shape")); 244 | 245 | Tensor * grad_points_tensor = nullptr; 246 | OP_REQUIRES_OK(context, context->allocate_output(0,TensorShape{b,m,c}, &grad_points_tensor)); 247 | 248 | auto points_flat = points_tensor.flat(); 249 | const float *points = &(points_flat(0)); 250 | auto idx_flat = idx_tensor.flat(); 251 | const int *idx = &(idx_flat(0)); 252 | auto weight_flat = weight_tensor.flat(); 253 | const float *weight = &(weight_flat(0)); 254 | auto grad_out_flat = grad_out_tensor.flat(); 255 | const float *grad_out = &(grad_out_flat(0)); 256 | auto grad_points_flat = grad_points_tensor->flat(); 257 | float *grad_points = &(grad_points_flat(0)); 258 | memset(grad_points, 0, sizeof(float)*b*m*c); 259 | threeinterpolate_grad_cpu(b,n,c,m,grad_out,idx,weight,grad_points); 260 | } 261 | }; 262 | REGISTER_KERNEL_BUILDER(Name("ThreeInterpolateGrad").Device(DEVICE_CPU),ThreeInterpolateGradOp); 263 | 264 | 265 | -------------------------------------------------------------------------------- /tf_ops/3d_interpolation/tf_interpolate.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | from tensorflow.python.framework import ops 3 | import sys 4 | import os 5 | BASE_DIR = os.path.dirname(__file__) 6 | sys.path.append(BASE_DIR) 7 | interpolate_module=tf.load_op_library(os.path.join(BASE_DIR, 'tf_interpolate_so.so')) 8 | def three_nn(xyz1, xyz2): 9 | ''' 10 | Input: 11 | xyz1: (b,n,3) float32 array, unknown points 12 | xyz2: (b,m,3) float32 array, known points 13 | Output: 14 | dist: (b,n,3) float32 array, distances to known points 15 | idx: (b,n,3) int32 array, indices to known points 16 | ''' 17 | return interpolate_module.three_nn(xyz1, xyz2) 18 | ops.NoGradient('ThreeNN') 19 | def three_interpolate(points, idx, weight): 20 | ''' 21 | Input: 22 | points: (b,m,c) float32 array, known points 23 | idx: (b,n,3) int32 array, indices to known points 24 | weight: (b,n,3) float32 array, weights on known points 25 | Output: 26 | out: (b,n,c) float32 array, interpolated point values 27 | ''' 28 | return interpolate_module.three_interpolate(points, idx, weight) 29 | @tf.RegisterGradient('ThreeInterpolate') 30 | def _three_interpolate_grad(op, grad_out): 31 | points = op.inputs[0] 32 | idx = op.inputs[1] 33 | weight = op.inputs[2] 34 | return [interpolate_module.three_interpolate_grad(points, idx, weight, grad_out), None, None] 35 | 36 | if __name__=='__main__': 37 | import numpy as np 38 | import time 39 | np.random.seed(100) 40 | pts = np.random.random((32,128,64)).astype('float32') 41 | tmp1 = np.random.random((32,512,3)).astype('float32') 42 | tmp2 = np.random.random((32,128,3)).astype('float32') 43 | with tf.device('/cpu:0'): 44 | points = tf.constant(pts) 45 | xyz1 = tf.constant(tmp1) 46 | xyz2 = tf.constant(tmp2) 47 | dist, idx = three_nn(xyz1, xyz2) 48 | weight = tf.ones_like(dist)/3.0 49 | interpolated_points = three_interpolate(points, idx, weight) 50 | with tf.Session('') as sess: 51 | now = time.time() 52 | for _ in range(100): 53 | ret = sess.run(interpolated_points) 54 | print (time.time() - now) 55 | print (ret.shape, ret.dtype) 56 | #print ret 57 | -------------------------------------------------------------------------------- /tf_ops/3d_interpolation/tf_interpolate_compile.sh: -------------------------------------------------------------------------------- 1 | #/bin/bash 2 | g++ -std=c++11 tf_interpolate.cpp -o tf_interpolate_so.so -shared -fPIC -I /../anaconda3/lib/python3.6/site-packages/tensorflow/include -I /usr/local/cuda-8.0/include -lcudart -L /usr/local/cuda-8.0/lib64/ -O2 -D_GLIBCXX_USE_CXX11_ABI=0 -------------------------------------------------------------------------------- /tf_ops/grouping_ring/.gitignore: -------------------------------------------------------------------------------- 1 | *.o 2 | *.so 3 | __pycache__/ -------------------------------------------------------------------------------- /tf_ops/grouping_ring/tf_grouping_ring.cpp: -------------------------------------------------------------------------------- 1 | #include "tensorflow/core/framework/op.h" 2 | #include "tensorflow/core/framework/op_kernel.h" 3 | #include "tensorflow/core/framework/shape_inference.h" 4 | #include "tensorflow/core/framework/common_shape_fns.h" 5 | #include 6 | 7 | using namespace tensorflow; 8 | 9 | REGISTER_OP("RingPoint") 10 | .Attr("radius_in: float") 11 | .Attr("radius_out: float") 12 | .Attr("nsample: int") 13 | .Input("xyz1: float32") 14 | .Input("xyz2: float32") 15 | .Input("idx2: int32") 16 | .Output("idx: int32") 17 | .Output("pts_cnt: int32") 18 | .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) { 19 | ::tensorflow::shape_inference::ShapeHandle dims2; 20 | c->WithRank(c->input(1), 3, &dims2); 21 | int nsample; 22 | TF_RETURN_IF_ERROR(c->GetAttr("nsample", &nsample)); 23 | ::tensorflow::shape_inference::ShapeHandle output1 = c->MakeShape({c->Dim(dims2, 0), c->Dim(dims2, 1), nsample}); 24 | c->set_output(0, output1); 25 | ::tensorflow::shape_inference::ShapeHandle output2 = c->MakeShape({c->Dim(dims2, 0), c->Dim(dims2, 1)}); 26 | c->set_output(1, output2); 27 | return Status::OK(); 28 | }); 29 | REGISTER_OP("GroupPoint") 30 | .Input("points: float32") 31 | .Input("idx: int32") 32 | .Output("out: float32") 33 | .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) { 34 | ::tensorflow::shape_inference::ShapeHandle dims1; // batch_size * ndataset * channels 35 | c->WithRank(c->input(0), 3, &dims1); 36 | ::tensorflow::shape_inference::ShapeHandle dims2; // batch_size * npoints * nsample 37 | c->WithRank(c->input(1), 3, &dims2); 38 | // batch_size * npoints * nsample * channels 39 | ::tensorflow::shape_inference::ShapeHandle output = c->MakeShape({c->Dim(dims2, 0), c->Dim(dims2, 1), c->Dim(dims2, 2), c->Dim(dims1, 2)}); 40 | c->set_output(0, output); 41 | return Status::OK(); 42 | }); 43 | REGISTER_OP("GroupPointGrad") 44 | .Input("points: float32") 45 | .Input("idx: int32") 46 | .Input("grad_out: float32") 47 | .Output("grad_points: float32") 48 | .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) { 49 | c->set_output(0, c->input(0)); 50 | return Status::OK(); 51 | }); 52 | 53 | 54 | void ringPointLauncher(int b, int n, int m, float radius_in, float radius_out, int nsample, const float *xyz1, const float *xyz2, const int * idx2, int *idx, int *pts_cnt); 55 | class RingPointGpuOp : public OpKernel { 56 | public: 57 | explicit RingPointGpuOp(OpKernelConstruction* context) : OpKernel(context) { 58 | OP_REQUIRES_OK(context, context->GetAttr("radius_in", &radius_in_)); 59 | OP_REQUIRES(context, radius_in_ >= 0, errors::InvalidArgument("RingPoint expects positive inner radius")); 60 | 61 | OP_REQUIRES_OK(context, context->GetAttr("radius_out", &radius_out_)); 62 | OP_REQUIRES(context, radius_out_ > 0, errors::InvalidArgument("RingPoint expects positive outter radius")); 63 | 64 | OP_REQUIRES_OK(context, context->GetAttr("nsample", &nsample_)); 65 | OP_REQUIRES(context, nsample_ > 0, errors::InvalidArgument("RingPoint expects positive nsample")); 66 | } 67 | 68 | void Compute(OpKernelContext* context) override { 69 | const Tensor& xyz1_tensor = context->input(0); 70 | OP_REQUIRES(context, xyz1_tensor.dims()==3 && xyz1_tensor.shape().dim_size(2)==3, errors::InvalidArgument("RingPoint expects (batch_size, ndataset, 3) xyz1 shape.")); 71 | int b = xyz1_tensor.shape().dim_size(0); 72 | int n = xyz1_tensor.shape().dim_size(1); 73 | 74 | const Tensor& xyz2_tensor = context->input(1); 75 | OP_REQUIRES(context, xyz2_tensor.dims()==3 && xyz2_tensor.shape().dim_size(2)==3, errors::InvalidArgument("RingPoint expects (batch_size, npoint, 3) xyz2 shape.")); 76 | int m = xyz2_tensor.shape().dim_size(1); 77 | 78 | const Tensor& idx2_tensor = context->input(2); 79 | OP_REQUIRES(context, idx2_tensor.dims()==2, errors::InvalidArgument("RingPoint expects (batch_size, npoint) idx2 shape.")); 80 | 81 | Tensor *idx_tensor = nullptr; 82 | OP_REQUIRES_OK(context, context->allocate_output(0, TensorShape{b,m,nsample_}, &idx_tensor)); 83 | Tensor *pts_cnt_tensor = nullptr; 84 | OP_REQUIRES_OK(context, context->allocate_output(1, TensorShape{b,m}, &pts_cnt_tensor)); 85 | 86 | auto xyz1_flat = xyz1_tensor.flat(); 87 | const float *xyz1 = &(xyz1_flat(0)); 88 | auto xyz2_flat = xyz2_tensor.flat(); 89 | const float *xyz2 = &(xyz2_flat(0)); 90 | auto idx2_flat = idx2_tensor.flat(); 91 | const int *idx2 = &(idx2_flat(0)); 92 | auto idx_flat = idx_tensor->flat(); 93 | int *idx = &(idx_flat(0)); 94 | auto pts_cnt_flat = pts_cnt_tensor->flat(); 95 | int *pts_cnt = &(pts_cnt_flat(0)); 96 | ringPointLauncher(b,n,m,radius_in_,radius_out_,nsample_,xyz1,xyz2,idx2,idx,pts_cnt); 97 | } 98 | private: 99 | float radius_in_; 100 | float radius_out_; 101 | int nsample_; 102 | }; 103 | REGISTER_KERNEL_BUILDER(Name("RingPoint").Device(DEVICE_GPU), RingPointGpuOp); 104 | 105 | void groupPointLauncher(int b, int n, int c, int m, int nsample, const float *points, const int *idx, float *out); 106 | class GroupPointGpuOp: public OpKernel{ 107 | public: 108 | explicit GroupPointGpuOp(OpKernelConstruction * context):OpKernel(context){} 109 | 110 | void Compute(OpKernelContext * context) override { 111 | const Tensor& points_tensor=context->input(0); 112 | OP_REQUIRES(context, points_tensor.dims()==3, errors::InvalidArgument("GroupPoint expects (batch_size, num_points, channel) points shape")); 113 | int b = points_tensor.shape().dim_size(0); 114 | int n = points_tensor.shape().dim_size(1); 115 | int c = points_tensor.shape().dim_size(2); 116 | 117 | const Tensor& idx_tensor=context->input(1); 118 | OP_REQUIRES(context,idx_tensor.dims()==3 && idx_tensor.shape().dim_size(0)==b, errors::InvalidArgument("GroupPoint expects (batch_size, npoints, nsample) idx shape")); 119 | int m = idx_tensor.shape().dim_size(1); 120 | int nsample = idx_tensor.shape().dim_size(2); 121 | 122 | Tensor * out_tensor = nullptr; 123 | OP_REQUIRES_OK(context, context->allocate_output(0,TensorShape{b,m,nsample,c}, &out_tensor)); 124 | 125 | auto points_flat = points_tensor.flat(); 126 | const float *points = &(points_flat(0)); 127 | auto idx_flat = idx_tensor.flat(); 128 | const int *idx = &(idx_flat(0)); 129 | auto out_flat = out_tensor->flat(); 130 | float *out = &(out_flat(0)); 131 | groupPointLauncher(b,n,c,m,nsample,points,idx,out); 132 | } 133 | }; 134 | REGISTER_KERNEL_BUILDER(Name("GroupPoint").Device(DEVICE_GPU),GroupPointGpuOp); 135 | 136 | void groupPointGradLauncher(int b, int n, int c, int m, int nsample, const float *grad_out, const int *idx, float *grad_points); 137 | class GroupPointGradGpuOp: public OpKernel{ 138 | public: 139 | explicit GroupPointGradGpuOp(OpKernelConstruction * context):OpKernel(context){} 140 | 141 | void Compute(OpKernelContext * context) override { 142 | const Tensor& points_tensor=context->input(0); 143 | OP_REQUIRES(context, points_tensor.dims()==3, errors::InvalidArgument("GroupPointGrad expects (batch_size, num_points, channel) points shape")); 144 | int b = points_tensor.shape().dim_size(0); 145 | int n = points_tensor.shape().dim_size(1); 146 | int c = points_tensor.shape().dim_size(2); 147 | 148 | const Tensor& idx_tensor=context->input(1); 149 | OP_REQUIRES(context,idx_tensor.dims()==3 && idx_tensor.shape().dim_size(0)==b, errors::InvalidArgument("GroupPointGrad expects (batch_size, npoints, nsample) idx shape")); 150 | int m = idx_tensor.shape().dim_size(1); 151 | int nsample = idx_tensor.shape().dim_size(2); 152 | 153 | const Tensor& grad_out_tensor=context->input(2); 154 | OP_REQUIRES(context,grad_out_tensor.dims()==4 && grad_out_tensor.shape().dim_size(0)==b && grad_out_tensor.shape().dim_size(1)==m && grad_out_tensor.shape().dim_size(2)==nsample && grad_out_tensor.shape().dim_size(3)==c, errors::InvalidArgument("GroupPointGrad expects (batch_size, npoints, nsample, channel) grad_out shape")); 155 | 156 | Tensor * grad_points_tensor = nullptr; 157 | OP_REQUIRES_OK(context, context->allocate_output(0,TensorShape{b,n,c}, &grad_points_tensor)); 158 | 159 | auto points_flat = points_tensor.flat(); 160 | const float *points = &(points_flat(0)); 161 | auto idx_flat = idx_tensor.flat(); 162 | const int *idx = &(idx_flat(0)); 163 | auto grad_out_flat = grad_out_tensor.flat(); 164 | const float *grad_out = &(grad_out_flat(0)); 165 | auto grad_points_flat = grad_points_tensor->flat(); 166 | float *grad_points = &(grad_points_flat(0)); 167 | cudaMemset(grad_points, 0, sizeof(float)*b*n*c); 168 | groupPointGradLauncher(b,n,c,m,nsample,grad_out,idx,grad_points); 169 | } 170 | }; 171 | REGISTER_KERNEL_BUILDER(Name("GroupPointGrad").Device(DEVICE_GPU),GroupPointGradGpuOp); 172 | -------------------------------------------------------------------------------- /tf_ops/grouping_ring/tf_grouping_ring.py: -------------------------------------------------------------------------------- 1 | ''' Grouping implemented by Charles R. Qi 2 | All Rights Reserved. 2017. 3 | 4 | Rings 5 | Author: Artem Komarichev 6 | All Rights Reserved. 2018. 7 | ''' 8 | 9 | import tensorflow as tf 10 | from tensorflow.python.framework import ops 11 | import sys 12 | import os 13 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 14 | sys.path.append(BASE_DIR) 15 | grouping_ring_module=tf.load_op_library(os.path.join(BASE_DIR, 'tf_grouping_ring_so.so')) 16 | def ring_point(radius_in, radius_out, nsample, xyz1, xyz2, idx2): 17 | ''' 18 | Input: 19 | radius_in: float32, ball search inner radius 20 | radius_out: float32, ball search outter radius 21 | nsample: int32, number of points selected in each ball region 22 | xyz1: (batch_size, ndataset, 3) float32 array, input points 23 | xyz2: (batch_size, npoint, 3) float32 array, query points 24 | idx2: (batch_size, npoint) int32, indecies of query points 25 | Output: 26 | idx: (batch_size, npoint, nsample) int32 array, indices to input points 27 | pts_cnt: (batch_size, npoint) int32 array, number of unique points in each local region 28 | ''' 29 | #return grouping_module.query_ball_point(radius, nsample, xyz1, xyz2) 30 | return grouping_ring_module.ring_point(xyz1, xyz2, idx2, radius_in, radius_out, nsample) 31 | ops.NoGradient('RingPoint') 32 | 33 | def group_point(points, idx): 34 | ''' 35 | Input: 36 | points: (batch_size, ndataset, channel) float32 array, points to sample from 37 | idx: (batch_size, npoint, nsample) int32 array, indices to points 38 | Output: 39 | out: (batch_size, npoint, nsample, channel) float32 array, values sampled from points 40 | ''' 41 | return grouping_ring_module.group_point(points, idx) 42 | @tf.RegisterGradient('GroupPoint') 43 | def _group_point_grad(op, grad_out): 44 | points = op.inputs[0] 45 | idx = op.inputs[1] 46 | return [grouping_ring_module.group_point_grad(points, idx, grad_out), None] 47 | 48 | def knn_point(k, xyz1, xyz2): 49 | ''' 50 | Input: 51 | k: int32, number of k in k-nn search 52 | xyz1: (batch_size, ndataset, c) float32 array, input points 53 | xyz2: (batch_size, npoint, c) float32 array, query points 54 | Output: 55 | val: (batch_size, npoint, k) float32 array, L2 distances 56 | idx: (batch_size, npoint, k) int32 array, indices to input points 57 | ''' 58 | b = xyz1.get_shape()[0].value 59 | n = xyz1.get_shape()[1].value 60 | c = xyz1.get_shape()[2].value 61 | m = xyz2.get_shape()[1].value 62 | print (b, n, c, m) 63 | print (xyz1, (b,1,n,c)) 64 | xyz1 = tf.tile(tf.reshape(xyz1, (b,1,n,c)), [1,m,1,1]) 65 | xyz2 = tf.tile(tf.reshape(xyz2, (b,m,1,c)), [1,1,n,1]) 66 | dist = tf.reduce_sum((xyz1-xyz2)**2, -1) 67 | print (dist, k) 68 | outi, out = select_top_k(k, dist) 69 | idx = tf.slice(outi, [0,0,0], [-1,-1,k]) 70 | val = tf.slice(out, [0,0,0], [-1,-1,k]) 71 | print (idx, val) 72 | print(idx[0,0,:]) 73 | #val, idx = tf.nn.top_k(-dist, k=k) # ONLY SUPPORT CPU 74 | return val, idx 75 | -------------------------------------------------------------------------------- /tf_ops/grouping_ring/tf_grouping_ring_compile.sh: -------------------------------------------------------------------------------- 1 | #/bin/bash 2 | /usr/local/cuda-8.0/bin/nvcc tf_grouping_ring_g.cu -o tf_grouping_ring_g.cu.o -c -O2 -DGOOGLE_CUDA=1 -x cu -Xcompiler -fPIC 3 | g++ -std=c++11 tf_grouping_ring.cpp tf_grouping_ring_g.cu.o -o tf_grouping_ring_so.so -shared -fPIC -I /../anaconda3/lib/python3.6/site-packages/tensorflow/include -I /usr/local/cuda-8.0/include -lcudart -L /usr/local/cuda-8.0/lib64/ -O2 -D_GLIBCXX_USE_CXX11_ABI=0 -------------------------------------------------------------------------------- /tf_ops/grouping_ring/tf_grouping_ring_g.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | __global__ void ring_point_gpu(int b, int n, int m, float radius_in, float radius_out, int nsample, const float *xyz1, const float *xyz2, const int *idx2, int *idx, int *pts_cnt) { 5 | int batch_index = blockIdx.x; 6 | xyz1 += n*3*batch_index; 7 | xyz2 += m*3*batch_index; 8 | idx2 += m*batch_index; 9 | idx += m*nsample*batch_index; 10 | pts_cnt += m*batch_index; // counting how many unique points selected in the local region 11 | 12 | int index = threadIdx.x; 13 | int stride = blockDim.x; 14 | 15 | for (int j=index;j= radius_in && d farthest_dist){ 38 | farthest_ind = cnt; 39 | farthest_dist = d; 40 | } 41 | idx[j*nsample+cnt] = k; 42 | dist[cnt] = d; 43 | cnt+=1; 44 | } 45 | else if(d < farthest_dist && d >= radius_in){ // && d > 1e-5 // found closer point than fatherst one found so far 46 | idx[j*nsample+farthest_ind] = k; 47 | dist[farthest_ind] = d; 48 | 49 | // update farthest distance and farthest index 50 | farthest_dist = FLT_MIN; 51 | farthest_ind = -1; 52 | for(int l=0; l farthest_dist){ 54 | farthest_ind = l; 55 | farthest_dist = dist[l]; 56 | } 57 | } 58 | } 59 | } 60 | 61 | if(cnt==0){ 62 | for (int l=0;l>>(b,n,m,radius_in,radius_out,nsample,xyz1,xyz2,idx2,idx,pts_cnt); 124 | //cudaDeviceSynchronize(); 125 | } 126 | void groupPointLauncher(int b, int n, int c, int m, int nsample, const float *points, const int *idx, float *out){ 127 | group_point_gpu<<>>(b,n,c,m,nsample,points,idx,out); 128 | //cudaDeviceSynchronize(); 129 | } 130 | void groupPointGradLauncher(int b, int n, int c, int m, int nsample, const float *grad_out, const int *idx, float *grad_points){ 131 | group_point_grad_gpu<<>>(b,n,c,m,nsample,grad_out,idx,grad_points); 132 | //cudaDeviceSynchronize(); 133 | } 134 | -------------------------------------------------------------------------------- /tf_ops/ordering/.gitignore: -------------------------------------------------------------------------------- 1 | *.o 2 | *.so 3 | __pycache__/ 4 | -------------------------------------------------------------------------------- /tf_ops/ordering/tf_ordering.cpp: -------------------------------------------------------------------------------- 1 | /* The module orders neighbors in counterclockwise manner. 2 | * Author: Artem Komarichev 3 | * All Rights Reserved. 2018. 4 | */ 5 | 6 | #include "tensorflow/core/framework/op.h" 7 | #include "tensorflow/core/framework/op_kernel.h" 8 | #include "tensorflow/core/framework/shape_inference.h" 9 | #include "tensorflow/core/framework/common_shape_fns.h" 10 | #include 11 | 12 | using namespace tensorflow; 13 | using namespace std; 14 | 15 | REGISTER_OP("OrderNeighbors") 16 | .Attr("k: int") 17 | .Input("input_xyz: float32") 18 | .Input("query_xyz: float32") 19 | .Input("query_normals: float32") 20 | .Input("idx: int32") 21 | .Output("outi: int32") 22 | .Output("proj: float32") 23 | .Output("angles: float32") 24 | .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c){ 25 | int k; 26 | TF_RETURN_IF_ERROR(c->GetAttr("k", &k)); 27 | ::tensorflow::shape_inference::ShapeHandle dims2; 28 | c->WithRank(c->input(3), 3, &dims2); 29 | ::tensorflow::shape_inference::ShapeHandle dims1; 30 | c->WithRank(c->input(0), 3, &dims1); 31 | ::tensorflow::shape_inference::ShapeHandle output2 = c->MakeShape({c->Dim(dims2, 0), c->Dim(dims2, 1), k, c->Dim(dims1,2)}); 32 | c->set_output(0, c->input(3)); 33 | c->set_output(1, output2); 34 | c->set_output(2, c->input(3)); 35 | return Status::OK(); 36 | }); 37 | 38 | void orderNeighborsLauncher(int b, int m, int n, int m_q, int k, const float *input, const float *queries, const float *queries_norm, const int *idx, float *proj, int *outi, float *angles); 39 | class OrderNeighborsGpuOp : public OpKernel { 40 | public: 41 | explicit OrderNeighborsGpuOp(OpKernelConstruction * context):OpKernel(context){ 42 | OP_REQUIRES_OK(context, context->GetAttr("k", &k_)); 43 | OP_REQUIRES(context, k_ > 0, errors::InvalidArgument("OrderNeighbors expects positive k")); 44 | } 45 | 46 | void Compute(OpKernelContext* context) override { 47 | const Tensor& input_xyz_tensor = context->input(0); 48 | OP_REQUIRES(context, input_xyz_tensor.dims() == 3, errors::InvalidArgument("OrderNeighbors expects (b,m,n) input_xyz shape")); 49 | 50 | int b = input_xyz_tensor.shape().dim_size(0); 51 | int m = input_xyz_tensor.shape().dim_size(1); 52 | int n = input_xyz_tensor.shape().dim_size(2); 53 | 54 | const Tensor& query_xyz_tensor = context->input(1); 55 | OP_REQUIRES(context, query_xyz_tensor.dims() == 3, errors::InvalidArgument("OrderNeighbors expects (b,m_q,n) query_xyz shape")); 56 | 57 | int m_q = query_xyz_tensor.shape().dim_size(1); 58 | 59 | const Tensor& query_normals_tensor = context->input(2); 60 | OP_REQUIRES(context, query_normals_tensor.dims() == 3, errors::InvalidArgument("OrderNeighbors expects (b,m_q,n) query_normals shape")); 61 | 62 | const Tensor& idx_tensor = context->input(3); 63 | OP_REQUIRES(context, idx_tensor.dims() == 3, errors::InvalidArgument("OrderNeighbors expects (b,m_q,k) idx shape")); 64 | 65 | int k = idx_tensor.shape().dim_size(2); 66 | 67 | Tensor *outi_tensor = nullptr; 68 | OP_REQUIRES_OK(context, context->allocate_output(0, TensorShape{b,m_q,k}, &outi_tensor)); 69 | Tensor *proj_tensor = nullptr; 70 | OP_REQUIRES_OK(context, context->allocate_output(1, TensorShape{b,m_q,k,n}, &proj_tensor)); 71 | Tensor *angles_tensor = nullptr; 72 | OP_REQUIRES_OK(context, context->allocate_output(2, TensorShape{b,m_q,k}, &angles_tensor)); 73 | 74 | auto input_flat = input_xyz_tensor.flat(); 75 | const float *input = &(input_flat(0)); 76 | auto queries_flat = query_xyz_tensor.flat(); 77 | const float *queries = &(queries_flat(0)); 78 | auto queries_norm_flat = query_normals_tensor.flat(); 79 | const float *queries_norm = &(queries_norm_flat(0)); 80 | auto idx_flat = idx_tensor.flat(); 81 | const int *idx = &(idx_flat(0)); 82 | auto outi_flat = outi_tensor->flat(); 83 | int *outi = &(outi_flat(0)); 84 | auto proj_flat = proj_tensor->flat(); 85 | float *proj = &(proj_flat(0)); 86 | cudaMemset(proj, 0.0, sizeof(float)*b*m_q*k*n); 87 | auto angles_flat = angles_tensor->flat(); 88 | float *angles = &(angles_flat(0)); 89 | cudaMemset(angles, 0.0, sizeof(float)*b*m_q*k); 90 | orderNeighborsLauncher(b, m, n, m_q, k, input, queries, queries_norm, idx, proj, outi, angles); 91 | } 92 | 93 | private: 94 | int k_; 95 | }; 96 | REGISTER_KERNEL_BUILDER(Name("OrderNeighbors").Device(DEVICE_GPU), OrderNeighborsGpuOp); 97 | -------------------------------------------------------------------------------- /tf_ops/ordering/tf_ordering.py: -------------------------------------------------------------------------------- 1 | ''' The module orders neighbors in counterclockwise manner. 2 | Author: Artem Komarichev 3 | All Rights Reserved. 2018. 4 | ''' 5 | 6 | import tensorflow as tf 7 | from tensorflow.python.framework import ops 8 | import sys 9 | import os 10 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 11 | sys.path.append(BASE_DIR) 12 | ordering_module=tf.load_op_library(os.path.join(BASE_DIR, 'tf_ordering_so.so')) 13 | import numpy as np 14 | 15 | def order_neighbors(k, idx, input_xyz, query_xyz, query_normals): 16 | ''' 17 | Order neighbors in counterclockwise manner. 18 | 19 | Input: 20 | input_xyz: (batch_size, ndataset, c) float32 array, input points 21 | query_xyz: (batch_size, npoint, c) float32 array, query points 22 | idx: (batch_size, npoint, k) int32 array, indecies of the k neighbor points 23 | Output: 24 | outi: (batch_size, npoint, k) int32 array, points orderred courterclockwise 25 | proj: (batch_size, npoint, k, 3) float32 array, projected neighbors on the local tangent plane 26 | angles: (batch_size, npoint, k) float32 array, values represents angles [0, 360) 27 | ''' 28 | return ordering_module.order_neighbors(input_xyz, query_xyz, query_normals, idx, k) 29 | ops.NoGradient("OrderNeighbors") -------------------------------------------------------------------------------- /tf_ops/ordering/tf_ordering_compile.sh: -------------------------------------------------------------------------------- 1 | #/bin/bash 2 | /usr/local/cuda-8.0/bin/nvcc tf_ordering_g.cu -o tf_ordering_g.cu.o -c -O2 -DGOOGLE_CUDA=1 -x cu -Xcompiler -fPIC 3 | g++ -std=c++11 tf_ordering.cpp tf_ordering_g.cu.o -o tf_ordering_so.so -shared -fPIC -I /../anaconda3/lib/python3.6/site-packages/tensorflow/include -I /usr/local/cuda-8.0/include -lcudart -L /usr/local/cuda-8.0/lib64/ -O2 -D_GLIBCXX_USE_CXX11_ABI=0 -------------------------------------------------------------------------------- /tf_ops/ordering/tf_ordering_g.cu: -------------------------------------------------------------------------------- 1 | /* The module orders neighbors in counterclockwise manner. 2 | * Author: Artem Komarichev 3 | * All Rights Reserved. 2018. 4 | */ 5 | 6 | #include 7 | #include 8 | #include 9 | #include 10 | 11 | __global__ void order_neighbors_gpu(int b, int m, int n, int m_q, int k, const float *input, const float *queries, const float *queries_norm, const int *idx, float *proj, int *outi, float *angles){ 12 | int batch_index = blockIdx.x; 13 | queries+=m_q*n*batch_index; 14 | queries_norm+=m_q*n*batch_index; 15 | idx+=m_q*k*batch_index; 16 | angles+=m_q*k*batch_index; 17 | outi+=m_q*k*batch_index; 18 | input+=m*n*batch_index; 19 | proj+=m_q*k*n*batch_index; 20 | 21 | int index = threadIdx.x; 22 | int stride = blockDim.x; 23 | 24 | // copy indecies from idx to outi 25 | for (int i=index; iangles[i*k+max-1]) { 92 | max = t; 93 | } 94 | } 95 | 96 | if (max!=s) { 97 | // swap angles 98 | float tmp = angles[i*k+max-1]; 99 | angles[i*k+max-1] = angles[i*k+s-1]; 100 | angles[i*k+s-1] = tmp; 101 | 102 | // swap indecies 103 | int tmpi = outi[i*k+max]; 104 | outi[i*k+max] = outi[i*k+s]; 105 | outi[i*k+s] = tmpi; 106 | 107 | // swap projections as well 108 | for(int l=0;l<3;++l){ 109 | float tmpl = proj[(i*k+max)*n+l]; 110 | proj[(i*k+max)*n+l] = proj[(i*k+s)*n+l]; 111 | proj[(i*k+s)*n+l] = tmpl; 112 | } 113 | } 114 | } 115 | } 116 | } 117 | 118 | 119 | void orderNeighborsLauncher(int b, int m, int n, int m_q, int k, const float *input, const float *queries, const float *queries_norm, const int *idx, float *proj, int *outi, float *angles){ 120 | order_neighbors_gpu<<>>(b,m,n,m_q,k,input,queries,queries_norm,idx,proj,outi,angles); 121 | cudaDeviceSynchronize(); 122 | } 123 | -------------------------------------------------------------------------------- /tf_ops/sampling/.gitignore: -------------------------------------------------------------------------------- 1 | *.o 2 | *.so 3 | __pycache__/ -------------------------------------------------------------------------------- /tf_ops/sampling/tf_sampling.cpp: -------------------------------------------------------------------------------- 1 | /* Furthest point sampling 2 | * Original author: Haoqiang Fan 3 | * Modified by Charles R. Qi 4 | * All Rights Reserved. 2017. 5 | */ 6 | #include "tensorflow/core/framework/op.h" 7 | #include "tensorflow/core/framework/op_kernel.h" 8 | #include "tensorflow/core/framework/shape_inference.h" 9 | #include "tensorflow/core/framework/common_shape_fns.h" 10 | #include 11 | 12 | using namespace tensorflow; 13 | 14 | REGISTER_OP("ProbSample") 15 | .Input("inp: float32") 16 | .Input("inpr: float32") 17 | .Output("out: int32") 18 | .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) { 19 | ::tensorflow::shape_inference::ShapeHandle dims1; // batch_size * ncategory 20 | c->WithRank(c->input(0), 2, &dims1); 21 | ::tensorflow::shape_inference::ShapeHandle dims2; // batch_size * npoints 22 | c->WithRank(c->input(1), 2, &dims2); 23 | // batch_size * npoints 24 | ::tensorflow::shape_inference::ShapeHandle output = c->MakeShape({c->Dim(dims2, 0), c->Dim(dims2, 1)}); 25 | c->set_output(0, output); 26 | return Status::OK(); 27 | }); 28 | REGISTER_OP("FarthestPointSample") 29 | .Attr("npoint: int") 30 | .Input("inp: float32") 31 | .Output("out: int32") 32 | .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) { 33 | ::tensorflow::shape_inference::ShapeHandle dims1; // batch_size * npoint * 3 34 | c->WithRank(c->input(0), 3, &dims1); 35 | int npoint; 36 | TF_RETURN_IF_ERROR(c->GetAttr("npoint", &npoint)); 37 | ::tensorflow::shape_inference::ShapeHandle output = c->MakeShape({c->Dim(dims1, 0), npoint}); 38 | c->set_output(0, output); 39 | return Status::OK(); 40 | }); 41 | REGISTER_OP("GatherPoint") 42 | .Input("inp: float32") 43 | .Input("idx: int32") 44 | .Output("out: float32") 45 | .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) { 46 | ::tensorflow::shape_inference::ShapeHandle dims1; // batch_size * ndataset * 3 47 | c->WithRank(c->input(0), 3, &dims1); 48 | ::tensorflow::shape_inference::ShapeHandle dims2; // batch_size * npoints 49 | c->WithRank(c->input(1), 2, &dims2); 50 | // batch_size * npoints * 3 51 | ::tensorflow::shape_inference::ShapeHandle output = c->MakeShape({c->Dim(dims1, 0), c->Dim(dims2, 1), c->Dim(dims1, 2)}); 52 | c->set_output(0, output); 53 | return Status::OK(); 54 | }); 55 | REGISTER_OP("GatherPointGrad") 56 | .Input("inp: float32") 57 | .Input("idx: int32") 58 | .Input("out_g: float32") 59 | .Output("inp_g: float32") 60 | .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) { 61 | c->set_output(0, c->input(0)); 62 | return Status::OK(); 63 | }); 64 | 65 | void probsampleLauncher(int b,int n,int m,const float * inp_p,const float * inp_r,float * temp,int * out); 66 | class ProbSampleGpuOp: public OpKernel{ 67 | public: 68 | explicit ProbSampleGpuOp(OpKernelConstruction* context):OpKernel(context){} 69 | void Compute(OpKernelContext * context)override{ 70 | const Tensor& inp_tensor=context->input(0); 71 | const Tensor& inpr_tensor=context->input(1); 72 | auto inp_flat=inp_tensor.flat(); 73 | auto inpr_flat=inpr_tensor.flat(); 74 | const float * inp=&(inp_flat(0)); 75 | const float * inpr=&(inpr_flat(0)); 76 | OP_REQUIRES(context,inp_tensor.dims()==2,errors::InvalidArgument("ProbSample expects (batch_size,num_choices) inp shape")); 77 | int b=inp_tensor.shape().dim_size(0); 78 | int n=inp_tensor.shape().dim_size(1); 79 | OP_REQUIRES(context,inpr_tensor.dims()==2 && inpr_tensor.shape().dim_size(0)==b,errors::InvalidArgument("ProbSample expects (batch_size,num_points) inpr shape")); 80 | int m=inpr_tensor.shape().dim_size(1); 81 | Tensor * out_tensor=NULL; 82 | OP_REQUIRES_OK(context,context->allocate_output(0,TensorShape{b,m},&out_tensor)); 83 | auto out_flat=out_tensor->flat(); 84 | int * out=&(out_flat(0)); 85 | Tensor temp_tensor; 86 | OP_REQUIRES_OK(context,context->allocate_temp(DataTypeToEnum::value,TensorShape{b,n},&temp_tensor)); 87 | auto temp_flat=temp_tensor.flat(); 88 | float * temp=&(temp_flat(0)); 89 | probsampleLauncher(b,n,m,inp,inpr,temp,out); 90 | } 91 | }; 92 | REGISTER_KERNEL_BUILDER(Name("ProbSample").Device(DEVICE_GPU), ProbSampleGpuOp); 93 | 94 | void farthestpointsamplingLauncher(int b,int n,int m,const float * inp,float * temp,int * out); 95 | class FarthestPointSampleGpuOp: public OpKernel{ 96 | public: 97 | explicit FarthestPointSampleGpuOp(OpKernelConstruction* context):OpKernel(context) { 98 | OP_REQUIRES_OK(context, context->GetAttr("npoint", &npoint_)); 99 | OP_REQUIRES(context, npoint_ > 0, errors::InvalidArgument("FarthestPointSample expects positive npoint")); 100 | } 101 | void Compute(OpKernelContext * context)override{ 102 | int m = npoint_; 103 | 104 | const Tensor& inp_tensor=context->input(0); 105 | OP_REQUIRES(context,inp_tensor.dims()==3 && inp_tensor.shape().dim_size(2)==3,errors::InvalidArgument("FarthestPointSample expects (batch_size,num_points,3) inp shape")); 106 | int b=inp_tensor.shape().dim_size(0); 107 | int n=inp_tensor.shape().dim_size(1); 108 | auto inp_flat=inp_tensor.flat(); 109 | const float * inp=&(inp_flat(0)); 110 | Tensor * out_tensor; 111 | OP_REQUIRES_OK(context,context->allocate_output(0,TensorShape{b,m},&out_tensor)); 112 | auto out_flat=out_tensor->flat(); 113 | int * out=&(out_flat(0)); 114 | Tensor temp_tensor; 115 | OP_REQUIRES_OK(context,context->allocate_temp(DataTypeToEnum::value,TensorShape{32,n},&temp_tensor)); 116 | auto temp_flat=temp_tensor.flat(); 117 | float * temp=&(temp_flat(0)); 118 | farthestpointsamplingLauncher(b,n,m,inp,temp,out); 119 | } 120 | private: 121 | int npoint_; 122 | }; 123 | REGISTER_KERNEL_BUILDER(Name("FarthestPointSample").Device(DEVICE_GPU),FarthestPointSampleGpuOp); 124 | 125 | void gatherpointLauncher(int b,int n,int m,const float * inp,const int * idx,float * out); 126 | class GatherPointGpuOp: public OpKernel{ 127 | public: 128 | explicit GatherPointGpuOp(OpKernelConstruction * context):OpKernel(context){} 129 | void Compute(OpKernelContext * context)override{ 130 | const Tensor& inp_tensor=context->input(0); 131 | OP_REQUIRES(context,inp_tensor.dims()==3 && inp_tensor.shape().dim_size(2)==3,errors::InvalidArgument("GatherPoint expects (batch_size,num_points,3) inp shape")); 132 | int b=inp_tensor.shape().dim_size(0); 133 | int n=inp_tensor.shape().dim_size(1); 134 | const Tensor& idx_tensor=context->input(1); 135 | OP_REQUIRES(context,idx_tensor.dims()==2 && idx_tensor.shape().dim_size(0)==b,errors::InvalidArgument("GatherPoint expects (batch_size,num_result) idx shape")); 136 | int m=idx_tensor.shape().dim_size(1); 137 | auto inp_flat=inp_tensor.flat(); 138 | const float * inp=&(inp_flat(0)); 139 | auto idx_flat=idx_tensor.flat(); 140 | const int * idx=&(idx_flat(0)); 141 | Tensor * out_tensor=NULL; 142 | OP_REQUIRES_OK(context,context->allocate_output(0,TensorShape{b,m,3},&out_tensor)); 143 | auto out_flat=out_tensor->flat(); 144 | float * out=&(out_flat(0)); 145 | gatherpointLauncher(b,n,m,inp,idx,out); 146 | } 147 | }; 148 | REGISTER_KERNEL_BUILDER(Name("GatherPoint").Device(DEVICE_GPU),GatherPointGpuOp); 149 | 150 | void scatteraddpointLauncher(int b,int n,int m,const float * out_g,const int * idx,float * inp_g); 151 | class GatherPointGradGpuOp: public OpKernel{ 152 | public: 153 | explicit GatherPointGradGpuOp(OpKernelConstruction * context):OpKernel(context){} 154 | void Compute(OpKernelContext * context)override{ 155 | const Tensor& inp_tensor=context->input(0); 156 | OP_REQUIRES(context,inp_tensor.dims()==3 && inp_tensor.shape().dim_size(2)==3,errors::InvalidArgument("GatherPointGradGpuOp expects (batch_size,num_points,3) inp")); 157 | int b=inp_tensor.shape().dim_size(0); 158 | int n=inp_tensor.shape().dim_size(1); 159 | const Tensor& idx_tensor=context->input(1); 160 | OP_REQUIRES(context,idx_tensor.dims()==2 && idx_tensor.shape().dim_size(0)==b,errors::InvalidArgument("GatherPointGradGpuOp expects (batch_size,num_result) idx shape")); 161 | int m=idx_tensor.shape().dim_size(1); 162 | auto inp_flat=inp_tensor.flat(); 163 | const float * inp=&(inp_flat(0)); 164 | auto idx_flat=idx_tensor.flat(); 165 | const int * idx=&(idx_flat(0)); 166 | const Tensor& out_g_tensor=context->input(2); 167 | OP_REQUIRES(context,out_g_tensor.dims()==3 && out_g_tensor.shape().dim_size(0)==b && out_g_tensor.shape().dim_size(1)==m && out_g_tensor.shape().dim_size(2)==3,errors::InvalidArgument("GatherPointGradGpuOp expects (batch_size,num_result,3) out_g shape")); 168 | auto out_g_flat=out_g_tensor.flat(); 169 | const float * out_g=&(out_g_flat(0)); 170 | Tensor * inp_g_tensor=NULL; 171 | OP_REQUIRES_OK(context,context->allocate_output(0,TensorShape{b,n,3},&inp_g_tensor)); 172 | auto inp_g_flat=inp_g_tensor->flat(); 173 | float * inp_g=&(inp_g_flat(0)); 174 | cudaMemset(inp_g,0,b*n*3*4); 175 | scatteraddpointLauncher(b,n,m,out_g,idx,inp_g); 176 | } 177 | }; 178 | REGISTER_KERNEL_BUILDER(Name("GatherPointGrad").Device(DEVICE_GPU),GatherPointGradGpuOp); 179 | 180 | -------------------------------------------------------------------------------- /tf_ops/sampling/tf_sampling.py: -------------------------------------------------------------------------------- 1 | ''' Furthest point sampling 2 | Original author: Haoqiang Fan 3 | Modified by Charles R. Qi 4 | All Rights Reserved. 2017. 5 | ''' 6 | import tensorflow as tf 7 | from tensorflow.python.framework import ops 8 | import sys 9 | import os 10 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 11 | sys.path.append(BASE_DIR) 12 | sampling_module=tf.load_op_library(os.path.join(BASE_DIR, 'tf_sampling_so.so')) 13 | def prob_sample(inp,inpr): 14 | ''' 15 | input: 16 | batch_size * ncategory float32 17 | batch_size * npoints float32 18 | returns: 19 | batch_size * npoints int32 20 | ''' 21 | return sampling_module.prob_sample(inp,inpr) 22 | ops.NoGradient('ProbSample') 23 | # TF1.0 API requires set shape in C++ 24 | #@tf.RegisterShape('ProbSample') 25 | #def _prob_sample_shape(op): 26 | # shape1=op.inputs[0].get_shape().with_rank(2) 27 | # shape2=op.inputs[1].get_shape().with_rank(2) 28 | # return [tf.TensorShape([shape2.dims[0],shape2.dims[1]])] 29 | def gather_point(inp,idx): 30 | ''' 31 | input: 32 | batch_size * ndataset * 3 float32 33 | batch_size * npoints int32 34 | returns: 35 | batch_size * npoints * 3 float32 36 | ''' 37 | return sampling_module.gather_point(inp,idx) 38 | #@tf.RegisterShape('GatherPoint') 39 | #def _gather_point_shape(op): 40 | # shape1=op.inputs[0].get_shape().with_rank(3) 41 | # shape2=op.inputs[1].get_shape().with_rank(2) 42 | # return [tf.TensorShape([shape1.dims[0],shape2.dims[1],shape1.dims[2]])] 43 | @tf.RegisterGradient('GatherPoint') 44 | def _gather_point_grad(op,out_g): 45 | inp=op.inputs[0] 46 | idx=op.inputs[1] 47 | return [sampling_module.gather_point_grad(inp,idx,out_g),None] 48 | def farthest_point_sample(npoint,inp): 49 | ''' 50 | input: 51 | int32 52 | batch_size * ndataset * 3 float32 53 | returns: 54 | batch_size * npoint int32 55 | ''' 56 | return sampling_module.farthest_point_sample(inp, npoint) 57 | ops.NoGradient('FarthestPointSample') 58 | 59 | 60 | if __name__=='__main__': 61 | import numpy as np 62 | np.random.seed(100) 63 | triangles=np.random.rand(1,5,3,3).astype('float32') 64 | with tf.device('/gpu:1'): 65 | inp=tf.constant(triangles) 66 | tria=inp[:,:,0,:] 67 | trib=inp[:,:,1,:] 68 | tric=inp[:,:,2,:] 69 | areas=tf.sqrt(tf.reduce_sum(tf.cross(trib-tria,tric-tria)**2,2)+1e-9) 70 | randomnumbers=tf.random_uniform((1,8192)) 71 | triids=prob_sample(areas,randomnumbers) 72 | tria_sample=gather_point(tria,triids) 73 | trib_sample=gather_point(trib,triids) 74 | tric_sample=gather_point(tric,triids) 75 | us=tf.random_uniform((1,8192)) 76 | vs=tf.random_uniform((1,8192)) 77 | uplusv=1-tf.abs(us+vs-1) 78 | uminusv=us-vs 79 | us=(uplusv+uminusv)*0.5 80 | vs=(uplusv-uminusv)*0.5 81 | pt_sample=tria_sample+(trib_sample-tria_sample)*tf.expand_dims(us,-1)+(tric_sample-tria_sample)*tf.expand_dims(vs,-1) 82 | print('pt_sample: ', pt_sample) 83 | reduced_sample=gather_point(pt_sample,farthest_point_sample(1024,pt_sample)) 84 | print(reduced_sample) 85 | with tf.Session('') as sess: 86 | ret=sess.run(reduced_sample) 87 | print(ret.shape,ret.dtype) 88 | import cPickle as pickle 89 | pickle.dump(ret,open('1.pkl','wb'),-1) 90 | -------------------------------------------------------------------------------- /tf_ops/sampling/tf_sampling_compile.sh: -------------------------------------------------------------------------------- 1 | #/bin/bash 2 | /usr/local/cuda-8.0/bin/nvcc tf_sampling_g.cu -o tf_sampling_g.cu.o -c -O2 -DGOOGLE_CUDA=1 -x cu -Xcompiler -fPIC 3 | g++ -std=c++11 tf_sampling.cpp tf_sampling_g.cu.o -o tf_sampling_so.so -shared -fPIC -I /../anaconda3/lib/python3.6/site-packages/tensorflow/include -I /usr/local/cuda-8.0/include -lcudart -L /usr/local/cuda-8.0/lib64/ -O2 -D_GLIBCXX_USE_CXX11_ABI=0 -------------------------------------------------------------------------------- /tf_ops/sampling/tf_sampling_g.cu: -------------------------------------------------------------------------------- 1 | /* Furthest point sampling GPU implementation 2 | * Original author: Haoqiang Fan 3 | * Modified by Charles R. Qi 4 | * All Rights Reserved. 2017. 5 | */ 6 | 7 | __global__ void cumsumKernel(int b,int n,const float * __restrict__ inp,float * __restrict__ out){ 8 | const int BlockSize=2048; 9 | const int paddingLevel=5; 10 | __shared__ float buffer4[BlockSize*4]; 11 | __shared__ float buffer[BlockSize+(BlockSize>>paddingLevel)]; 12 | for (int i=blockIdx.x;i>2; 18 | for (int k=threadIdx.x*4;k>2)+(k>>(2+paddingLevel))]=v4; 33 | }else{ 34 | float v=0; 35 | for (int k2=k;k2>2)+(k>>(2+paddingLevel))]=v; 43 | } 44 | } 45 | int u=0; 46 | for (;(2<>(u+1));k+=blockDim.x){ 49 | int i1=(((k<<1)+2)<>paddingLevel; 52 | i2+=i2>>paddingLevel; 53 | buffer[i1]+=buffer[i2]; 54 | } 55 | } 56 | u--; 57 | for (;u>=0;u--){ 58 | __syncthreads(); 59 | for (int k=threadIdx.x;k>(u+1));k+=blockDim.x){ 60 | int i1=(((k<<1)+3)<>paddingLevel; 63 | i2+=i2>>paddingLevel; 64 | buffer[i1]+=buffer[i2]; 65 | } 66 | } 67 | __syncthreads(); 68 | for (int k=threadIdx.x*4;k>2)-1)+(((k>>2)-1)>>paddingLevel); 71 | buffer4[k]+=buffer[k2]; 72 | buffer4[k+1]+=buffer[k2]; 73 | buffer4[k+2]+=buffer[k2]; 74 | buffer4[k+3]+=buffer[k2]; 75 | } 76 | } 77 | __syncthreads(); 78 | for (int k=threadIdx.x;k>paddingLevel)]+runningsum2; 82 | float r2=runningsum+t; 83 | runningsum2=t-(r2-runningsum); 84 | runningsum=r2; 85 | __syncthreads(); 86 | } 87 | } 88 | } 89 | 90 | __global__ void binarysearchKernel(int b,int n,int m,const float * __restrict__ dataset,const float * __restrict__ query, int * __restrict__ result){ 91 | int base=1; 92 | while (base=1;k>>=1) 99 | if (r>=k && dataset[i*n+r-k]>=q) 100 | r-=k; 101 | result[i*m+j]=r; 102 | } 103 | } 104 | } 105 | __global__ void farthestpointsamplingKernel(int b,int n,int m,const float * __restrict__ dataset,float * __restrict__ temp,int * __restrict__ idxs){ 106 | if (m<=0) 107 | return; 108 | const int BlockSize=512; 109 | __shared__ float dists[BlockSize]; 110 | __shared__ int dists_i[BlockSize]; 111 | const int BufferSize=3072; 112 | __shared__ float buf[BufferSize*3]; 113 | for (int i=blockIdx.x;ibest){ 147 | best=d2; 148 | besti=k; 149 | } 150 | } 151 | dists[threadIdx.x]=best; 152 | dists_i[threadIdx.x]=besti; 153 | for (int u=0;(1<>(u+1))){ 156 | int i1=(threadIdx.x*2)<>>(b,n,inp,out); 196 | } 197 | //require b*n working space 198 | void probsampleLauncher(int b,int n,int m,const float * inp_p,const float * inp_r,float * temp,int * out){ 199 | cumsumKernel<<<32,512>>>(b,n,inp_p,temp); 200 | binarysearchKernel<<>>(b,n,m,temp,inp_r,out); 201 | } 202 | //require 32*n working space 203 | void farthestpointsamplingLauncher(int b,int n,int m,const float * inp,float * temp,int * out){ 204 | farthestpointsamplingKernel<<<32,512>>>(b,n,m,inp,temp,out); 205 | } 206 | void gatherpointLauncher(int b,int n,int m,const float * inp,const int * idx,float * out){ 207 | gatherpointKernel<<>>(b,n,m,inp,idx,out); 208 | } 209 | void scatteraddpointLauncher(int b,int n,int m,const float * out_g,const int * idx,float * inp_g){ 210 | scatteraddpointKernel<<>>(b,n,m,out_g,idx,inp_g); 211 | } 212 | 213 | -------------------------------------------------------------------------------- /train.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import math 3 | from datetime import datetime 4 | import h5py 5 | import numpy as np 6 | import tensorflow as tf 7 | import socket 8 | import importlib 9 | import os 10 | import sys 11 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 12 | ROOT_DIR = os.path.dirname(BASE_DIR) 13 | DATA_DIR = os.path.dirname(ROOT_DIR) 14 | sys.path.append(BASE_DIR) 15 | sys.path.append(os.path.join(BASE_DIR, 'models')) 16 | sys.path.append(os.path.join(BASE_DIR, 'utils')) 17 | import provider 18 | import tf_util 19 | import modelnet_dataset 20 | 21 | parser = argparse.ArgumentParser() 22 | parser.add_argument('--gpu', type=int, default=0, help='GPU to use [default: GPU 0]') 23 | parser.add_argument('--model', default='acnn_cls_rings', help='Model name') 24 | parser.add_argument('--log_dir', default='log', help='Log dir [default: log]') 25 | parser.add_argument('--num_point', type=int, default=1024, help='Point Number [default: 1024]') 26 | parser.add_argument('--max_epoch', type=int, default=251, help='Epoch to run [default: 251]') 27 | parser.add_argument('--batch_size', type=int, default=16, help='Batch Size during training [default: 16]') 28 | parser.add_argument('--learning_rate', type=float, default=0.001, help='Initial learning rate [default: 0.001]') 29 | parser.add_argument('--momentum', type=float, default=0.9, help='Initial learning rate [default: 0.9]') 30 | parser.add_argument('--optimizer', default='adam', help='adam or momentum [default: adam]') 31 | parser.add_argument('--decay_step', type=int, default=200000, help='Decay step for lr decay [default: 200000]') 32 | parser.add_argument('--decay_rate', type=float, default=0.7, help='Decay rate for lr decay [default: 0.7]') 33 | parser.add_argument('--normal', action='store_true', default='store_true', help='Whether to use normal information') 34 | FLAGS = parser.parse_args() 35 | 36 | EPOCH_CNT = 0 37 | 38 | BATCH_SIZE = FLAGS.batch_size 39 | NUM_POINT = FLAGS.num_point 40 | MAX_EPOCH = FLAGS.max_epoch 41 | BASE_LEARNING_RATE = FLAGS.learning_rate 42 | GPU_INDEX = FLAGS.gpu 43 | MOMENTUM = FLAGS.momentum 44 | OPTIMIZER = FLAGS.optimizer 45 | DECAY_STEP = FLAGS.decay_step 46 | DECAY_RATE = FLAGS.decay_rate 47 | 48 | MODEL = importlib.import_module(FLAGS.model) # import network module 49 | MODEL_FILE = os.path.join(ROOT_DIR, 'models', FLAGS.model+'.py') 50 | LOG_DIR = FLAGS.log_dir 51 | if not os.path.exists(LOG_DIR): os.mkdir(LOG_DIR) 52 | os.system('cp %s %s' % (MODEL_FILE, LOG_DIR)) # bkp of model def 53 | os.system('cp train.py %s' % (LOG_DIR)) # bkp of train procedure 54 | LOG_FOUT = open(os.path.join(LOG_DIR, 'log_train.txt'), 'w') 55 | LOG_FOUT.write(str(FLAGS)+'\n') 56 | 57 | BN_INIT_DECAY = 0.5 58 | BN_DECAY_DECAY_RATE = 0.5 59 | BN_DECAY_DECAY_STEP = float(DECAY_STEP) 60 | BN_DECAY_CLIP = 0.99 61 | 62 | HOSTNAME = socket.gethostname() 63 | 64 | NUM_CLASSES = 40 65 | 66 | # Shapenet official train/test split 67 | if FLAGS.normal: 68 | assert(NUM_POINT<=10000) 69 | DATA_PATH = os.path.join(DATA_DIR, 'data_10K/modelnet40_normal_resampled') 70 | TRAIN_DATASET = modelnet_dataset.ModelNetDataset(root=DATA_PATH, npoints=NUM_POINT, split='train', normal_channel=FLAGS.normal, batch_size=BATCH_SIZE) 71 | TEST_DATASET = modelnet_dataset.ModelNetDataset(root=DATA_PATH, npoints=NUM_POINT, split='test', normal_channel=FLAGS.normal, batch_size=BATCH_SIZE) 72 | 73 | def log_string(out_str): 74 | LOG_FOUT.write(out_str+'\n') 75 | LOG_FOUT.flush() 76 | print(out_str) 77 | 78 | def get_learning_rate(batch): 79 | learning_rate = tf.train.exponential_decay( 80 | BASE_LEARNING_RATE, # Base learning rate. 81 | batch * BATCH_SIZE, # Current index into the dataset. 82 | DECAY_STEP, # Decay step. 83 | DECAY_RATE, # Decay rate. 84 | staircase=True) 85 | learning_rate = tf.maximum(learning_rate, 0.00001) # CLIP THE LEARNING RATE! 86 | return learning_rate 87 | 88 | def get_bn_decay(batch): 89 | bn_momentum = tf.train.exponential_decay( 90 | BN_INIT_DECAY, 91 | batch*BATCH_SIZE, 92 | BN_DECAY_DECAY_STEP, 93 | BN_DECAY_DECAY_RATE, 94 | staircase=True) 95 | bn_decay = tf.minimum(BN_DECAY_CLIP, 1 - bn_momentum) 96 | return bn_decay 97 | 98 | def train(): 99 | with tf.Graph().as_default(): 100 | with tf.device('/gpu:'+str(GPU_INDEX)): 101 | pointclouds_pl, labels_pl, normals_pl = MODEL.placeholder_inputs(BATCH_SIZE, NUM_POINT) 102 | is_training_pl = tf.placeholder(tf.bool, shape=()) 103 | 104 | # Note the global_step=batch parameter to minimize. 105 | # That tells the optimizer to helpfully increment the 'batch' parameter 106 | # for you every time it trains. 107 | batch = tf.get_variable('batch', [], 108 | initializer=tf.constant_initializer(0), trainable=False) 109 | bn_decay = get_bn_decay(batch) 110 | tf.summary.scalar('bn_decay', bn_decay) 111 | 112 | # Get model and loss 113 | pred, end_points = MODEL.get_model(pointclouds_pl, normals_pl, is_training_pl, bn_decay=bn_decay) 114 | MODEL.get_loss(pred, labels_pl, end_points) 115 | losses = tf.get_collection('losses') 116 | total_loss = tf.add_n(losses, name='total_loss') 117 | tf.summary.scalar('total_loss', total_loss) 118 | for l in losses + [total_loss]: 119 | tf.summary.scalar(l.op.name, l) 120 | 121 | correct = tf.equal(tf.argmax(pred, 1), tf.to_int64(labels_pl)) 122 | accuracy = tf.reduce_sum(tf.cast(correct, tf.float32)) / float(BATCH_SIZE) 123 | tf.summary.scalar('accuracy', accuracy) 124 | 125 | print("--- Get training operator") 126 | # Get training operator 127 | learning_rate = get_learning_rate(batch) 128 | tf.summary.scalar('learning_rate', learning_rate) 129 | if OPTIMIZER == 'momentum': 130 | optimizer = tf.train.MomentumOptimizer(learning_rate, momentum=MOMENTUM) 131 | elif OPTIMIZER == 'adam': 132 | optimizer = tf.train.AdamOptimizer(learning_rate) 133 | train_op = optimizer.minimize(total_loss, global_step=batch) 134 | 135 | # Add ops to save and restore all the variables. 136 | saver = tf.train.Saver() 137 | 138 | # Create a session 139 | config = tf.ConfigProto() 140 | config.gpu_options.allow_growth = True 141 | config.allow_soft_placement = True 142 | config.log_device_placement = False 143 | sess = tf.Session(config=config) 144 | 145 | # Add summary writers 146 | merged = tf.summary.merge_all() 147 | train_writer = tf.summary.FileWriter(os.path.join(LOG_DIR, 'train'), sess.graph) 148 | test_writer = tf.summary.FileWriter(os.path.join(LOG_DIR, 'test'), sess.graph) 149 | 150 | # Init variables 151 | init = tf.global_variables_initializer() 152 | sess.run(init) 153 | 154 | ops = {'pointclouds_pl': pointclouds_pl, 155 | 'normals_pl': normals_pl, 156 | 'labels_pl': labels_pl, 157 | 'is_training_pl': is_training_pl, 158 | 'pred': pred, 159 | 'loss': total_loss, 160 | 'train_op': train_op, 161 | 'merged': merged, 162 | 'step': batch, 163 | 'end_points': end_points} 164 | 165 | best_acc = -1 166 | best_acc_avg = -1 167 | for epoch in range(MAX_EPOCH): 168 | log_string('**** EPOCH %03d ****' % (epoch)) 169 | sys.stdout.flush() 170 | 171 | train_one_epoch(sess, ops, train_writer) 172 | acc, acc_avg = eval_one_epoch(sess, ops, test_writer) 173 | 174 | if acc > best_acc: 175 | best_acc = acc 176 | save_path = saver.save(sess, os.path.join(LOG_DIR, "model_best_acc.ckpt")) 177 | log_string("Model saved in file: %s" % save_path) 178 | log_string("Best category accuracy: %f" % best_acc) 179 | 180 | if acc_avg > best_acc_avg: 181 | best_acc_avg = acc_avg 182 | save_path = saver.save(sess, os.path.join(LOG_DIR, "model_best_acc_avg.ckpt")) 183 | log_string("Model saved in file: %s" % save_path) 184 | log_string("Best instance accuracy: %f" % best_acc_avg) 185 | 186 | # Save the variables to disk. 187 | if epoch % 10 == 0: 188 | save_path = saver.save(sess, os.path.join(LOG_DIR, "model.ckpt")) 189 | log_string("Model saved in file: %s" % save_path) 190 | 191 | 192 | def train_one_epoch(sess, ops, train_writer): 193 | """ ops: dict mapping from string to tf ops """ 194 | is_training = True 195 | 196 | log_string(str(datetime.now())) 197 | 198 | # Make sure batch data is of same size 199 | cur_batch_data = np.zeros((BATCH_SIZE,NUM_POINT,3)) 200 | cur_batch_normals = np.zeros((BATCH_SIZE,NUM_POINT,3)) 201 | cur_batch_label = np.zeros((BATCH_SIZE), dtype=np.int32) 202 | 203 | total_correct = 0 204 | total_seen = 0 205 | loss_sum = 0 206 | batch_idx = 0 207 | while TRAIN_DATASET.has_next_batch(): 208 | batch_data, batch_label = TRAIN_DATASET.next_batch(augment=True) 209 | bsize = batch_data.shape[0] 210 | cur_batch_data[0:bsize,...] = batch_data[:,:,:3] 211 | cur_batch_normals[0:bsize,...] = batch_data[:,:,3:] 212 | cur_batch_label[0:bsize] = batch_label 213 | 214 | feed_dict = {ops['pointclouds_pl']: cur_batch_data, 215 | ops['labels_pl']: cur_batch_label, 216 | ops['normals_pl']: cur_batch_normals, 217 | ops['is_training_pl']: is_training,} 218 | summary, step, _, loss_val, pred_val = sess.run([ops['merged'], ops['step'], 219 | ops['train_op'], ops['loss'], ops['pred']], feed_dict=feed_dict) 220 | train_writer.add_summary(summary, step) 221 | pred_val = np.argmax(pred_val, 1) 222 | correct = np.sum(pred_val[0:bsize] == batch_label[0:bsize]) 223 | total_correct += correct 224 | total_seen += bsize 225 | loss_sum += loss_val 226 | 227 | if (batch_idx+1)%50 == 0: 228 | log_string(' ---- batch: %03d ----' % (batch_idx+1)) 229 | log_string('mean loss: %f' % (loss_sum / 50)) 230 | log_string('accuracy: %f' % (total_correct / float(total_seen))) 231 | total_correct = 0 232 | total_seen = 0 233 | loss_sum = 0 234 | batch_idx += 1 235 | 236 | TRAIN_DATASET.reset() 237 | 238 | def eval_one_epoch(sess, ops, test_writer): 239 | """ ops: dict mapping from string to tf ops """ 240 | global EPOCH_CNT 241 | is_training = False 242 | 243 | # Make sure batch data is of same size 244 | cur_batch_data = np.zeros((BATCH_SIZE,NUM_POINT,3)) 245 | cur_batch_normals = np.zeros((BATCH_SIZE,NUM_POINT,3)) 246 | cur_batch_label = np.zeros((BATCH_SIZE), dtype=np.int32) 247 | 248 | total_correct = 0 249 | total_seen = 0 250 | loss_sum = 0 251 | batch_idx = 0 252 | shape_ious = [] 253 | total_seen_class = [0 for _ in range(NUM_CLASSES)] 254 | total_correct_class = [0 for _ in range(NUM_CLASSES)] 255 | 256 | log_string(str(datetime.now())) 257 | log_string('---- EPOCH %03d EVALUATION ----'%(EPOCH_CNT)) 258 | 259 | while TEST_DATASET.has_next_batch(): 260 | batch_data, batch_label = TEST_DATASET.next_batch(augment=False) 261 | bsize = batch_data.shape[0] 262 | cur_batch_data[0:bsize,...] = batch_data[:,:,:3] 263 | cur_batch_normals[0:bsize,...] = batch_data[:,:,3:] 264 | cur_batch_label[0:bsize] = batch_label 265 | 266 | feed_dict = {ops['pointclouds_pl']: cur_batch_data, 267 | ops['labels_pl']: cur_batch_label, 268 | ops['normals_pl']: cur_batch_normals, 269 | ops['is_training_pl']: is_training} 270 | summary, step, loss_val, pred_val = sess.run([ops['merged'], ops['step'], 271 | ops['loss'], ops['pred']], feed_dict=feed_dict) 272 | test_writer.add_summary(summary, step) 273 | pred_val = np.argmax(pred_val, 1) 274 | correct = np.sum(pred_val[0:bsize] == batch_label[0:bsize]) 275 | total_correct += correct 276 | total_seen += bsize 277 | loss_sum += loss_val 278 | batch_idx += 1 279 | for i in range(0, bsize): 280 | l = batch_label[i] 281 | total_seen_class[l] += 1 282 | total_correct_class[l] += (pred_val[i] == l) 283 | 284 | log_string('eval mean loss: %f' % (loss_sum / float(batch_idx))) 285 | log_string('eval accuracy: %f'% (total_correct / float(total_seen))) 286 | log_string('eval avg class acc: %f' % (np.mean(np.array(total_correct_class)/np.array(total_seen_class,dtype=np.float)))) 287 | EPOCH_CNT += 1 288 | 289 | TEST_DATASET.reset() 290 | return total_correct/float(total_seen), np.mean(np.array(total_correct_class)/np.array(total_seen_class,dtype=np.float)) 291 | 292 | 293 | if __name__ == "__main__": 294 | log_string('pid: %s'%(str(os.getpid()))) 295 | train() 296 | LOG_FOUT.close() 297 | -------------------------------------------------------------------------------- /utils/pointnet_util.py: -------------------------------------------------------------------------------- 1 | """ PointNet++ Layers 2 | 3 | Author: Charles R. Qi 4 | Date: November 2017 5 | """ 6 | 7 | import os 8 | import sys 9 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 10 | ROOT_DIR = os.path.dirname(BASE_DIR) 11 | sys.path.append(os.path.join(ROOT_DIR, 'utils')) 12 | sys.path.append(os.path.join(ROOT_DIR, 'tf_ops/sampling')) 13 | sys.path.append(os.path.join(ROOT_DIR, 'tf_ops/grouping_ring')) 14 | sys.path.append(os.path.join(ROOT_DIR, 'tf_ops/3d_interpolation')) 15 | sys.path.append(os.path.join(ROOT_DIR, 'tf_ops/ordering')) 16 | from tf_sampling import farthest_point_sample, gather_point 17 | from tf_grouping_ring import ring_point, group_point, knn_point 18 | from tf_interpolate import three_nn, three_interpolate 19 | from tf_ordering import order_neighbors 20 | import tensorflow as tf 21 | import numpy as np 22 | import tf_util 23 | 24 | def sample_and_group(npoint, radius, nsample, xyz, points, tnet_spec=None, knn=False, use_xyz=True): 25 | ''' 26 | Input: 27 | npoint: int32 28 | radius: float32 29 | nsample: int32 30 | xyz: (batch_size, ndataset, 3) TF tensor 31 | points: (batch_size, ndataset, channel) TF tensor, if None will just use xyz as points 32 | tnet_spec: dict (keys: mlp, mlp2, is_training, bn_decay), if None do not apply tnet 33 | knn: bool, if True use kNN instead of radius search 34 | use_xyz: bool, if True concat XYZ with local point features, otherwise just use point features 35 | Output: 36 | new_xyz: (batch_size, npoint, 3) TF tensor 37 | new_points: (batch_size, npoint, nsample, 3+channel) TF tensor 38 | idx: (batch_size, npoint, nsample) TF tensor, indices of local points as in ndataset points 39 | grouped_xyz: (batch_size, npoint, nsample, 3) TF tensor, normalized point XYZs 40 | (subtracted by seed point XYZ) in local regions 41 | ''' 42 | 43 | indecies = farthest_point_sample(npoint, xyz) 44 | new_xyz = gather_point(xyz, indecies) # (batch_size, npoint, 3) 45 | new_normals = gather_point(normals, indecies) # (batch_size, npoint, 3) 46 | if knn: 47 | _,idx = knn_point(nsample, xyz, new_xyz) 48 | else: 49 | idx, pts_cnt = query_ball_point(radius, nsample, xyz, new_xyz) 50 | grouped_xyz = group_point(xyz, idx) # (batch_size, npoint, nsample, 3) 51 | grouped_xyz -= tf.tile(tf.expand_dims(new_xyz, 2), [1,1,nsample,1]) # translation normalization 52 | if tnet_spec is not None: 53 | grouped_xyz = tnet(grouped_xyz, tnet_spec) 54 | if points is not None: 55 | grouped_points = group_point(points, idx) # (batch_size, npoint, nsample, channel) 56 | if use_xyz: 57 | new_points = tf.concat([grouped_xyz, grouped_points], axis=-1) # (batch_size, npoint, nample, 3+channel) 58 | else: 59 | new_points = grouped_points 60 | else: 61 | new_points = grouped_xyz 62 | 63 | return new_xyz, new_points, new_normals, idx, grouped_xyz 64 | 65 | 66 | def sample_and_group_all(xyz, points, use_xyz=True): 67 | ''' 68 | Inputs: 69 | xyz: (batch_size, ndataset, 3) TF tensor 70 | points: (batch_size, ndataset, channel) TF tensor, if None will just use xyz as points 71 | use_xyz: bool, if True concat XYZ with local point features, otherwise just use point features 72 | Outputs: 73 | new_xyz: (batch_size, 1, 3) as (0,0,0) 74 | new_points: (batch_size, 1, ndataset, 3+channel) TF tensor 75 | Note: 76 | Equivalent to sample_and_group with npoint=1, radius=inf, use (0,0,0) as the centroid 77 | ''' 78 | batch_size = xyz.get_shape()[0].value 79 | nsample = xyz.get_shape()[1].value 80 | new_xyz = tf.constant(np.tile(np.array([0,0,0]).reshape((1,1,3)), (batch_size,1,1)),dtype=tf.float32) # (batch_size, 1, 3) 81 | idx = tf.constant(np.tile(np.array(range(nsample)).reshape((1,1,nsample)), (batch_size,1,1))) 82 | grouped_xyz = tf.reshape(xyz, (batch_size, 1, nsample, 3)) # (batch_size, npoint=1, nsample, 3) 83 | if points is not None: 84 | if use_xyz: 85 | new_points = tf.concat([xyz, points], axis=2) # (batch_size, 16, 259) 86 | else: 87 | new_points = points 88 | new_points = tf.expand_dims(new_points, 1) # (batch_size, 1, 16, 259) 89 | else: 90 | new_points = grouped_xyz 91 | return new_xyz, new_points, idx, grouped_xyz 92 | 93 | 94 | def pointnet_sa_module(xyz, points, npoint, radius, nsample, mlp, mlp2, group_all, is_training, bn_decay, scope, bn=True, pooling='max', tnet_spec=None, knn=False, use_xyz=True): 95 | ''' PointNet Set Abstraction (SA) Module 96 | Input: 97 | xyz: (batch_size, ndataset, 3) TF tensor 98 | points: (batch_size, ndataset, channel) TF tensor 99 | npoint: int32 -- #points sampled in farthest point sampling 100 | radius: float32 -- search radius in local region 101 | nsample: int32 -- how many points in each local region 102 | mlp: list of int32 -- output size for MLP on each point 103 | mlp2: list of int32 -- output size for MLP on each region 104 | group_all: bool -- group all points into one PC if set true, OVERRIDE 105 | npoint, radius and nsample settings 106 | use_xyz: bool, if True concat XYZ with local point features, otherwise just use point features 107 | Return: 108 | new_xyz: (batch_size, npoint, 3) TF tensor 109 | new_points: (batch_size, npoint, mlp[-1] or mlp2[-1]) TF tensor 110 | idx: (batch_size, npoint, nsample) int32 -- indices for local regions 111 | ''' 112 | with tf.variable_scope(scope) as sc: 113 | if group_all: 114 | nsample = xyz.get_shape()[1].value 115 | new_xyz, new_points, idx, grouped_xyz = sample_and_group_all(xyz, points, use_xyz) 116 | 117 | for i, num_out_channel in enumerate(mlp): 118 | new_points = tf_util.conv2d(new_points, num_out_channel, [1,1], 119 | padding='SAME', stride=[1,1], 120 | bn=bn, is_training=is_training, 121 | scope='conv%d'%(i), bn_decay=bn_decay) 122 | 123 | if pooling=='avg': 124 | new_points = tf_util.avg_pool2d(new_points, [1,nsample], stride=[1,1], padding='VALID', scope='avgpool1') 125 | elif pooling=='weighted_avg': 126 | with tf.variable_scope('weighted_avg1'): 127 | dists = tf.norm(grouped_xyz,axis=-1,ord=2,keep_dims=True) 128 | exp_dists = tf.exp(-dists * 5) 129 | weights = exp_dists/tf.reduce_sum(exp_dists,axis=2,keep_dims=True) # (batch_size, npoint, nsample, 1) 130 | new_points *= weights # (batch_size, npoint, nsample, mlp[-1]) 131 | new_points = tf.reduce_sum(new_points, axis=2, keep_dims=True) 132 | elif pooling=='max': 133 | new_points = tf.reduce_max(new_points, axis=[2], keep_dims=True) 134 | elif pooling=='min': 135 | new_points = tf_util.max_pool2d(-1*new_points, [1,nsample], stride=[1,1], padding='VALID', scope='minpool1') 136 | elif pooling=='max_and_avg': 137 | avg_points = tf_util.max_pool2d(new_points, [1,nsample], stride=[1,1], padding='VALID', scope='maxpool1') 138 | max_points = tf_util.avg_pool2d(new_points, [1,nsample], stride=[1,1], padding='VALID', scope='avgpool1') 139 | new_points = tf.concat([avg_points, max_points], axis=-1) 140 | 141 | if mlp2 is None: mlp2 = [] 142 | for i, num_out_channel in enumerate(mlp2): 143 | new_points = tf_util.conv2d(new_points, num_out_channel, [1,1], 144 | padding='VALID', stride=[1,1], 145 | bn=bn, is_training=is_training, 146 | scope='conv_post_%d'%(i), bn_decay=bn_decay) 147 | new_points = tf.squeeze(new_points, [2]) # (batch_size, npoints, mlp2[-1]) 148 | return new_xyz, new_points, idx 149 | 150 | 151 | def acnn_module_rings(xyz, points, normals, npoint, radius_list, nsample_list, mlp_list, is_training, bn_decay, scope, bn=True, use_xyz=True): 152 | ''' A-CNN module with rings 153 | Input: 154 | xyz: (batch_size, ndataset, 3) TF tensor 155 | points: (batch_size, ndataset, channel) TF tensor 156 | normals: (batch_size, ndataset, 3) TF tensor 157 | npoint: int32 -- #points sampled in farthest point sampling 158 | radius_list: list of float32 -- search radiuses (inner and outer) represent ring in local region 159 | nsample_list: list of int32 -- how many points in each local region 160 | mlp: list of list of int32 -- output size for MLP on each point 161 | use_xyz: bool, if True concat XYZ with local point features, otherwise just use point features 162 | Return: 163 | new_xyz: (batch_size, npoint, 3) TF tensor 164 | new_points: (batch_size, npoint, \sum_k{mlp[k][-1]}) TF tensor 165 | ''' 166 | # data_format = 'NCHW' if use_nchw else 'NHWC' 167 | with tf.variable_scope(scope) as sc: 168 | indecies = farthest_point_sample(npoint, xyz) 169 | new_xyz = gather_point(xyz, indecies) # (batch_size, npoint, 3) 170 | new_normals = gather_point(normals, indecies) # (batch_size, npoint, 3) 171 | 172 | batch_size = xyz.get_shape()[0].value 173 | new_points_list = [] 174 | for i in range(len(radius_list)): 175 | radius_in = radius_list[i][0] 176 | radius_out = radius_list[i][1] 177 | nsample = nsample_list[i] 178 | idx, _ = ring_point(radius_in, radius_out, nsample, xyz, new_xyz, indecies) 179 | idx, _, _ = order_neighbors(nsample, idx, xyz, new_xyz, new_normals) 180 | grouped_xyz = group_point(xyz, idx) 181 | grouped_xyz -= tf.tile(tf.expand_dims(new_xyz, 2), [1,1,nsample,1]) 182 | if points is not None: 183 | grouped_points = group_point(points, idx) 184 | if use_xyz: 185 | grouped_points = tf.concat([grouped_points, grouped_xyz], axis=-1) 186 | else: 187 | grouped_points = grouped_xyz 188 | for j,num_out_channel in enumerate(mlp_list[i]): 189 | grouped_points = tf.concat([grouped_points, grouped_points[:,:,:2,:]], axis=2) 190 | grouped_points = tf_util.conv2d(grouped_points, num_out_channel, [1,3], 191 | padding='VALID', stride=[1,1], bn=bn, is_training=is_training, 192 | scope='conv%d_%d'%(i,j), bn_decay=bn_decay) 193 | new_points = tf.reduce_max(grouped_points, axis=[2]) 194 | new_points_list.append(new_points) 195 | new_points_concat = tf.concat(new_points_list, axis=-1) 196 | return new_xyz, new_points_concat, new_normals 197 | 198 | 199 | def pointnet_fp_module(xyz1, xyz2, points1, points2, mlp, is_training, bn_decay, scope, bn=True): 200 | ''' PointNet Feature Propogation (FP) Module 201 | Input: 202 | xyz1: (batch_size, ndataset1, 3) TF tensor 203 | xyz2: (batch_size, ndataset2, 3) TF tensor, sparser than xyz1 204 | points1: (batch_size, ndataset1, nchannel1) TF tensor 205 | points2: (batch_size, ndataset2, nchannel2) TF tensor 206 | mlp: list of int32 -- output size for MLP on each point 207 | Return: 208 | new_points: (batch_size, ndataset1, mlp[-1]) TF tensor 209 | ''' 210 | with tf.variable_scope(scope) as sc: 211 | dist, idx = three_nn(xyz1, xyz2) 212 | dist = tf.maximum(dist, 1e-10) 213 | norm = tf.reduce_sum((1.0/dist),axis=2,keep_dims=True) 214 | norm = tf.tile(norm,[1,1,3]) 215 | weight = (1.0/dist) / norm 216 | interpolated_points = three_interpolate(points2, idx, weight) 217 | 218 | if points1 is not None: 219 | new_points1 = tf.concat(axis=2, values=[interpolated_points, points1]) # B,ndataset1,nchannel1+nchannel2 220 | else: 221 | new_points1 = interpolated_points 222 | new_points1 = tf.expand_dims(new_points1, 2) 223 | for i, num_out_channel in enumerate(mlp): 224 | new_points1 = tf_util.conv2d(new_points1, num_out_channel, [1,1], 225 | padding='VALID', stride=[1,1], 226 | bn=bn, is_training=is_training, 227 | scope='conv_%d'%(i), bn_decay=bn_decay) 228 | new_points1 = tf.squeeze(new_points1, [2]) # B,ndataset1,mlp[-1] 229 | return new_points1 230 | -------------------------------------------------------------------------------- /utils/tf_util.py: -------------------------------------------------------------------------------- 1 | """ Wrapper functions for TensorFlow layers. 2 | 3 | Author: Charles R. Qi 4 | Date: November 2017 5 | """ 6 | 7 | import numpy as np 8 | import tensorflow as tf 9 | 10 | def _variable_on_cpu(name, shape, initializer, use_fp16=False): 11 | """Helper to create a Variable stored on CPU memory. 12 | Args: 13 | name: name of the variable 14 | shape: list of ints 15 | initializer: initializer for Variable 16 | Returns: 17 | Variable Tensor 18 | """ 19 | dtype = tf.float16 if use_fp16 else tf.float32 20 | var = tf.get_variable(name, shape, initializer=initializer, dtype=dtype) 21 | return var 22 | 23 | def _variable_with_weight_decay(name, shape, stddev, wd, use_xavier=True): 24 | """Helper to create an initialized Variable with weight decay. 25 | 26 | Note that the Variable is initialized with a truncated normal distribution. 27 | A weight decay is added only if one is specified. 28 | 29 | Args: 30 | name: name of the variable 31 | shape: list of ints 32 | stddev: standard deviation of a truncated Gaussian 33 | wd: add L2Loss weight decay multiplied by this float. If None, weight 34 | decay is not added for this Variable. 35 | use_xavier: bool, whether to use xavier initializer 36 | 37 | Returns: 38 | Variable Tensor 39 | """ 40 | if use_xavier: 41 | initializer = tf.contrib.layers.xavier_initializer() 42 | else: 43 | initializer = tf.truncated_normal_initializer(stddev=stddev) 44 | var = _variable_on_cpu(name, shape, initializer) 45 | if wd is not None: 46 | weight_decay = tf.multiply(tf.nn.l2_loss(var), wd, name='weight_loss') 47 | tf.add_to_collection('losses', weight_decay) 48 | return var 49 | 50 | 51 | def conv1d(inputs, 52 | num_output_channels, 53 | kernel_size, 54 | scope, 55 | stride=1, 56 | padding='SAME', 57 | use_xavier=True, 58 | stddev=1e-3, 59 | weight_decay=0.0, 60 | activation_fn=tf.nn.relu, 61 | bn=False, 62 | bn_decay=None, 63 | is_training=None): 64 | """ 1D convolution with non-linear operation. 65 | 66 | Args: 67 | inputs: 3-D tensor variable BxLxC 68 | num_output_channels: int 69 | kernel_size: int 70 | scope: string 71 | stride: int 72 | padding: 'SAME' or 'VALID' 73 | use_xavier: bool, use xavier_initializer if true 74 | stddev: float, stddev for truncated_normal init 75 | weight_decay: float 76 | activation_fn: function 77 | bn: bool, whether to use batch norm 78 | bn_decay: float or float tensor variable in [0,1] 79 | is_training: bool Tensor variable 80 | 81 | Returns: 82 | Variable tensor 83 | """ 84 | with tf.variable_scope(scope) as sc: 85 | num_in_channels = inputs.get_shape()[-1].value 86 | kernel_shape = [kernel_size, 87 | num_in_channels, num_output_channels] 88 | kernel = _variable_with_weight_decay('weights', 89 | shape=kernel_shape, 90 | use_xavier=use_xavier, 91 | stddev=stddev, 92 | wd=weight_decay) 93 | outputs = tf.nn.conv1d(inputs, kernel, 94 | stride=stride, 95 | padding=padding) 96 | biases = _variable_on_cpu('biases', [num_output_channels], 97 | tf.constant_initializer(0.0)) 98 | outputs = tf.nn.bias_add(outputs, biases) 99 | 100 | if bn: 101 | outputs = batch_norm_for_conv1d(outputs, is_training, 102 | bn_decay=bn_decay, scope='bn') 103 | 104 | if activation_fn is not None: 105 | outputs = activation_fn(outputs) 106 | return outputs 107 | 108 | 109 | 110 | 111 | def conv2d(inputs, 112 | num_output_channels, 113 | kernel_size, 114 | scope, 115 | stride=[1, 1], 116 | padding='SAME', 117 | use_xavier=True, 118 | stddev=1e-3, 119 | weight_decay=0.0, 120 | activation_fn=tf.nn.relu, 121 | bn=False, 122 | bn_decay=None, 123 | is_training=None): 124 | """ 2D convolution with non-linear operation. 125 | 126 | Args: 127 | inputs: 4-D tensor variable BxHxWxC 128 | num_output_channels: int 129 | kernel_size: a list of 2 ints 130 | scope: string 131 | stride: a list of 2 ints 132 | padding: 'SAME' or 'VALID' 133 | use_xavier: bool, use xavier_initializer if true 134 | stddev: float, stddev for truncated_normal init 135 | weight_decay: float 136 | activation_fn: function 137 | bn: bool, whether to use batch norm 138 | bn_decay: float or float tensor variable in [0,1] 139 | is_training: bool Tensor variable 140 | 141 | Returns: 142 | Variable tensor 143 | """ 144 | with tf.variable_scope(scope) as sc: 145 | kernel_h, kernel_w = kernel_size 146 | num_in_channels = inputs.get_shape()[-1].value 147 | kernel_shape = [kernel_h, kernel_w, 148 | num_in_channels, num_output_channels] 149 | kernel = _variable_with_weight_decay('weights', 150 | shape=kernel_shape, 151 | use_xavier=use_xavier, 152 | stddev=stddev, 153 | wd=weight_decay) 154 | stride_h, stride_w = stride 155 | outputs = tf.nn.conv2d(inputs, kernel, 156 | [1, stride_h, stride_w, 1], 157 | padding=padding) 158 | biases = _variable_on_cpu('biases', [num_output_channels], 159 | tf.constant_initializer(0.0)) 160 | outputs = tf.nn.bias_add(outputs, biases) 161 | 162 | if bn: 163 | outputs = batch_norm_for_conv2d(outputs, is_training, 164 | bn_decay=bn_decay, scope='bn') 165 | 166 | if activation_fn is not None: 167 | outputs = activation_fn(outputs) 168 | return outputs 169 | 170 | 171 | def conv2d_transpose(inputs, 172 | num_output_channels, 173 | kernel_size, 174 | scope, 175 | stride=[1, 1], 176 | padding='SAME', 177 | use_xavier=True, 178 | stddev=1e-3, 179 | weight_decay=0.0, 180 | activation_fn=tf.nn.relu, 181 | bn=False, 182 | bn_decay=None, 183 | is_training=None): 184 | """ 2D convolution transpose with non-linear operation. 185 | 186 | Args: 187 | inputs: 4-D tensor variable BxHxWxC 188 | num_output_channels: int 189 | kernel_size: a list of 2 ints 190 | scope: string 191 | stride: a list of 2 ints 192 | padding: 'SAME' or 'VALID' 193 | use_xavier: bool, use xavier_initializer if true 194 | stddev: float, stddev for truncated_normal init 195 | weight_decay: float 196 | activation_fn: function 197 | bn: bool, whether to use batch norm 198 | bn_decay: float or float tensor variable in [0,1] 199 | is_training: bool Tensor variable 200 | 201 | Returns: 202 | Variable tensor 203 | 204 | Note: conv2d(conv2d_transpose(a, num_out, ksize, stride), a.shape[-1], ksize, stride) == a 205 | """ 206 | with tf.variable_scope(scope) as sc: 207 | kernel_h, kernel_w = kernel_size 208 | num_in_channels = inputs.get_shape()[-1].value 209 | kernel_shape = [kernel_h, kernel_w, 210 | num_output_channels, num_in_channels] # reversed to conv2d 211 | kernel = _variable_with_weight_decay('weights', 212 | shape=kernel_shape, 213 | use_xavier=use_xavier, 214 | stddev=stddev, 215 | wd=weight_decay) 216 | stride_h, stride_w = stride 217 | 218 | # from slim.convolution2d_transpose 219 | def get_deconv_dim(dim_size, stride_size, kernel_size, padding): 220 | dim_size *= stride_size 221 | 222 | if padding == 'VALID' and dim_size is not None: 223 | dim_size += max(kernel_size - stride_size, 0) 224 | return dim_size 225 | 226 | # caculate output shape 227 | batch_size = inputs.get_shape()[0].value 228 | height = inputs.get_shape()[1].value 229 | width = inputs.get_shape()[2].value 230 | out_height = get_deconv_dim(height, stride_h, kernel_h, padding) 231 | out_width = get_deconv_dim(width, stride_w, kernel_w, padding) 232 | output_shape = [batch_size, out_height, out_width, num_output_channels] 233 | 234 | outputs = tf.nn.conv2d_transpose(inputs, kernel, output_shape, 235 | [1, stride_h, stride_w, 1], 236 | padding=padding) 237 | biases = _variable_on_cpu('biases', [num_output_channels], 238 | tf.constant_initializer(0.0)) 239 | outputs = tf.nn.bias_add(outputs, biases) 240 | 241 | if bn: 242 | outputs = batch_norm_for_conv2d(outputs, is_training, 243 | bn_decay=bn_decay, scope='bn') 244 | 245 | if activation_fn is not None: 246 | outputs = activation_fn(outputs) 247 | return outputs 248 | 249 | 250 | 251 | def conv3d(inputs, 252 | num_output_channels, 253 | kernel_size, 254 | scope, 255 | stride=[1, 1, 1], 256 | padding='SAME', 257 | use_xavier=True, 258 | stddev=1e-3, 259 | weight_decay=0.0, 260 | activation_fn=tf.nn.relu, 261 | bn=False, 262 | bn_decay=None, 263 | is_training=None): 264 | """ 3D convolution with non-linear operation. 265 | 266 | Args: 267 | inputs: 5-D tensor variable BxDxHxWxC 268 | num_output_channels: int 269 | kernel_size: a list of 3 ints 270 | scope: string 271 | stride: a list of 3 ints 272 | padding: 'SAME' or 'VALID' 273 | use_xavier: bool, use xavier_initializer if true 274 | stddev: float, stddev for truncated_normal init 275 | weight_decay: float 276 | activation_fn: function 277 | bn: bool, whether to use batch norm 278 | bn_decay: float or float tensor variable in [0,1] 279 | is_training: bool Tensor variable 280 | 281 | Returns: 282 | Variable tensor 283 | """ 284 | with tf.variable_scope(scope) as sc: 285 | kernel_d, kernel_h, kernel_w = kernel_size 286 | num_in_channels = inputs.get_shape()[-1].value 287 | kernel_shape = [kernel_d, kernel_h, kernel_w, 288 | num_in_channels, num_output_channels] 289 | kernel = _variable_with_weight_decay('weights', 290 | shape=kernel_shape, 291 | use_xavier=use_xavier, 292 | stddev=stddev, 293 | wd=weight_decay) 294 | stride_d, stride_h, stride_w = stride 295 | outputs = tf.nn.conv3d(inputs, kernel, 296 | [1, stride_d, stride_h, stride_w, 1], 297 | padding=padding) 298 | biases = _variable_on_cpu('biases', [num_output_channels], 299 | tf.constant_initializer(0.0)) 300 | outputs = tf.nn.bias_add(outputs, biases) 301 | 302 | if bn: 303 | outputs = batch_norm_for_conv3d(outputs, is_training, 304 | bn_decay=bn_decay, scope='bn') 305 | 306 | if activation_fn is not None: 307 | outputs = activation_fn(outputs) 308 | return outputs 309 | 310 | def fully_connected(inputs, 311 | num_outputs, 312 | scope, 313 | use_xavier=True, 314 | stddev=1e-3, 315 | weight_decay=0.0, 316 | activation_fn=tf.nn.relu, 317 | bn=False, 318 | bn_decay=None, 319 | is_training=None): 320 | """ Fully connected layer with non-linear operation. 321 | 322 | Args: 323 | inputs: 2-D tensor BxN 324 | num_outputs: int 325 | 326 | Returns: 327 | Variable tensor of size B x num_outputs. 328 | """ 329 | with tf.variable_scope(scope) as sc: 330 | num_input_units = inputs.get_shape()[-1].value 331 | weights = _variable_with_weight_decay('weights', 332 | shape=[num_input_units, num_outputs], 333 | use_xavier=use_xavier, 334 | stddev=stddev, 335 | wd=weight_decay) 336 | outputs = tf.matmul(inputs, weights) 337 | biases = _variable_on_cpu('biases', [num_outputs], 338 | tf.constant_initializer(0.0)) 339 | outputs = tf.nn.bias_add(outputs, biases) 340 | 341 | if bn: 342 | outputs = batch_norm_for_fc(outputs, is_training, bn_decay, 'bn') 343 | 344 | if activation_fn is not None: 345 | outputs = activation_fn(outputs) 346 | return outputs 347 | 348 | 349 | def max_pool2d(inputs, 350 | kernel_size, 351 | scope, 352 | stride=[2, 2], 353 | padding='VALID'): 354 | """ 2D max pooling. 355 | 356 | Args: 357 | inputs: 4-D tensor BxHxWxC 358 | kernel_size: a list of 2 ints 359 | stride: a list of 2 ints 360 | 361 | Returns: 362 | Variable tensor 363 | """ 364 | with tf.variable_scope(scope) as sc: 365 | kernel_h, kernel_w = kernel_size 366 | stride_h, stride_w = stride 367 | outputs = tf.nn.max_pool(inputs, 368 | ksize=[1, kernel_h, kernel_w, 1], 369 | strides=[1, stride_h, stride_w, 1], 370 | padding=padding, 371 | name=sc.name) 372 | return outputs 373 | 374 | def avg_pool2d(inputs, 375 | kernel_size, 376 | scope, 377 | stride=[2, 2], 378 | padding='VALID'): 379 | """ 2D avg pooling. 380 | 381 | Args: 382 | inputs: 4-D tensor BxHxWxC 383 | kernel_size: a list of 2 ints 384 | stride: a list of 2 ints 385 | 386 | Returns: 387 | Variable tensor 388 | """ 389 | with tf.variable_scope(scope) as sc: 390 | kernel_h, kernel_w = kernel_size 391 | stride_h, stride_w = stride 392 | outputs = tf.nn.avg_pool(inputs, 393 | ksize=[1, kernel_h, kernel_w, 1], 394 | strides=[1, stride_h, stride_w, 1], 395 | padding=padding, 396 | name=sc.name) 397 | return outputs 398 | 399 | 400 | def max_pool3d(inputs, 401 | kernel_size, 402 | scope, 403 | stride=[2, 2, 2], 404 | padding='VALID'): 405 | """ 3D max pooling. 406 | 407 | Args: 408 | inputs: 5-D tensor BxDxHxWxC 409 | kernel_size: a list of 3 ints 410 | stride: a list of 3 ints 411 | 412 | Returns: 413 | Variable tensor 414 | """ 415 | with tf.variable_scope(scope) as sc: 416 | kernel_d, kernel_h, kernel_w = kernel_size 417 | stride_d, stride_h, stride_w = stride 418 | outputs = tf.nn.max_pool3d(inputs, 419 | ksize=[1, kernel_d, kernel_h, kernel_w, 1], 420 | strides=[1, stride_d, stride_h, stride_w, 1], 421 | padding=padding, 422 | name=sc.name) 423 | return outputs 424 | 425 | def avg_pool3d(inputs, 426 | kernel_size, 427 | scope, 428 | stride=[2, 2, 2], 429 | padding='VALID'): 430 | """ 3D avg pooling. 431 | 432 | Args: 433 | inputs: 5-D tensor BxDxHxWxC 434 | kernel_size: a list of 3 ints 435 | stride: a list of 3 ints 436 | 437 | Returns: 438 | Variable tensor 439 | """ 440 | with tf.variable_scope(scope) as sc: 441 | kernel_d, kernel_h, kernel_w = kernel_size 442 | stride_d, stride_h, stride_w = stride 443 | outputs = tf.nn.avg_pool3d(inputs, 444 | ksize=[1, kernel_d, kernel_h, kernel_w, 1], 445 | strides=[1, stride_d, stride_h, stride_w, 1], 446 | padding=padding, 447 | name=sc.name) 448 | return outputs 449 | 450 | 451 | 452 | 453 | 454 | def batch_norm_template(inputs, is_training, scope, moments_dims, bn_decay): 455 | """ Batch normalization on convolutional maps and beyond... 456 | Ref.: http://stackoverflow.com/questions/33949786/how-could-i-use-batch-normalization-in-tensorflow 457 | 458 | Args: 459 | inputs: Tensor, k-D input ... x C could be BC or BHWC or BDHWC 460 | is_training: boolean tf.Varialbe, true indicates training phase 461 | scope: string, variable scope 462 | moments_dims: a list of ints, indicating dimensions for moments calculation 463 | bn_decay: float or float tensor variable, controling moving average weight 464 | Return: 465 | normed: batch-normalized maps 466 | """ 467 | # For support of GAN 468 | #bn_decay = bn_decay if bn_decay is not None else 0.9 469 | #return tf.contrib.layers.batch_norm(inputs, 470 | # center=True, scale=True, 471 | # is_training=is_training, decay=bn_decay,updates_collections=None, 472 | # scope=scope) 473 | with tf.variable_scope(scope) as sc: 474 | num_channels = inputs.get_shape()[-1].value 475 | beta = tf.Variable(tf.constant(0.0, shape=[num_channels]), 476 | name='beta', trainable=True) 477 | gamma = tf.Variable(tf.constant(1.0, shape=[num_channels]), 478 | name='gamma', trainable=True) 479 | batch_mean, batch_var = tf.nn.moments(inputs, moments_dims, name='moments') 480 | decay = bn_decay if bn_decay is not None else 0.9 481 | ema = tf.train.ExponentialMovingAverage(decay=decay) 482 | # Operator that maintains moving averages of variables. 483 | # Need to set reuse=False, otherwise if reuse, will see moments_1/mean/ExponentialMovingAverage/ does not exist 484 | # https://github.com/shekkizh/WassersteinGAN.tensorflow/issues/3 485 | with tf.variable_scope(tf.get_variable_scope(), reuse=False): 486 | ema_apply_op = tf.cond(is_training, 487 | lambda: ema.apply([batch_mean, batch_var]), 488 | lambda: tf.no_op()) 489 | 490 | # Update moving average and return current batch's avg and var. 491 | def mean_var_with_update(): 492 | with tf.control_dependencies([ema_apply_op]): 493 | return tf.identity(batch_mean), tf.identity(batch_var) 494 | 495 | # ema.average returns the Variable holding the average of var. 496 | mean, var = tf.cond(is_training, 497 | mean_var_with_update, 498 | lambda: (ema.average(batch_mean), ema.average(batch_var))) 499 | normed = tf.nn.batch_normalization(inputs, mean, var, beta, gamma, 1e-3) 500 | return normed 501 | 502 | def batch_norm_for_fc(inputs, is_training, bn_decay, scope): 503 | """ Batch normalization on FC data. 504 | 505 | Args: 506 | inputs: Tensor, 2D BxC input 507 | is_training: boolean tf.Varialbe, true indicates training phase 508 | bn_decay: float or float tensor variable, controling moving average weight 509 | scope: string, variable scope 510 | Return: 511 | normed: batch-normalized maps 512 | """ 513 | return batch_norm_template(inputs, is_training, scope, [0,], bn_decay) 514 | 515 | 516 | def batch_norm_for_conv1d(inputs, is_training, bn_decay, scope): 517 | """ Batch normalization on 1D convolutional maps. 518 | 519 | Args: 520 | inputs: Tensor, 3D BLC input maps 521 | is_training: boolean tf.Varialbe, true indicates training phase 522 | bn_decay: float or float tensor variable, controling moving average weight 523 | scope: string, variable scope 524 | Return: 525 | normed: batch-normalized maps 526 | """ 527 | return batch_norm_template(inputs, is_training, scope, [0,1], bn_decay) 528 | 529 | 530 | 531 | 532 | def batch_norm_for_conv2d(inputs, is_training, bn_decay, scope): 533 | """ Batch normalization on 2D convolutional maps. 534 | 535 | Args: 536 | inputs: Tensor, 4D BHWC input maps 537 | is_training: boolean tf.Varialbe, true indicates training phase 538 | bn_decay: float or float tensor variable, controling moving average weight 539 | scope: string, variable scope 540 | Return: 541 | normed: batch-normalized maps 542 | """ 543 | return batch_norm_template(inputs, is_training, scope, [0,1,2], bn_decay) 544 | 545 | 546 | 547 | def batch_norm_for_conv3d(inputs, is_training, bn_decay, scope): 548 | """ Batch normalization on 3D convolutional maps. 549 | 550 | Args: 551 | inputs: Tensor, 5D BDHWC input maps 552 | is_training: boolean tf.Varialbe, true indicates training phase 553 | bn_decay: float or float tensor variable, controling moving average weight 554 | scope: string, variable scope 555 | Return: 556 | normed: batch-normalized maps 557 | """ 558 | return batch_norm_template(inputs, is_training, scope, [0,1,2,3], bn_decay) 559 | 560 | 561 | def dropout(inputs, 562 | is_training, 563 | scope, 564 | keep_prob=0.5, 565 | noise_shape=None): 566 | """ Dropout layer. 567 | 568 | Args: 569 | inputs: tensor 570 | is_training: boolean tf.Variable 571 | scope: string 572 | keep_prob: float in [0,1] 573 | noise_shape: list of ints 574 | 575 | Returns: 576 | tensor variable 577 | """ 578 | with tf.variable_scope(scope) as sc: 579 | outputs = tf.cond(is_training, 580 | lambda: tf.nn.dropout(inputs, keep_prob, noise_shape), 581 | lambda: inputs) 582 | return outputs 583 | --------------------------------------------------------------------------------