├── utils ├── pc_distance │ ├── __init__.py │ ├── tf_approxmatch.cu.o │ ├── tf_nndistance.cu.o │ ├── tf_approxmatch_so.so │ ├── makefile │ ├── tf_nndistance.py │ ├── tf_approxmatch.py │ ├── tf_nndistance.cu │ ├── tf_approxmatch.cu │ ├── tf_nndistance.cpp │ └── tf_approxmatch.cpp ├── tf_util_loss.py ├── data_prep_util.py ├── pc_util.py ├── eulerangles.py ├── tf_util.py └── plyfile.py ├── results ├── result1.png ├── result2.png ├── result3.png ├── result4.png └── network_structure.png ├── show_pc.py ├── README.md ├── models ├── pointnet_pose_test.py └── pointnet_pose.py ├── provider.py ├── pointnet_autoencoder_train.py └── helper.py /utils/pc_distance/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /results/result1.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/vinits5/pc_autoencoder/HEAD/results/result1.png -------------------------------------------------------------------------------- /results/result2.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/vinits5/pc_autoencoder/HEAD/results/result2.png -------------------------------------------------------------------------------- /results/result3.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/vinits5/pc_autoencoder/HEAD/results/result3.png -------------------------------------------------------------------------------- /results/result4.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/vinits5/pc_autoencoder/HEAD/results/result4.png -------------------------------------------------------------------------------- /results/network_structure.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/vinits5/pc_autoencoder/HEAD/results/network_structure.png -------------------------------------------------------------------------------- /utils/pc_distance/tf_approxmatch.cu.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/vinits5/pc_autoencoder/HEAD/utils/pc_distance/tf_approxmatch.cu.o -------------------------------------------------------------------------------- /utils/pc_distance/tf_nndistance.cu.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/vinits5/pc_autoencoder/HEAD/utils/pc_distance/tf_nndistance.cu.o -------------------------------------------------------------------------------- /utils/pc_distance/tf_approxmatch_so.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/vinits5/pc_autoencoder/HEAD/utils/pc_distance/tf_approxmatch_so.so -------------------------------------------------------------------------------- /show_pc.py: -------------------------------------------------------------------------------- 1 | import provider 2 | import helper 3 | 4 | TRAIN_FILES = provider.getDataFiles('data/modelnet40_ply_hdf5_2048/train_files.txt') 5 | 6 | current_data, current_label = provider.loadDataFile(TRAIN_FILES[0]) 7 | 8 | idx = int(sys.argv[1]) 9 | 10 | helper.display_clouds_data(current_data[idx]) -------------------------------------------------------------------------------- /utils/pc_distance/makefile: -------------------------------------------------------------------------------- 1 | cuda_inc = /usr/local/cuda-8.0/include/ 2 | cuda_lib = /usr/local/cuda-8.0/lib64/ 3 | nsync = /usr/local/lib/python2.7/dist-packages/external/nsync/public 4 | nvcc = /usr/local/cuda-8.0/bin/nvcc 5 | tf_inc = /usr/local/lib/python2.7/dist-packages/tensorflow/include 6 | tf_lib = /usr/local/lib/python2.7/dist-packages/tensorflow 7 | 8 | all: tf_nndistance_so.so tf_approxmatch_so.so 9 | 10 | tf_nndistance.cu.o: tf_nndistance.cu 11 | $(nvcc) tf_nndistance.cu -o tf_nndistance.cu.o -c -O2 -DGOOGLE_CUDA=1 -x cu -Xcompiler -fPIC 12 | 13 | tf_nndistance_so.so: tf_nndistance.cpp tf_nndistance.cu.o 14 | g++ tf_nndistance.cpp tf_nndistance.cu.o -o tf_nndistance_so.so \ 15 | -I $(cuda_inc) -I $(tf_inc) -I $(nsync) \ 16 | -L $(cuda_lib) -lcudart -L $(tf_lib) -ltensorflow_framework \ 17 | -shared -D_GLIBCXX_USE_CXX11_ABI=0 -std=c++11 -fPIC -O2 18 | 19 | tf_approxmatch.cu.o: tf_approxmatch.cu 20 | $(nvcc) tf_approxmatch.cu -o tf_approxmatch.cu.o -c -O2 -DGOOGLE_CUDA=1 -x cu -Xcompiler -fPIC 21 | 22 | tf_approxmatch_so.so: tf_approxmatch.cpp tf_approxmatch.cu.o 23 | g++ -shared $(CPPFLAGS) tf_approxmatch.cpp tf_approxmatch.cu.o -o tf_approxmatch_so.so \ 24 | -I $(cuda_inc) -I $(tf_inc) -I $(nsync) \ 25 | -L $(cuda_lib) -lcudart -L $(tf_lib) -ltensorflow_framework \ 26 | -shared -D_GLIBCXX_USE_CXX11_ABI=0 -std=c++11 -fPIC -O2 27 | 28 | clean: 29 | rm -rf *.o *.so 30 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Auto Encoder for 3D Point Clouds 2 | 3 | ### Network Structure: 4 |

5 | 6 |

7 | 8 | ### Code: 9 | Steps to train the auto-encoder: 10 | 1. Download ModelNet40 Dataset [[Link]](https://shapenet.cs.stanford.edu/media/modelnet40_ply_hdf5_2048.zip) 11 | 2. Clone repository. 12 | 3. Extract the zip file and copy *modelnet40_ply_hdf5_2048* folder to *pc_autoencoder/data*. 13 | 4. *python pointnet_autoencoder_train.py --mode train* 14 | 15 | Steps to test the auto-encoder: 16 | 1. Download dataset as given in training steps. 17 | 2. Download weights for the trained network. [[Link]](https://drive.google.com/drive/folders/17k0mWR65eHQbnWcvNWJKlZ1hXeVYVhm7?usp=sharing) 18 | 3. *python pointnet_autoencoder_train.py --mode test* 19 | 20 | Visualise the Dataset: 21 | *python show_pc.py idx* 22 | idx: Index of Point Cloud in ModelNet40 Dataset. 23 | 24 | ### Results: 25 | **Red colored point clouds are input to the network and blue point clouds are the output.** 26 | 27 | [Note: A translation has been applied to blue point clouds during testing for a better visualisation purpose.] 28 |

29 | 30 |

31 | 32 | #### Additional Results: 33 | 34 |

35 | 36 | 37 | 38 |

39 | 40 | ### References: 41 | 1. PointNet: Deep Learning on Point Sets for 3D Classification and Segmentation [[Link]](https://arxiv.org/abs/1612.00593) 42 | 2. PCN: Point Completion Network [[Link]](https://arxiv.org/abs/1808.00671) -------------------------------------------------------------------------------- /utils/tf_util_loss.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | from pc_distance import tf_nndistance, tf_approxmatch 3 | 4 | 5 | def mlp(features, layer_dims, bn=None, bn_params=None): 6 | for i, num_outputs in enumerate(layer_dims[:-1]): 7 | features = tf.contrib.layers.fully_connected( 8 | features, num_outputs, 9 | normalizer_fn=bn, 10 | normalizer_params=bn_params, 11 | scope='fc_%d' % i) 12 | outputs = tf.contrib.layers.fully_connected( 13 | features, layer_dims[-1], 14 | activation_fn=None, 15 | scope='fc_%d' % (len(layer_dims) - 1)) 16 | return outputs 17 | 18 | 19 | def mlp_conv(inputs, layer_dims, bn=None, bn_params=None): 20 | for i, num_out_channel in enumerate(layer_dims[:-1]): 21 | inputs = tf.contrib.layers.conv2d( 22 | inputs, num_out_channel, 23 | kernel_size=1, 24 | normalizer_fn=bn, 25 | normalizer_params=bn_params, 26 | scope='conv_%d' % i) 27 | outputs = tf.contrib.layers.conv2d( 28 | inputs, layer_dims[-1], 29 | kernel_size=1, 30 | activation_fn=None, 31 | scope='conv_%d' % (len(layer_dims) - 1)) 32 | return outputs 33 | 34 | 35 | def chamfer(pcd1, pcd2): 36 | dist1, _, dist2, _ = tf_nndistance.nn_distance(pcd1, pcd2) 37 | dist1 = tf.reduce_mean(tf.sqrt(dist1)) 38 | dist2 = tf.reduce_mean(tf.sqrt(dist2)) 39 | return (dist1 + dist2) / 2 40 | 41 | 42 | def earth_mover(pcd1, pcd2): 43 | assert pcd1.shape[1] == pcd2.shape[1] 44 | num_points = tf.cast(pcd1.shape[1], tf.float32) 45 | match = tf_approxmatch.approx_match(pcd1, pcd2) 46 | cost = tf_approxmatch.match_cost(pcd1, pcd2, match) 47 | return tf.reduce_mean(cost / num_points) 48 | 49 | 50 | def add_train_summary(name, value): 51 | tf.summary.scalar(name, value, collections=['train_summary']) 52 | 53 | 54 | def add_valid_summary(name, value): 55 | avg, update = tf.metrics.mean(value) 56 | tf.summary.scalar(name, avg, collections=['valid_summary']) 57 | return update 58 | -------------------------------------------------------------------------------- /models/pointnet_pose_test.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import numpy as np 3 | import math 4 | import sys 5 | import os 6 | import numpy as np 7 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 8 | sys.path.append(BASE_DIR) 9 | sys.path.append(os.path.join(BASE_DIR, '../utils')) 10 | import tf_util 11 | # from transform_nets import input_transform_net, feature_transform_net 12 | # import tf_util_loss 13 | 14 | class Network: 15 | def placeholder_inputs(self,batch_size, num_point): 16 | # with tf.variable_scope('inputs') as ip: 17 | source_pointclouds_pl = tf.placeholder(tf.float32, shape=(batch_size, num_point, 3)) 18 | return source_pointclouds_pl 19 | 20 | def get_model(self, source_pointclouds_pl, feature_size, is_training, bn_decay=None): 21 | """ Classification PointNet, input is BxNx3, output Bx40 """ 22 | # with tf.variable_scope('PointNet') as pn: 23 | 24 | # Comment above two lines to have same points for loss and features and also change the variable names in the next line. 25 | batch_size = source_pointclouds_pl.get_shape()[0].value 26 | num_point = source_pointclouds_pl.get_shape()[1].value 27 | end_points = {} 28 | 29 | input_image = tf.expand_dims(source_pointclouds_pl, -1) 30 | 31 | net = tf_util.conv2d(input_image, 128, [1,3], 32 | padding='VALID', stride=[1,1], 33 | bn=True, is_training=is_training, 34 | scope='conv1', bn_decay=bn_decay) 35 | 36 | net = tf_util.conv2d(net, 256, [1,1], 37 | padding='VALID', stride=[1,1], 38 | bn=True, is_training=is_training, 39 | scope='conv2', bn_decay=bn_decay, activation_fn=None) 40 | 41 | # Symmetric function: max pooling 42 | source_feature = tf_util.max_pool2d(net, [num_point, 1], 43 | padding='VALID', scope='maxpool') 44 | source_feature = tf.tile(source_feature, [1, num_point, 1, 1]) 45 | source_feature = tf.concat([net, source_feature], axis=3) 46 | 47 | net = tf_util.conv2d(source_feature, 512, [1,1], 48 | padding='VALID', stride=[1,1], 49 | bn=True, is_training=is_training, 50 | scope='conv3', bn_decay=bn_decay) 51 | 52 | net = tf_util.conv2d(net, 1024, [1,1], 53 | padding='VALID', stride=[1,1], 54 | bn=True, is_training=is_training, 55 | scope='conv4', bn_decay=bn_decay, activation_fn=None) 56 | source_global_feature = tf_util.max_pool2d(net, [num_point, 1], 57 | padding='VALID', scope='maxpool') 58 | source_global_feature = tf.reshape(source_global_feature, [batch_size, -1]) 59 | 60 | return source_global_feature 61 | 62 | def decode_data(self, source_global_feature, is_training, bn_decay=None): 63 | batch_size = source_global_feature.get_shape()[0].value 64 | net = tf_util.fully_connected(source_global_feature, 1024, bn=True, is_training=is_training, scope='fc1', bn_decay=bn_decay) 65 | net = tf_util.fully_connected(net, 1024, bn=True, is_training=is_training, scope='fc2', bn_decay=bn_decay) 66 | net = tf_util.fully_connected(net, 1024*3, activation_fn=None, scope='fc3') 67 | predicted_pointclouds_pl = tf.reshape(net, [batch_size, 1024, 3]) 68 | return predicted_pointclouds_pl 69 | 70 | def get_loss_b(self, predicted_pointclouds_pl, source_pointclouds_pl): 71 | loss = 0 72 | return loss 73 | 74 | if __name__=='__main__': 75 | with tf.Graph().as_default(): 76 | net = Network() 77 | inputs = tf.zeros((1,1024,3)) 78 | outputs = net.get_model(inputs, 1024, tf.constant(True)) 79 | pt_cloud = net.decode_data(outputs, tf.constant(True)) 80 | loss = net.get_loss_b(pt_cloud, inputs) -------------------------------------------------------------------------------- /utils/pc_distance/tf_nndistance.py: -------------------------------------------------------------------------------- 1 | import os, sys 2 | import tensorflow as tf 3 | from tensorflow.python.framework import ops 4 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 5 | nn_distance_module=tf.load_op_library(os.path.join(BASE_DIR, 'tf_nndistance_so.so')) 6 | 7 | def nn_distance(xyz1,xyz2): 8 | ''' 9 | Computes the distance of nearest neighbors for a pair of point clouds 10 | input: xyz1: (batch_size,#points_1,3) the first point cloud 11 | input: xyz2: (batch_size,#points_2,3) the second point cloud 12 | output: dist1: (batch_size,#point_1) distance from first to second 13 | output: idx1: (batch_size,#point_1) nearest neighbor from first to second 14 | output: dist2: (batch_size,#point_2) distance from second to first 15 | output: idx2: (batch_size,#point_2) nearest neighbor from second to first 16 | ''' 17 | return nn_distance_module.nn_distance(xyz1,xyz2) 18 | #@tf.RegisterShape('NnDistance') 19 | #def _nn_distance_shape(op): 20 | #shape1=op.inputs[0].get_shape().with_rank(3) 21 | #shape2=op.inputs[1].get_shape().with_rank(3) 22 | #return [tf.TensorShape([shape1.dims[0],shape1.dims[1]]),tf.TensorShape([shape1.dims[0],shape1.dims[1]]), 23 | #tf.TensorShape([shape2.dims[0],shape2.dims[1]]),tf.TensorShape([shape2.dims[0],shape2.dims[1]])] 24 | @ops.RegisterGradient('NnDistance') 25 | def _nn_distance_grad(op,grad_dist1,grad_idx1,grad_dist2,grad_idx2): 26 | xyz1=op.inputs[0] 27 | xyz2=op.inputs[1] 28 | idx1=op.outputs[1] 29 | idx2=op.outputs[3] 30 | return nn_distance_module.nn_distance_grad(xyz1,xyz2,grad_dist1,idx1,grad_dist2,idx2) 31 | 32 | 33 | if __name__=='__main__': 34 | import numpy as np 35 | import random 36 | import time 37 | from tensorflow.python.ops.gradient_checker import compute_gradient 38 | random.seed(100) 39 | np.random.seed(100) 40 | with tf.Session('') as sess: 41 | xyz1=np.random.randn(32,16384,3).astype('float32') 42 | xyz2=np.random.randn(32,1024,3).astype('float32') 43 | #with tf.device('/gpu:0'): 44 | if True: 45 | inp1=tf.Variable(xyz1) 46 | inp2=tf.constant(xyz2) 47 | reta,retb,retc,retd=nn_distance(inp1,inp2) 48 | loss=tf.reduce_sum(reta)+tf.reduce_sum(retc) 49 | train=tf.train.GradientDescentOptimizer(learning_rate=0.05).minimize(loss) 50 | sess.run(tf.initialize_all_variables()) 51 | t0=time.time() 52 | t1=t0 53 | best=1e100 54 | for i in xrange(100): 55 | trainloss,_=sess.run([loss,train]) 56 | newt=time.time() 57 | best=min(best,newt-t1) 58 | print(i,trainloss,(newt-t0)/(i+1),best) 59 | t1=newt 60 | #print sess.run([inp1,retb,inp2,retd]) 61 | #grads=compute_gradient([inp1,inp2],[(16,32,3),(16,32,3)],loss,(1,),[xyz1,xyz2]) 62 | #for i,j in grads: 63 | #print i.shape,j.shape,np.mean(np.abs(i-j)),np.mean(np.abs(i)),np.mean(np.abs(j)) 64 | #for i in xrange(10): 65 | #t0=time.time() 66 | #a,b,c,d=sess.run([reta,retb,retc,retd],feed_dict={inp1:xyz1,inp2:xyz2}) 67 | #print 'time',time.time()-t0 68 | #print a.shape,b.shape,c.shape,d.shape 69 | #print a.dtype,b.dtype,c.dtype,d.dtype 70 | #samples=np.array(random.sample(range(xyz2.shape[1]),100),dtype='int32') 71 | #dist1=((xyz1[:,samples,None,:]-xyz2[:,None,:,:])**2).sum(axis=-1).min(axis=-1) 72 | #idx1=((xyz1[:,samples,None,:]-xyz2[:,None,:,:])**2).sum(axis=-1).argmin(axis=-1) 73 | #print np.abs(dist1-a[:,samples]).max() 74 | #print np.abs(idx1-b[:,samples]).max() 75 | #dist2=((xyz2[:,samples,None,:]-xyz1[:,None,:,:])**2).sum(axis=-1).min(axis=-1) 76 | #idx2=((xyz2[:,samples,None,:]-xyz1[:,None,:,:])**2).sum(axis=-1).argmin(axis=-1) 77 | #print np.abs(dist2-c[:,samples]).max() 78 | #print np.abs(idx2-d[:,samples]).max() 79 | 80 | -------------------------------------------------------------------------------- /models/pointnet_pose.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import numpy as np 3 | import math 4 | import sys 5 | import os 6 | import numpy as np 7 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 8 | sys.path.append(BASE_DIR) 9 | sys.path.append(os.path.join(BASE_DIR, '../utils')) 10 | import tf_util 11 | # from transform_nets import input_transform_net, feature_transform_net 12 | import tf_util_loss 13 | 14 | class Network: 15 | def placeholder_inputs(self,batch_size, num_point): 16 | # with tf.variable_scope('inputs') as ip: 17 | source_pointclouds_pl = tf.placeholder(tf.float32, shape=(batch_size, num_point, 3)) 18 | return source_pointclouds_pl 19 | 20 | def get_model(self, source_pointclouds_pl, feature_size, is_training, bn_decay=None): 21 | """ Classification PointNet, input is BxNx3, output Bx40 """ 22 | # with tf.variable_scope('PointNet') as pn: 23 | 24 | # Comment above two lines to have same points for loss and features and also change the variable names in the next line. 25 | batch_size = source_pointclouds_pl.get_shape()[0].value 26 | num_point = source_pointclouds_pl.get_shape()[1].value 27 | end_points = {} 28 | 29 | input_image = tf.expand_dims(source_pointclouds_pl, -1) 30 | 31 | net = tf_util.conv2d(input_image, 128, [1,3], 32 | padding='VALID', stride=[1,1], 33 | bn=True, is_training=is_training, 34 | scope='conv1', bn_decay=bn_decay) 35 | 36 | net = tf_util.conv2d(net, 256, [1,1], 37 | padding='VALID', stride=[1,1], 38 | bn=True, is_training=is_training, 39 | scope='conv2', bn_decay=bn_decay, activation_fn=None) 40 | 41 | # Symmetric function: max pooling 42 | source_feature = tf_util.max_pool2d(net, [num_point, 1], 43 | padding='VALID', scope='maxpool') 44 | source_feature = tf.tile(source_feature, [1, num_point, 1, 1]) 45 | source_feature = tf.concat([net, source_feature], axis=3) 46 | 47 | net = tf_util.conv2d(source_feature, 512, [1,1], 48 | padding='VALID', stride=[1,1], 49 | bn=True, is_training=is_training, 50 | scope='conv3', bn_decay=bn_decay) 51 | 52 | net = tf_util.conv2d(net, 1024, [1,1], 53 | padding='VALID', stride=[1,1], 54 | bn=True, is_training=is_training, 55 | scope='conv4', bn_decay=bn_decay, activation_fn=None) 56 | source_global_feature = tf_util.max_pool2d(net, [num_point, 1], 57 | padding='VALID', scope='maxpool') 58 | source_global_feature = tf.reshape(source_global_feature, [batch_size, -1]) 59 | 60 | return source_global_feature 61 | 62 | def decode_data(self, source_global_feature, is_training, bn_decay=None): 63 | batch_size = source_global_feature.get_shape()[0].value 64 | net = tf_util.fully_connected(source_global_feature, 1024, bn=True, is_training=is_training, scope='fc1', bn_decay=bn_decay) 65 | net = tf_util.fully_connected(net, 1024, bn=True, is_training=is_training, scope='fc2', bn_decay=bn_decay) 66 | net = tf_util.fully_connected(net, 1024*3, activation_fn=None, scope='fc3') 67 | predicted_pointclouds_pl = tf.reshape(net, [batch_size, 1024, 3]) 68 | return predicted_pointclouds_pl 69 | 70 | def get_loss_b(self, predicted_pointclouds_pl, source_pointclouds_pl): 71 | with tf.variable_scope('loss') as LossEvaluation: 72 | # loss = tf.reduce_mean(tf.square(tf.subtract(predicted_pointclouds_pl, source_pointclouds_pl))) 73 | loss = tf_util_loss.chamfer(predicted_pointclouds_pl, source_pointclouds_pl) 74 | return loss 75 | 76 | if __name__=='__main__': 77 | with tf.Graph().as_default(): 78 | net = Network() 79 | inputs = tf.zeros((32,1024,3)) 80 | outputs = net.get_model(inputs, 1024, tf.constant(True)) 81 | print(outputs) -------------------------------------------------------------------------------- /utils/pc_distance/tf_approxmatch.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | from tensorflow.python.framework import ops 3 | import os.path as osp 4 | 5 | base_dir = osp.dirname(osp.abspath(__file__)) 6 | 7 | approxmatch_module = tf.load_op_library(osp.join(base_dir, 'tf_approxmatch_so.so')) 8 | 9 | 10 | def approx_match(xyz1,xyz2): 11 | ''' 12 | input: 13 | xyz1 : batch_size * #dataset_points * 3 14 | xyz2 : batch_size * #query_points * 3 15 | returns: 16 | match : batch_size * #query_points * #dataset_points 17 | ''' 18 | return approxmatch_module.approx_match(xyz1,xyz2) 19 | ops.NoGradient('ApproxMatch') 20 | #@tf.RegisterShape('ApproxMatch') 21 | @ops.RegisterShape('ApproxMatch') 22 | def _approx_match_shape(op): 23 | shape1=op.inputs[0].get_shape().with_rank(3) 24 | shape2=op.inputs[1].get_shape().with_rank(3) 25 | return [tf.TensorShape([shape1.dims[0],shape2.dims[1],shape1.dims[1]])] 26 | 27 | def match_cost(xyz1,xyz2,match): 28 | ''' 29 | input: 30 | xyz1 : batch_size * #dataset_points * 3 31 | xyz2 : batch_size * #query_points * 3 32 | match : batch_size * #query_points * #dataset_points 33 | returns: 34 | cost : batch_size 35 | ''' 36 | return approxmatch_module.match_cost(xyz1,xyz2,match) 37 | #@tf.RegisterShape('MatchCost') 38 | @ops.RegisterShape('MatchCost') 39 | def _match_cost_shape(op): 40 | shape1=op.inputs[0].get_shape().with_rank(3) 41 | shape2=op.inputs[1].get_shape().with_rank(3) 42 | shape3=op.inputs[2].get_shape().with_rank(3) 43 | return [tf.TensorShape([shape1.dims[0]])] 44 | @tf.RegisterGradient('MatchCost') 45 | def _match_cost_grad(op,grad_cost): 46 | xyz1=op.inputs[0] 47 | xyz2=op.inputs[1] 48 | match=op.inputs[2] 49 | grad_1,grad_2=approxmatch_module.match_cost_grad(xyz1,xyz2,match) 50 | return [grad_1*tf.expand_dims(tf.expand_dims(grad_cost,1),2),grad_2*tf.expand_dims(tf.expand_dims(grad_cost,1),2),None] 51 | 52 | if __name__=='__main__': 53 | alpha=0.5 54 | beta=2.0 55 | import bestmatch 56 | import numpy as np 57 | import math 58 | import random 59 | import cv2 60 | 61 | import tf_nndistance 62 | 63 | npoint=100 64 | 65 | with tf.device('/gpu:2'): 66 | pt_in=tf.placeholder(tf.float32,shape=(1,npoint*4,3)) 67 | mypoints=tf.Variable(np.random.randn(1,npoint,3).astype('float32')) 68 | match=approx_match(pt_in,mypoints) 69 | loss=tf.reduce_sum(match_cost(pt_in,mypoints,match)) 70 | #match=approx_match(mypoints,pt_in) 71 | #loss=tf.reduce_sum(match_cost(mypoints,pt_in,match)) 72 | #distf,_,distb,_=tf_nndistance.nn_distance(pt_in,mypoints) 73 | #loss=tf.reduce_sum((distf+1e-9)**0.5)*0.5+tf.reduce_sum((distb+1e-9)**0.5)*0.5 74 | #loss=tf.reduce_max((distf+1e-9)**0.5)*0.5*npoint+tf.reduce_max((distb+1e-9)**0.5)*0.5*npoint 75 | 76 | optimizer=tf.train.GradientDescentOptimizer(1e-4).minimize(loss) 77 | with tf.Session('') as sess: 78 | sess.run(tf.initialize_all_variables()) 79 | while True: 80 | meanloss=0 81 | meantrueloss=0 82 | for i in xrange(1001): 83 | #phi=np.random.rand(4*npoint)*math.pi*2 84 | #tpoints=(np.hstack([np.cos(phi)[:,None],np.sin(phi)[:,None],(phi*0)[:,None]])*random.random())[None,:,:] 85 | #tpoints=((np.random.rand(400)-0.5)[:,None]*[0,2,0]+[(random.random()-0.5)*2,0,0]).astype('float32')[None,:,:] 86 | tpoints=np.hstack([np.linspace(-1,1,400)[:,None],(random.random()*2*np.linspace(1,0,400)**2)[:,None],np.zeros((400,1))])[None,:,:] 87 | trainloss,_=sess.run([loss,optimizer],feed_dict={pt_in:tpoints.astype('float32')}) 88 | trainloss,trainmatch=sess.run([loss,match],feed_dict={pt_in:tpoints.astype('float32')}) 89 | #trainmatch=trainmatch.transpose((0,2,1)) 90 | show=np.zeros((400,400,3),dtype='uint8')^255 91 | trainmypoints=sess.run(mypoints) 92 | for i in xrange(len(tpoints[0])): 93 | u=np.random.choice(range(len(trainmypoints[0])),p=trainmatch[0].T[i]) 94 | cv2.line(show, 95 | (int(tpoints[0][i,1]*100+200),int(tpoints[0][i,0]*100+200)), 96 | (int(trainmypoints[0][u,1]*100+200),int(trainmypoints[0][u,0]*100+200)), 97 | cv2.cv.CV_RGB(0,255,0)) 98 | for x,y,z in tpoints[0]: 99 | cv2.circle(show,(int(y*100+200),int(x*100+200)),2,cv2.cv.CV_RGB(255,0,0)) 100 | for x,y,z in trainmypoints[0]: 101 | cv2.circle(show,(int(y*100+200),int(x*100+200)),3,cv2.cv.CV_RGB(0,0,255)) 102 | cost=((tpoints[0][:,None,:]-np.repeat(trainmypoints[0][None,:,:],4,axis=1))**2).sum(axis=2)**0.5 103 | #trueloss=bestmatch.bestmatch(cost)[0] 104 | print(trainloss) #,trueloss 105 | cv2.imshow('show',show) 106 | cmd=cv2.waitKey(10)%256 107 | if cmd==ord('q'): 108 | break 109 | -------------------------------------------------------------------------------- /utils/pc_distance/tf_nndistance.cu: -------------------------------------------------------------------------------- 1 | #if GOOGLE_CUDA 2 | #define EIGEN_USE_GPU 3 | // #include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor" 4 | 5 | __global__ void NmDistanceKernel(int b,int n,const float * xyz,int m,const float * xyz2,float * result,int * result_i){ 6 | const int batch=512; 7 | __shared__ float buf[batch*3]; 8 | for (int i=blockIdx.x;ibest){ 120 | result[(i*n+j)]=best; 121 | result_i[(i*n+j)]=best_i; 122 | } 123 | } 124 | __syncthreads(); 125 | } 126 | } 127 | } 128 | void NmDistanceKernelLauncher(int b,int n,const float * xyz,int m,const float * xyz2,float * result,int * result_i,float * result2,int * result2_i){ 129 | NmDistanceKernel<<>>(b,n,xyz,m,xyz2,result,result_i); 130 | NmDistanceKernel<<>>(b,m,xyz2,n,xyz,result2,result2_i); 131 | } 132 | __global__ void NmDistanceGradKernel(int b,int n,const float * xyz1,int m,const float * xyz2,const float * grad_dist1,const int * idx1,float * grad_xyz1,float * grad_xyz2){ 133 | for (int i=blockIdx.x;i>>(b,n,xyz1,m,xyz2,grad_dist1,idx1,grad_xyz1,grad_xyz2); 156 | NmDistanceGradKernel<<>>(b,m,xyz2,n,xyz1,grad_dist2,idx2,grad_xyz2,grad_xyz1); 157 | } 158 | 159 | #endif 160 | -------------------------------------------------------------------------------- /utils/data_prep_util.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 4 | sys.path.append(BASE_DIR) 5 | from plyfile import (PlyData, PlyElement, make2d, PlyParseError, PlyProperty) 6 | import numpy as np 7 | import h5py 8 | 9 | SAMPLING_BIN = os.path.join(BASE_DIR, 'third_party/mesh_sampling/build/pcsample') 10 | 11 | SAMPLING_POINT_NUM = 2048 12 | SAMPLING_LEAF_SIZE = 0.005 13 | 14 | MODELNET40_PATH = '../datasets/modelnet40' 15 | def export_ply(pc, filename): 16 | vertex = np.zeros(pc.shape[0], dtype=[('x', 'f4'), ('y', 'f4'), ('z', 'f4')]) 17 | for i in range(pc.shape[0]): 18 | vertex[i] = (pc[i][0], pc[i][1], pc[i][2]) 19 | ply_out = PlyData([PlyElement.describe(vertex, 'vertex', comments=['vertices'])]) 20 | ply_out.write(filename) 21 | 22 | # Sample points on the obj shape 23 | def get_sampling_command(obj_filename, ply_filename): 24 | cmd = SAMPLING_BIN + ' ' + obj_filename 25 | cmd += ' ' + ply_filename 26 | cmd += ' -n_samples %d ' % SAMPLING_POINT_NUM 27 | cmd += ' -leaf_size %f ' % SAMPLING_LEAF_SIZE 28 | return cmd 29 | 30 | # -------------------------------------------------------------- 31 | # Following are the helper functions to load MODELNET40 shapes 32 | # -------------------------------------------------------------- 33 | 34 | # Read in the list of categories in MODELNET40 35 | def get_category_names(): 36 | shape_names_file = os.path.join(MODELNET40_PATH, 'shape_names.txt') 37 | shape_names = [line.rstrip() for line in open(shape_names_file)] 38 | return shape_names 39 | 40 | # Return all the filepaths for the shapes in MODELNET40 41 | def get_obj_filenames(): 42 | obj_filelist_file = os.path.join(MODELNET40_PATH, 'filelist.txt') 43 | obj_filenames = [os.path.join(MODELNET40_PATH, line.rstrip()) for line in open(obj_filelist_file)] 44 | print('Got %d obj files in modelnet40.' % len(obj_filenames)) 45 | return obj_filenames 46 | 47 | # Helper function to create the father folder and all subdir folders if not exist 48 | def batch_mkdir(output_folder, subdir_list): 49 | if not os.path.exists(output_folder): 50 | os.mkdir(output_folder) 51 | for subdir in subdir_list: 52 | if not os.path.exists(os.path.join(output_folder, subdir)): 53 | os.mkdir(os.path.join(output_folder, subdir)) 54 | 55 | # ---------------------------------------------------------------- 56 | # Following are the helper functions to load save/load HDF5 files 57 | # ---------------------------------------------------------------- 58 | 59 | # Write numpy array data and label to h5_filename 60 | def save_h5_data_label_normal(h5_filename, data, label, normal, 61 | data_dtype='float32', label_dtype='uint8', noral_dtype='float32'): 62 | h5_fout = h5py.File(h5_filename) 63 | h5_fout.create_dataset( 64 | 'data', data=data, 65 | compression='gzip', compression_opts=4, 66 | dtype=data_dtype) 67 | h5_fout.create_dataset( 68 | 'normal', data=normal, 69 | compression='gzip', compression_opts=4, 70 | dtype=normal_dtype) 71 | h5_fout.create_dataset( 72 | 'label', data=label, 73 | compression='gzip', compression_opts=1, 74 | dtype=label_dtype) 75 | h5_fout.close() 76 | 77 | 78 | # Write numpy array data and label to h5_filename 79 | def save_h5(h5_filename, data, label, data_dtype='uint8', label_dtype='uint8'): 80 | h5_fout = h5py.File(h5_filename) 81 | h5_fout.create_dataset( 82 | 'data', data=data, 83 | compression='gzip', compression_opts=4, 84 | dtype=data_dtype) 85 | h5_fout.create_dataset( 86 | 'label', data=label, 87 | compression='gzip', compression_opts=1, 88 | dtype=label_dtype) 89 | h5_fout.close() 90 | 91 | # Read numpy array data and label from h5_filename 92 | def load_h5_data_label_normal(h5_filename): 93 | f = h5py.File(h5_filename) 94 | data = f['data'][:] 95 | label = f['label'][:] 96 | normal = f['normal'][:] 97 | return (data, label, normal) 98 | 99 | # Read numpy array data and label from h5_filename 100 | def load_h5_data_label_seg(h5_filename): 101 | f = h5py.File(h5_filename) 102 | data = f['data'][:] 103 | label = f['label'][:] 104 | seg = f['pid'][:] 105 | return (data, label, seg) 106 | 107 | # Read numpy array data and label from h5_filename 108 | def load_h5(h5_filename): 109 | f = h5py.File(h5_filename) 110 | data = f['data'][:] 111 | label = f['label'][:] 112 | return (data, label) 113 | 114 | # ---------------------------------------------------------------- 115 | # Following are the helper functions to load save/load PLY files 116 | # ---------------------------------------------------------------- 117 | 118 | # Load PLY file 119 | def load_ply_data(filename, point_num): 120 | plydata = PlyData.read(filename) 121 | pc = plydata['vertex'].data[:point_num] 122 | pc_array = np.array([[x, y, z] for x,y,z in pc]) 123 | return pc_array 124 | 125 | # Load PLY file 126 | def load_ply_normal(filename, point_num): 127 | plydata = PlyData.read(filename) 128 | pc = plydata['normal'].data[:point_num] 129 | pc_array = np.array([[x, y, z] for x,y,z in pc]) 130 | return pc_array 131 | 132 | # Make up rows for Nxk array 133 | # Input Pad is 'edge' or 'constant' 134 | def pad_arr_rows(arr, row, pad='edge'): 135 | assert(len(arr.shape) == 2) 136 | assert(arr.shape[0] <= row) 137 | assert(pad == 'edge' or pad == 'constant') 138 | if arr.shape[0] == row: 139 | return arr 140 | if pad == 'edge': 141 | return np.lib.pad(arr, ((0, row-arr.shape[0]), (0, 0)), 'edge') 142 | if pad == 'constant': 143 | return np.lib.pad(arr, ((0, row-arr.shape[0]), (0, 0)), 'constant', (0, 0)) 144 | 145 | 146 | -------------------------------------------------------------------------------- /utils/pc_util.py: -------------------------------------------------------------------------------- 1 | """ Utility functions for processing point clouds. 2 | 3 | Author: Charles R. Qi, Hao Su 4 | Date: November 2016 5 | """ 6 | 7 | import os 8 | import sys 9 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 10 | sys.path.append(BASE_DIR) 11 | 12 | # Draw point cloud 13 | from eulerangles import euler2mat 14 | 15 | # Point cloud IO 16 | import numpy as np 17 | from plyfile import PlyData, PlyElement 18 | 19 | 20 | # ---------------------------------------- 21 | # Point Cloud/Volume Conversions 22 | # ---------------------------------------- 23 | 24 | def point_cloud_to_volume_batch(point_clouds, vsize=12, radius=1.0, flatten=True): 25 | """ Input is BxNx3 batch of point cloud 26 | Output is Bx(vsize^3) 27 | """ 28 | vol_list = [] 29 | for b in range(point_clouds.shape[0]): 30 | vol = point_cloud_to_volume(np.squeeze(point_clouds[b,:,:]), vsize, radius) 31 | if flatten: 32 | vol_list.append(vol.flatten()) 33 | else: 34 | vol_list.append(np.expand_dims(np.expand_dims(vol, -1), 0)) 35 | if flatten: 36 | return np.vstack(vol_list) 37 | else: 38 | return np.concatenate(vol_list, 0) 39 | 40 | 41 | def point_cloud_to_volume(points, vsize, radius=1.0): 42 | """ input is Nx3 points. 43 | output is vsize*vsize*vsize 44 | assumes points are in range [-radius, radius] 45 | """ 46 | vol = np.zeros((vsize,vsize,vsize)) 47 | voxel = 2*radius/float(vsize) 48 | locations = (points + radius)/voxel 49 | locations = locations.astype(int) 50 | vol[locations[:,0],locations[:,1],locations[:,2]] = 1.0 51 | return vol 52 | 53 | #a = np.zeros((16,1024,3)) 54 | #print point_cloud_to_volume_batch(a, 12, 1.0, False).shape 55 | 56 | def volume_to_point_cloud(vol): 57 | """ vol is occupancy grid (value = 0 or 1) of size vsize*vsize*vsize 58 | return Nx3 numpy array. 59 | """ 60 | vsize = vol.shape[0] 61 | assert(vol.shape[1] == vsize and vol.shape[1] == vsize) 62 | points = [] 63 | for a in range(vsize): 64 | for b in range(vsize): 65 | for c in range(vsize): 66 | if vol[a,b,c] == 1: 67 | points.append(np.array([a,b,c])) 68 | if len(points) == 0: 69 | return np.zeros((0,3)) 70 | points = np.vstack(points) 71 | return points 72 | 73 | # ---------------------------------------- 74 | # Point cloud IO 75 | # ---------------------------------------- 76 | 77 | def read_ply(filename): 78 | """ read XYZ point cloud from filename PLY file """ 79 | plydata = PlyData.read(filename) 80 | pc = plydata['vertex'].data 81 | pc_array = np.array([[x, y, z] for x,y,z in pc]) 82 | return pc_array 83 | 84 | 85 | def write_ply(points, filename, text=True): 86 | """ input: Nx3, write points to filename as PLY format. """ 87 | points = [(points[i,0], points[i,1], points[i,2]) for i in range(points.shape[0])] 88 | vertex = np.array(points, dtype=[('x', 'f4'), ('y', 'f4'),('z', 'f4')]) 89 | el = PlyElement.describe(vertex, 'vertex', comments=['vertices']) 90 | PlyData([el], text=text).write(filename) 91 | 92 | 93 | # ---------------------------------------- 94 | # Simple Point cloud and Volume Renderers 95 | # ---------------------------------------- 96 | 97 | def draw_point_cloud(input_points, canvasSize=500, space=200, diameter=25, 98 | xrot=0, yrot=0, zrot=0, switch_xyz=[0,1,2], normalize=True): 99 | """ Render point cloud to image with alpha channel. 100 | Input: 101 | points: Nx3 numpy array (+y is up direction) 102 | Output: 103 | gray image as numpy array of size canvasSizexcanvasSize 104 | """ 105 | image = np.zeros((canvasSize, canvasSize)) 106 | if input_points is None or input_points.shape[0] == 0: 107 | return image 108 | 109 | points = input_points[:, switch_xyz] 110 | M = euler2mat(zrot, yrot, xrot) 111 | points = (np.dot(M, points.transpose())).transpose() 112 | 113 | # Normalize the point cloud 114 | # We normalize scale to fit points in a unit sphere 115 | if normalize: 116 | centroid = np.mean(points, axis=0) 117 | points -= centroid 118 | furthest_distance = np.max(np.sqrt(np.sum(abs(points)**2,axis=-1))) 119 | points /= furthest_distance 120 | 121 | # Pre-compute the Gaussian disk 122 | radius = (diameter-1)/2.0 123 | disk = np.zeros((diameter, diameter)) 124 | for i in range(diameter): 125 | for j in range(diameter): 126 | if (i - radius) * (i-radius) + (j-radius) * (j-radius) <= radius * radius: 127 | disk[i, j] = np.exp((-(i-radius)**2 - (j-radius)**2)/(radius**2)) 128 | mask = np.argwhere(disk > 0) 129 | dx = mask[:, 0] 130 | dy = mask[:, 1] 131 | dv = disk[disk > 0] 132 | 133 | # Order points by z-buffer 134 | zorder = np.argsort(points[:, 2]) 135 | points = points[zorder, :] 136 | points[:, 2] = (points[:, 2] - np.min(points[:, 2])) / (np.max(points[:, 2] - np.min(points[:, 2]))) 137 | max_depth = np.max(points[:, 2]) 138 | 139 | for i in range(points.shape[0]): 140 | j = points.shape[0] - i - 1 141 | x = points[j, 0] 142 | y = points[j, 1] 143 | xc = canvasSize/2 + (x*space) 144 | yc = canvasSize/2 + (y*space) 145 | xc = int(np.round(xc)) 146 | yc = int(np.round(yc)) 147 | 148 | px = dx + xc 149 | py = dy + yc 150 | 151 | image[px, py] = image[px, py] * 0.7 + dv * (max_depth - points[j, 2]) * 0.3 152 | 153 | image = image / np.max(image) 154 | return image 155 | 156 | def point_cloud_three_views(points): 157 | """ input points Nx3 numpy array (+y is up direction). 158 | return an numpy array gray image of size 500x1500. """ 159 | # +y is up direction 160 | # xrot is azimuth 161 | # yrot is in-plane 162 | # zrot is elevation 163 | img1 = draw_point_cloud(points, zrot=110/180.0*np.pi, xrot=45/180.0*np.pi, yrot=0/180.0*np.pi) 164 | img2 = draw_point_cloud(points, zrot=70/180.0*np.pi, xrot=135/180.0*np.pi, yrot=0/180.0*np.pi) 165 | img3 = draw_point_cloud(points, zrot=180.0/180.0*np.pi, xrot=90/180.0*np.pi, yrot=0/180.0*np.pi) 166 | image_large = np.concatenate([img1, img2, img3], 1) 167 | return image_large 168 | 169 | 170 | from PIL import Image 171 | def point_cloud_three_views_demo(): 172 | """ Demo for draw_point_cloud function """ 173 | points = read_ply('../third_party/mesh_sampling/piano.ply') 174 | im_array = point_cloud_three_views(points) 175 | img = Image.fromarray(np.uint8(im_array*255.0)) 176 | img.save('piano.jpg') 177 | 178 | if __name__=="__main__": 179 | point_cloud_three_views_demo() 180 | 181 | 182 | import matplotlib.pyplot as plt 183 | def pyplot_draw_point_cloud(points, output_filename): 184 | """ points is a Nx3 numpy array """ 185 | fig = plt.figure() 186 | ax = fig.add_subplot(111, projection='3d') 187 | ax.scatter(points[:,0], points[:,1], points[:,2]) 188 | ax.set_xlabel('x') 189 | ax.set_ylabel('y') 190 | ax.set_zlabel('z') 191 | #savefig(output_filename) 192 | 193 | def pyplot_draw_volume(vol, output_filename): 194 | """ vol is of size vsize*vsize*vsize 195 | output an image to output_filename 196 | """ 197 | points = volume_to_point_cloud(vol) 198 | pyplot_draw_point_cloud(points, output_filename) 199 | -------------------------------------------------------------------------------- /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 | # Download dataset for point cloud classification 9 | DATA_DIR = os.path.join(BASE_DIR, 'data') 10 | if not os.path.exists(DATA_DIR): 11 | os.mkdir(DATA_DIR) 12 | if not os.path.exists(os.path.join(DATA_DIR, 'modelnet40_ply_hdf5_2048')): 13 | www = 'https://shapenet.cs.stanford.edu/media/modelnet40_ply_hdf5_2048.zip' 14 | zipfile = os.path.basename(www) 15 | os.system('wget %s; unzip %s' % (www, zipfile)) 16 | os.system('mv %s %s' % (zipfile[:-4], DATA_DIR)) 17 | os.system('rm %s' % (zipfile)) 18 | 19 | 20 | def shuffle_data(data, labels): 21 | """ Shuffle data and labels. 22 | Input: 23 | data: B,N,... numpy array 24 | label: B,... numpy array 25 | Return: 26 | shuffled data, label and shuffle indices 27 | """ 28 | idx = np.arange(len(labels)) 29 | np.random.shuffle(idx) 30 | return data[idx, ...], labels[idx], idx 31 | 32 | 33 | def rotate_point_cloud(batch_data): 34 | """ Randomly rotate the point clouds to augument the dataset 35 | rotation is per shape based along up direction 36 | Input: 37 | BxNx3 array, original batch of point clouds 38 | Return: 39 | BxNx3 array, rotated batch of point clouds 40 | """ 41 | rotated_data = np.zeros(batch_data.shape, dtype=np.float32) 42 | for k in range(batch_data.shape[0]): 43 | rotation_angle = np.random.uniform() * 2 * np.pi 44 | cosval = np.cos(rotation_angle) 45 | sinval = np.sin(rotation_angle) 46 | rotation_matrix = np.array([[cosval, 0, sinval], 47 | [0, 1, 0], 48 | [-sinval, 0, cosval]]) 49 | shape_pc = batch_data[k, ...] 50 | rotated_data[k, ...] = np.dot(shape_pc.reshape((-1, 3)), rotation_matrix) 51 | return rotated_data 52 | 53 | 54 | def rotate_point_cloud_by_angle(batch_data, rotation_angle): 55 | """ Rotate the point cloud along up direction with certain angle. 56 | Input: 57 | BxNx3 array, original batch of point clouds 58 | Return: 59 | BxNx3 array, rotated batch of point clouds 60 | """ 61 | rotated_data = np.zeros(batch_data.shape, dtype=np.float32) 62 | for k in range(batch_data.shape[0]): 63 | #rotation_angle = np.random.uniform() * 2 * np.pi 64 | cosval = np.cos(rotation_angle) 65 | sinval = np.sin(rotation_angle) 66 | rotation_matrix = np.array([[cosval, 0, sinval], 67 | [0, 1, 0], 68 | [-sinval, 0, cosval]]) 69 | shape_pc = batch_data[k, ...] 70 | rotated_data[k, ...] = np.dot(shape_pc.reshape((-1, 3)), rotation_matrix) 71 | return rotated_data 72 | 73 | def rotate_point_cloud_by_angle_x_axis(batch_data, rotation_angle): 74 | """ Rotate the point cloud along up direction with certain angle. 75 | Input: 76 | BxNx3 array, original batch of point clouds 77 | Return: 78 | BxNx3 array, rotated batch of point clouds 79 | """ 80 | rotated_data = np.zeros(batch_data.shape, dtype=np.float32) 81 | for k in range(batch_data.shape[0]): 82 | #rotation_angle = np.random.uniform() * 2 * np.pi 83 | cosval = np.cos(rotation_angle) 84 | sinval = np.sin(rotation_angle) 85 | rotation_matrix = np.array([[1, 0, 0], 86 | [0, cosval, -sinval], 87 | [0, sinval, cosval]]) 88 | shape_pc = batch_data[k, ...] 89 | rotated_data[k, ...] = np.dot(shape_pc.reshape((-1, 3)), rotation_matrix) 90 | return rotated_data 91 | 92 | def rotate_point_cloud_by_angle_z_axis(batch_data, rotation_angle): 93 | """ Rotate the point cloud along up direction with certain angle. 94 | Input: 95 | BxNx3 array, original batch of point clouds 96 | Return: 97 | BxNx3 array, rotated batch of point clouds 98 | """ 99 | rotated_data = np.zeros(batch_data.shape, dtype=np.float32) 100 | for k in range(batch_data.shape[0]): 101 | #rotation_angle = np.random.uniform() * 2 * np.pi 102 | cosval = np.cos(rotation_angle) 103 | sinval = np.sin(rotation_angle) 104 | rotation_matrix = np.array([[cosval, -sinval, 0], 105 | [sinval, cosval, 0], 106 | [0, 0, 1]]) 107 | shape_pc = batch_data[k, ...] 108 | rotated_data[k, ...] = np.dot(shape_pc.reshape((-1, 3)), rotation_matrix) 109 | return rotated_data 110 | 111 | 112 | def jitter_point_cloud(batch_data, sigma=0.01, clip=0.05): 113 | """ Randomly jitter points. jittering is per point. 114 | Input: 115 | BxNx3 array, original batch of point clouds 116 | Return: 117 | BxNx3 array, jittered batch of point clouds 118 | """ 119 | B, N, C = batch_data.shape 120 | assert(clip > 0) 121 | jittered_data = np.clip(sigma * np.random.randn(B, N, C), -1*clip, clip) 122 | jittered_data += batch_data 123 | return jittered_data 124 | 125 | def getDataFiles(list_filename): 126 | return [line.rstrip() for line in open(list_filename)] 127 | 128 | def load_h5(h5_filename): 129 | f = h5py.File(h5_filename) 130 | data = f['data'][:] 131 | label = f['label'][:] 132 | return (data, label) 133 | 134 | def loadDataFile(filename): 135 | return load_h5(filename) 136 | 137 | def load_h5_data_label_seg(h5_filename): 138 | f = h5py.File(h5_filename) 139 | data = f['data'][:] 140 | label = f['label'][:] 141 | seg = f['pid'][:] 142 | return (data, label, seg) 143 | 144 | 145 | def loadDataFile_with_seg(filename): 146 | return load_h5_data_label_seg(filename) 147 | 148 | 149 | # Defined by Vinit 150 | def zero_mean(data): 151 | centroid = sum(data)/(data.shape[0]*1.0) 152 | return data-centroid 153 | 154 | def normalize(data): 155 | x,y,z = data[:,0],data[:,1],data[:,2] 156 | x_max,y_max,z_max,x_min,y_min,z_min = max(x),max(y),max(z),min(x),min(y),min(z) 157 | x = x-x_min 158 | y = y-y_min 159 | z = z-z_min 160 | x = x/((x_max-x_min)*1.0) 161 | y = y/((y_max-y_min)*1.0) 162 | z = z/((z_max-z_min)*1.0) 163 | data = np.zeros((1024,3)) 164 | print max(x),max(y),max(z),min(x),min(y),min(z) 165 | data[:,0],data[:,1],data[:,2]=x,y,z 166 | return data 167 | 168 | def scale(data,scale_x,scale_y,scale_z): 169 | try: 170 | data = np.asarray(data) 171 | except: 172 | pass 173 | scale = [scale_x,scale_y,scale_z] 174 | return data*scale 175 | 176 | def translate(data,shift): 177 | try: 178 | data = np.asarray(data) 179 | except: 180 | pass 181 | # shift = [shift_x,shift_y,shift_z] 182 | return data+shift 183 | 184 | def partial_data_plane(airplane_data): 185 | airplane_data_f1,airplane_data_f2 = [],[] 186 | for i in range(airplane_data.shape[0]): 187 | if airplane_data[i,2]>-0.0363: 188 | airplane_data_f1.append(airplane_data[i,:]) 189 | else: 190 | airplane_data_f2.append(airplane_data[i,:]) 191 | return np.asarray(airplane_data_f1),np.asarray(airplane_data_f2) 192 | 193 | def clone_cloud(data): 194 | clone_data = np.zeros((32,2048,3)) 195 | for i in range(32): 196 | clone_data[i,:,:]=current_data[1,:,:] 197 | return clone_data 198 | 199 | import numpy.linalg as alg 200 | def find_lambda(f1,f2,f): 201 | f1 = np.matrix(f1.reshape((1024,1))) 202 | f2 = np.matrix(f2.reshape((1024,1))) 203 | f = np.matrix(f.reshape((1024,1))) 204 | B = f-f2 205 | A = f1-f2 206 | Atrans = A.getT() 207 | term1 = alg.inv(Atrans*A) 208 | term2 = Atrans*B 209 | return float(term1*term2) 210 | 211 | def display_clouds_data(data): 212 | import matplotlib.pyplot as plt 213 | from mpl_toolkits.mplot3d import Axes3D 214 | fig = plt.figure() 215 | ax = fig.add_subplot(111,projection='3d') 216 | try: 217 | data = data.tolist() 218 | except: 219 | pass 220 | X,Y,Z = [],[],[] 221 | for row in data: 222 | X.append(row[0]) 223 | Y.append(row[1]) 224 | Z.append(row[2]) 225 | ax.scatter(X,Y,Z) 226 | plt.show() -------------------------------------------------------------------------------- /utils/pc_distance/tf_approxmatch.cu: -------------------------------------------------------------------------------- 1 | __global__ void approxmatch(int b,int n,int m,const float * __restrict__ xyz1,const float * __restrict__ xyz2,float * __restrict__ match,float * temp){ 2 | float * remainL=temp+blockIdx.x*(n+m)*2, * remainR=temp+blockIdx.x*(n+m)*2+n,*ratioL=temp+blockIdx.x*(n+m)*2+n+m,*ratioR=temp+blockIdx.x*(n+m)*2+n+m+n; 3 | float multiL,multiR; 4 | if (n>=m){ 5 | multiL=1; 6 | multiR=n/m; 7 | }else{ 8 | multiL=m/n; 9 | multiR=1; 10 | } 11 | const int Block=1024; 12 | __shared__ float buf[Block*4]; 13 | for (int i=blockIdx.x;i=-2;j--){ 22 | float level=-powf(4.0f,j); 23 | if (j==-2){ 24 | level=0; 25 | } 26 | for (int k0=0;k0>>(b,n,m,xyz1,xyz2,match,temp); 182 | } 183 | __global__ void matchcost(int b,int n,int m,const float * __restrict__ xyz1,const float * __restrict__ xyz2,const float * __restrict__ match,float * __restrict__ out){ 184 | __shared__ float allsum[512]; 185 | const int Block=1024; 186 | __shared__ float buf[Block*3]; 187 | for (int i=blockIdx.x;i>>(b,n,m,xyz1,xyz2,match,out); 228 | } 229 | __global__ void matchcostgrad2(int b,int n,int m,const float * __restrict__ xyz1,const float * __restrict__ xyz2,const float * __restrict__ match,float * __restrict__ grad2){ 230 | __shared__ float sum_grad[256*3]; 231 | for (int i=blockIdx.x;i>>(b,n,m,xyz1,xyz2,match,grad1); 294 | matchcostgrad2<<>>(b,n,m,xyz1,xyz2,match,grad2); 295 | } 296 | 297 | -------------------------------------------------------------------------------- /utils/pc_distance/tf_nndistance.cpp: -------------------------------------------------------------------------------- 1 | #include "tensorflow/core/framework/op.h" 2 | #include "tensorflow/core/framework/op_kernel.h" 3 | REGISTER_OP("NnDistance") 4 | .Input("xyz1: float32") 5 | .Input("xyz2: float32") 6 | .Output("dist1: float32") 7 | .Output("idx1: int32") 8 | .Output("dist2: float32") 9 | .Output("idx2: int32"); 10 | REGISTER_OP("NnDistanceGrad") 11 | .Input("xyz1: float32") 12 | .Input("xyz2: float32") 13 | .Input("grad_dist1: float32") 14 | .Input("idx1: int32") 15 | .Input("grad_dist2: float32") 16 | .Input("idx2: int32") 17 | .Output("grad_xyz1: float32") 18 | .Output("grad_xyz2: float32"); 19 | using namespace tensorflow; 20 | 21 | static void nnsearch(int b,int n,int m,const float * xyz1,const float * xyz2,float * dist,int * idx){ 22 | for (int i=0;iinput(0); 50 | const Tensor& xyz2_tensor=context->input(1); 51 | OP_REQUIRES(context,xyz1_tensor.dims()==3,errors::InvalidArgument("NnDistance requires xyz1 be of shape (batch,#points,3)")); 52 | OP_REQUIRES(context,xyz1_tensor.shape().dim_size(2)==3,errors::InvalidArgument("NnDistance only accepts 3d point set xyz1")); 53 | int b=xyz1_tensor.shape().dim_size(0); 54 | int n=xyz1_tensor.shape().dim_size(1); 55 | OP_REQUIRES(context,xyz2_tensor.dims()==3,errors::InvalidArgument("NnDistance requires xyz2 be of shape (batch,#points,3)")); 56 | OP_REQUIRES(context,xyz2_tensor.shape().dim_size(2)==3,errors::InvalidArgument("NnDistance only accepts 3d point set xyz2")); 57 | int m=xyz2_tensor.shape().dim_size(1); 58 | OP_REQUIRES(context,xyz2_tensor.shape().dim_size(0)==b,errors::InvalidArgument("NnDistance expects xyz1 and xyz2 have same batch size")); 59 | auto xyz1_flat=xyz1_tensor.flat(); 60 | const float * xyz1=&xyz1_flat(0); 61 | auto xyz2_flat=xyz2_tensor.flat(); 62 | const float * xyz2=&xyz2_flat(0); 63 | Tensor * dist1_tensor=NULL; 64 | Tensor * idx1_tensor=NULL; 65 | Tensor * dist2_tensor=NULL; 66 | Tensor * idx2_tensor=NULL; 67 | OP_REQUIRES_OK(context,context->allocate_output(0,TensorShape{b,n},&dist1_tensor)); 68 | OP_REQUIRES_OK(context,context->allocate_output(1,TensorShape{b,n},&idx1_tensor)); 69 | auto dist1_flat=dist1_tensor->flat(); 70 | auto idx1_flat=idx1_tensor->flat(); 71 | OP_REQUIRES_OK(context,context->allocate_output(2,TensorShape{b,m},&dist2_tensor)); 72 | OP_REQUIRES_OK(context,context->allocate_output(3,TensorShape{b,m},&idx2_tensor)); 73 | auto dist2_flat=dist2_tensor->flat(); 74 | auto idx2_flat=idx2_tensor->flat(); 75 | float * dist1=&(dist1_flat(0)); 76 | int * idx1=&(idx1_flat(0)); 77 | float * dist2=&(dist2_flat(0)); 78 | int * idx2=&(idx2_flat(0)); 79 | nnsearch(b,n,m,xyz1,xyz2,dist1,idx1); 80 | nnsearch(b,m,n,xyz2,xyz1,dist2,idx2); 81 | } 82 | }; 83 | REGISTER_KERNEL_BUILDER(Name("NnDistance").Device(DEVICE_CPU), NnDistanceOp); 84 | class NnDistanceGradOp : public OpKernel{ 85 | public: 86 | explicit NnDistanceGradOp(OpKernelConstruction* context):OpKernel(context){} 87 | void Compute(OpKernelContext * context)override{ 88 | const Tensor& xyz1_tensor=context->input(0); 89 | const Tensor& xyz2_tensor=context->input(1); 90 | const Tensor& grad_dist1_tensor=context->input(2); 91 | const Tensor& idx1_tensor=context->input(3); 92 | const Tensor& grad_dist2_tensor=context->input(4); 93 | const Tensor& idx2_tensor=context->input(5); 94 | OP_REQUIRES(context,xyz1_tensor.dims()==3,errors::InvalidArgument("NnDistanceGrad requires xyz1 be of shape (batch,#points,3)")); 95 | OP_REQUIRES(context,xyz1_tensor.shape().dim_size(2)==3,errors::InvalidArgument("NnDistanceGrad only accepts 3d point set xyz1")); 96 | int b=xyz1_tensor.shape().dim_size(0); 97 | int n=xyz1_tensor.shape().dim_size(1); 98 | OP_REQUIRES(context,xyz2_tensor.dims()==3,errors::InvalidArgument("NnDistanceGrad requires xyz2 be of shape (batch,#points,3)")); 99 | OP_REQUIRES(context,xyz2_tensor.shape().dim_size(2)==3,errors::InvalidArgument("NnDistanceGrad only accepts 3d point set xyz2")); 100 | int m=xyz2_tensor.shape().dim_size(1); 101 | OP_REQUIRES(context,xyz2_tensor.shape().dim_size(0)==b,errors::InvalidArgument("NnDistanceGrad expects xyz1 and xyz2 have same batch size")); 102 | OP_REQUIRES(context,grad_dist1_tensor.shape()==(TensorShape{b,n}),errors::InvalidArgument("NnDistanceGrad requires grad_dist1 be of shape(batch,#points)")); 103 | OP_REQUIRES(context,idx1_tensor.shape()==(TensorShape{b,n}),errors::InvalidArgument("NnDistanceGrad requires idx1 be of shape(batch,#points)")); 104 | OP_REQUIRES(context,grad_dist2_tensor.shape()==(TensorShape{b,m}),errors::InvalidArgument("NnDistanceGrad requires grad_dist2 be of shape(batch,#points)")); 105 | OP_REQUIRES(context,idx2_tensor.shape()==(TensorShape{b,m}),errors::InvalidArgument("NnDistanceGrad requires idx2 be of shape(batch,#points)")); 106 | auto xyz1_flat=xyz1_tensor.flat(); 107 | const float * xyz1=&xyz1_flat(0); 108 | auto xyz2_flat=xyz2_tensor.flat(); 109 | const float * xyz2=&xyz2_flat(0); 110 | auto idx1_flat=idx1_tensor.flat(); 111 | const int * idx1=&idx1_flat(0); 112 | auto idx2_flat=idx2_tensor.flat(); 113 | const int * idx2=&idx2_flat(0); 114 | auto grad_dist1_flat=grad_dist1_tensor.flat(); 115 | const float * grad_dist1=&grad_dist1_flat(0); 116 | auto grad_dist2_flat=grad_dist2_tensor.flat(); 117 | const float * grad_dist2=&grad_dist2_flat(0); 118 | Tensor * grad_xyz1_tensor=NULL; 119 | OP_REQUIRES_OK(context,context->allocate_output(0,TensorShape{b,n,3},&grad_xyz1_tensor)); 120 | Tensor * grad_xyz2_tensor=NULL; 121 | OP_REQUIRES_OK(context,context->allocate_output(1,TensorShape{b,m,3},&grad_xyz2_tensor)); 122 | auto grad_xyz1_flat=grad_xyz1_tensor->flat(); 123 | float * grad_xyz1=&grad_xyz1_flat(0); 124 | auto grad_xyz2_flat=grad_xyz2_tensor->flat(); 125 | float * grad_xyz2=&grad_xyz2_flat(0); 126 | for (int i=0;iinput(0); 174 | const Tensor& xyz2_tensor=context->input(1); 175 | OP_REQUIRES(context,xyz1_tensor.dims()==3,errors::InvalidArgument("NnDistance requires xyz1 be of shape (batch,#points,3)")); 176 | OP_REQUIRES(context,xyz1_tensor.shape().dim_size(2)==3,errors::InvalidArgument("NnDistance only accepts 3d point set xyz1")); 177 | int b=xyz1_tensor.shape().dim_size(0); 178 | int n=xyz1_tensor.shape().dim_size(1); 179 | OP_REQUIRES(context,xyz2_tensor.dims()==3,errors::InvalidArgument("NnDistance requires xyz2 be of shape (batch,#points,3)")); 180 | OP_REQUIRES(context,xyz2_tensor.shape().dim_size(2)==3,errors::InvalidArgument("NnDistance only accepts 3d point set xyz2")); 181 | int m=xyz2_tensor.shape().dim_size(1); 182 | OP_REQUIRES(context,xyz2_tensor.shape().dim_size(0)==b,errors::InvalidArgument("NnDistance expects xyz1 and xyz2 have same batch size")); 183 | auto xyz1_flat=xyz1_tensor.flat(); 184 | const float * xyz1=&xyz1_flat(0); 185 | auto xyz2_flat=xyz2_tensor.flat(); 186 | const float * xyz2=&xyz2_flat(0); 187 | Tensor * dist1_tensor=NULL; 188 | Tensor * idx1_tensor=NULL; 189 | Tensor * dist2_tensor=NULL; 190 | Tensor * idx2_tensor=NULL; 191 | OP_REQUIRES_OK(context,context->allocate_output(0,TensorShape{b,n},&dist1_tensor)); 192 | OP_REQUIRES_OK(context,context->allocate_output(1,TensorShape{b,n},&idx1_tensor)); 193 | auto dist1_flat=dist1_tensor->flat(); 194 | auto idx1_flat=idx1_tensor->flat(); 195 | OP_REQUIRES_OK(context,context->allocate_output(2,TensorShape{b,m},&dist2_tensor)); 196 | OP_REQUIRES_OK(context,context->allocate_output(3,TensorShape{b,m},&idx2_tensor)); 197 | auto dist2_flat=dist2_tensor->flat(); 198 | auto idx2_flat=idx2_tensor->flat(); 199 | float * dist1=&(dist1_flat(0)); 200 | int * idx1=&(idx1_flat(0)); 201 | float * dist2=&(dist2_flat(0)); 202 | int * idx2=&(idx2_flat(0)); 203 | NmDistanceKernelLauncher(b,n,xyz1,m,xyz2,dist1,idx1,dist2,idx2); 204 | } 205 | }; 206 | REGISTER_KERNEL_BUILDER(Name("NnDistance").Device(DEVICE_GPU), NnDistanceGpuOp); 207 | 208 | void NmDistanceGradKernelLauncher(int b,int n,const float * xyz1,int m,const float * xyz2,const float * grad_dist1,const int * idx1,const float * grad_dist2,const int * idx2,float * grad_xyz1,float * grad_xyz2); 209 | class NnDistanceGradGpuOp : public OpKernel{ 210 | public: 211 | explicit NnDistanceGradGpuOp(OpKernelConstruction* context):OpKernel(context){} 212 | void Compute(OpKernelContext * context)override{ 213 | const Tensor& xyz1_tensor=context->input(0); 214 | const Tensor& xyz2_tensor=context->input(1); 215 | const Tensor& grad_dist1_tensor=context->input(2); 216 | const Tensor& idx1_tensor=context->input(3); 217 | const Tensor& grad_dist2_tensor=context->input(4); 218 | const Tensor& idx2_tensor=context->input(5); 219 | OP_REQUIRES(context,xyz1_tensor.dims()==3,errors::InvalidArgument("NnDistanceGrad requires xyz1 be of shape (batch,#points,3)")); 220 | OP_REQUIRES(context,xyz1_tensor.shape().dim_size(2)==3,errors::InvalidArgument("NnDistanceGrad only accepts 3d point set xyz1")); 221 | int b=xyz1_tensor.shape().dim_size(0); 222 | int n=xyz1_tensor.shape().dim_size(1); 223 | OP_REQUIRES(context,xyz2_tensor.dims()==3,errors::InvalidArgument("NnDistanceGrad requires xyz2 be of shape (batch,#points,3)")); 224 | OP_REQUIRES(context,xyz2_tensor.shape().dim_size(2)==3,errors::InvalidArgument("NnDistanceGrad only accepts 3d point set xyz2")); 225 | int m=xyz2_tensor.shape().dim_size(1); 226 | OP_REQUIRES(context,xyz2_tensor.shape().dim_size(0)==b,errors::InvalidArgument("NnDistanceGrad expects xyz1 and xyz2 have same batch size")); 227 | OP_REQUIRES(context,grad_dist1_tensor.shape()==(TensorShape{b,n}),errors::InvalidArgument("NnDistanceGrad requires grad_dist1 be of shape(batch,#points)")); 228 | OP_REQUIRES(context,idx1_tensor.shape()==(TensorShape{b,n}),errors::InvalidArgument("NnDistanceGrad requires idx1 be of shape(batch,#points)")); 229 | OP_REQUIRES(context,grad_dist2_tensor.shape()==(TensorShape{b,m}),errors::InvalidArgument("NnDistanceGrad requires grad_dist2 be of shape(batch,#points)")); 230 | OP_REQUIRES(context,idx2_tensor.shape()==(TensorShape{b,m}),errors::InvalidArgument("NnDistanceGrad requires idx2 be of shape(batch,#points)")); 231 | auto xyz1_flat=xyz1_tensor.flat(); 232 | const float * xyz1=&xyz1_flat(0); 233 | auto xyz2_flat=xyz2_tensor.flat(); 234 | const float * xyz2=&xyz2_flat(0); 235 | auto idx1_flat=idx1_tensor.flat(); 236 | const int * idx1=&idx1_flat(0); 237 | auto idx2_flat=idx2_tensor.flat(); 238 | const int * idx2=&idx2_flat(0); 239 | auto grad_dist1_flat=grad_dist1_tensor.flat(); 240 | const float * grad_dist1=&grad_dist1_flat(0); 241 | auto grad_dist2_flat=grad_dist2_tensor.flat(); 242 | const float * grad_dist2=&grad_dist2_flat(0); 243 | Tensor * grad_xyz1_tensor=NULL; 244 | OP_REQUIRES_OK(context,context->allocate_output(0,TensorShape{b,n,3},&grad_xyz1_tensor)); 245 | Tensor * grad_xyz2_tensor=NULL; 246 | OP_REQUIRES_OK(context,context->allocate_output(1,TensorShape{b,m,3},&grad_xyz2_tensor)); 247 | auto grad_xyz1_flat=grad_xyz1_tensor->flat(); 248 | float * grad_xyz1=&grad_xyz1_flat(0); 249 | auto grad_xyz2_flat=grad_xyz2_tensor->flat(); 250 | float * grad_xyz2=&grad_xyz2_flat(0); 251 | NmDistanceGradKernelLauncher(b,n,xyz1,m,xyz2,grad_dist1,idx1,grad_dist2,idx2,grad_xyz1,grad_xyz2); 252 | } 253 | }; 254 | REGISTER_KERNEL_BUILDER(Name("NnDistanceGrad").Device(DEVICE_GPU), NnDistanceGradGpuOp); 255 | -------------------------------------------------------------------------------- /utils/eulerangles.py: -------------------------------------------------------------------------------- 1 | # emacs: -*- mode: python-mode; py-indent-offset: 4; indent-tabs-mode: nil -*- 2 | # vi: set ft=python sts=4 ts=4 sw=4 et: 3 | ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ## 4 | # 5 | # See COPYING file distributed along with the NiBabel package for the 6 | # copyright and license terms. 7 | # 8 | ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ## 9 | ''' Module implementing Euler angle rotations and their conversions 10 | 11 | See: 12 | 13 | * http://en.wikipedia.org/wiki/Rotation_matrix 14 | * http://en.wikipedia.org/wiki/Euler_angles 15 | * http://mathworld.wolfram.com/EulerAngles.html 16 | 17 | See also: *Representing Attitude with Euler Angles and Quaternions: A 18 | Reference* (2006) by James Diebel. A cached PDF link last found here: 19 | 20 | http://citeseerx.ist.psu.edu/viewdoc/summary?doi=10.1.1.110.5134 21 | 22 | Euler's rotation theorem tells us that any rotation in 3D can be 23 | described by 3 angles. Let's call the 3 angles the *Euler angle vector* 24 | and call the angles in the vector :math:`alpha`, :math:`beta` and 25 | :math:`gamma`. The vector is [ :math:`alpha`, 26 | :math:`beta`. :math:`gamma` ] and, in this description, the order of the 27 | parameters specifies the order in which the rotations occur (so the 28 | rotation corresponding to :math:`alpha` is applied first). 29 | 30 | In order to specify the meaning of an *Euler angle vector* we need to 31 | specify the axes around which each of the rotations corresponding to 32 | :math:`alpha`, :math:`beta` and :math:`gamma` will occur. 33 | 34 | There are therefore three axes for the rotations :math:`alpha`, 35 | :math:`beta` and :math:`gamma`; let's call them :math:`i` :math:`j`, 36 | :math:`k`. 37 | 38 | Let us express the rotation :math:`alpha` around axis `i` as a 3 by 3 39 | rotation matrix `A`. Similarly :math:`beta` around `j` becomes 3 x 3 40 | matrix `B` and :math:`gamma` around `k` becomes matrix `G`. Then the 41 | whole rotation expressed by the Euler angle vector [ :math:`alpha`, 42 | :math:`beta`. :math:`gamma` ], `R` is given by:: 43 | 44 | R = np.dot(G, np.dot(B, A)) 45 | 46 | See http://mathworld.wolfram.com/EulerAngles.html 47 | 48 | The order :math:`G B A` expresses the fact that the rotations are 49 | performed in the order of the vector (:math:`alpha` around axis `i` = 50 | `A` first). 51 | 52 | To convert a given Euler angle vector to a meaningful rotation, and a 53 | rotation matrix, we need to define: 54 | 55 | * the axes `i`, `j`, `k` 56 | * whether a rotation matrix should be applied on the left of a vector to 57 | be transformed (vectors are column vectors) or on the right (vectors 58 | are row vectors). 59 | * whether the rotations move the axes as they are applied (intrinsic 60 | rotations) - compared the situation where the axes stay fixed and the 61 | vectors move within the axis frame (extrinsic) 62 | * the handedness of the coordinate system 63 | 64 | See: http://en.wikipedia.org/wiki/Rotation_matrix#Ambiguities 65 | 66 | We are using the following conventions: 67 | 68 | * axes `i`, `j`, `k` are the `z`, `y`, and `x` axes respectively. Thus 69 | an Euler angle vector [ :math:`alpha`, :math:`beta`. :math:`gamma` ] 70 | in our convention implies a :math:`alpha` radian rotation around the 71 | `z` axis, followed by a :math:`beta` rotation around the `y` axis, 72 | followed by a :math:`gamma` rotation around the `x` axis. 73 | * the rotation matrix applies on the left, to column vectors on the 74 | right, so if `R` is the rotation matrix, and `v` is a 3 x N matrix 75 | with N column vectors, the transformed vector set `vdash` is given by 76 | ``vdash = np.dot(R, v)``. 77 | * extrinsic rotations - the axes are fixed, and do not move with the 78 | rotations. 79 | * a right-handed coordinate system 80 | 81 | The convention of rotation around ``z``, followed by rotation around 82 | ``y``, followed by rotation around ``x``, is known (confusingly) as 83 | "xyz", pitch-roll-yaw, Cardan angles, or Tait-Bryan angles. 84 | ''' 85 | 86 | import math 87 | 88 | import sys 89 | if sys.version_info >= (3,0): 90 | from functools import reduce 91 | 92 | import numpy as np 93 | 94 | 95 | _FLOAT_EPS_4 = np.finfo(float).eps * 4.0 96 | 97 | 98 | def euler2mat(z=0, y=0, x=0): 99 | ''' Return matrix for rotations around z, y and x axes 100 | 101 | Uses the z, then y, then x convention above 102 | 103 | Parameters 104 | ---------- 105 | z : scalar 106 | Rotation angle in radians around z-axis (performed first) 107 | y : scalar 108 | Rotation angle in radians around y-axis 109 | x : scalar 110 | Rotation angle in radians around x-axis (performed last) 111 | 112 | Returns 113 | ------- 114 | M : array shape (3,3) 115 | Rotation matrix giving same rotation as for given angles 116 | 117 | Examples 118 | -------- 119 | >>> zrot = 1.3 # radians 120 | >>> yrot = -0.1 121 | >>> xrot = 0.2 122 | >>> M = euler2mat(zrot, yrot, xrot) 123 | >>> M.shape == (3, 3) 124 | True 125 | 126 | The output rotation matrix is equal to the composition of the 127 | individual rotations 128 | 129 | >>> M1 = euler2mat(zrot) 130 | >>> M2 = euler2mat(0, yrot) 131 | >>> M3 = euler2mat(0, 0, xrot) 132 | >>> composed_M = np.dot(M3, np.dot(M2, M1)) 133 | >>> np.allclose(M, composed_M) 134 | True 135 | 136 | You can specify rotations by named arguments 137 | 138 | >>> np.all(M3 == euler2mat(x=xrot)) 139 | True 140 | 141 | When applying M to a vector, the vector should column vector to the 142 | right of M. If the right hand side is a 2D array rather than a 143 | vector, then each column of the 2D array represents a vector. 144 | 145 | >>> vec = np.array([1, 0, 0]).reshape((3,1)) 146 | >>> v2 = np.dot(M, vec) 147 | >>> vecs = np.array([[1, 0, 0],[0, 1, 0]]).T # giving 3x2 array 148 | >>> vecs2 = np.dot(M, vecs) 149 | 150 | Rotations are counter-clockwise. 151 | 152 | >>> zred = np.dot(euler2mat(z=np.pi/2), np.eye(3)) 153 | >>> np.allclose(zred, [[0, -1, 0],[1, 0, 0], [0, 0, 1]]) 154 | True 155 | >>> yred = np.dot(euler2mat(y=np.pi/2), np.eye(3)) 156 | >>> np.allclose(yred, [[0, 0, 1],[0, 1, 0], [-1, 0, 0]]) 157 | True 158 | >>> xred = np.dot(euler2mat(x=np.pi/2), np.eye(3)) 159 | >>> np.allclose(xred, [[1, 0, 0],[0, 0, -1], [0, 1, 0]]) 160 | True 161 | 162 | Notes 163 | ----- 164 | The direction of rotation is given by the right-hand rule (orient 165 | the thumb of the right hand along the axis around which the rotation 166 | occurs, with the end of the thumb at the positive end of the axis; 167 | curl your fingers; the direction your fingers curl is the direction 168 | of rotation). Therefore, the rotations are counterclockwise if 169 | looking along the axis of rotation from positive to negative. 170 | ''' 171 | Ms = [] 172 | if z: 173 | cosz = math.cos(z) 174 | sinz = math.sin(z) 175 | Ms.append(np.array( 176 | [[cosz, -sinz, 0], 177 | [sinz, cosz, 0], 178 | [0, 0, 1]])) 179 | if y: 180 | cosy = math.cos(y) 181 | siny = math.sin(y) 182 | Ms.append(np.array( 183 | [[cosy, 0, siny], 184 | [0, 1, 0], 185 | [-siny, 0, cosy]])) 186 | if x: 187 | cosx = math.cos(x) 188 | sinx = math.sin(x) 189 | Ms.append(np.array( 190 | [[1, 0, 0], 191 | [0, cosx, -sinx], 192 | [0, sinx, cosx]])) 193 | if Ms: 194 | return reduce(np.dot, Ms[::-1]) 195 | return np.eye(3) 196 | 197 | 198 | def mat2euler(M, cy_thresh=None): 199 | ''' Discover Euler angle vector from 3x3 matrix 200 | 201 | Uses the conventions above. 202 | 203 | Parameters 204 | ---------- 205 | M : array-like, shape (3,3) 206 | cy_thresh : None or scalar, optional 207 | threshold below which to give up on straightforward arctan for 208 | estimating x rotation. If None (default), estimate from 209 | precision of input. 210 | 211 | Returns 212 | ------- 213 | z : scalar 214 | y : scalar 215 | x : scalar 216 | Rotations in radians around z, y, x axes, respectively 217 | 218 | Notes 219 | ----- 220 | If there was no numerical error, the routine could be derived using 221 | Sympy expression for z then y then x rotation matrix, which is:: 222 | 223 | [ cos(y)*cos(z), -cos(y)*sin(z), sin(y)], 224 | [cos(x)*sin(z) + cos(z)*sin(x)*sin(y), cos(x)*cos(z) - sin(x)*sin(y)*sin(z), -cos(y)*sin(x)], 225 | [sin(x)*sin(z) - cos(x)*cos(z)*sin(y), cos(z)*sin(x) + cos(x)*sin(y)*sin(z), cos(x)*cos(y)] 226 | 227 | with the obvious derivations for z, y, and x 228 | 229 | z = atan2(-r12, r11) 230 | y = asin(r13) 231 | x = atan2(-r23, r33) 232 | 233 | Problems arise when cos(y) is close to zero, because both of:: 234 | 235 | z = atan2(cos(y)*sin(z), cos(y)*cos(z)) 236 | x = atan2(cos(y)*sin(x), cos(x)*cos(y)) 237 | 238 | will be close to atan2(0, 0), and highly unstable. 239 | 240 | The ``cy`` fix for numerical instability below is from: *Graphics 241 | Gems IV*, Paul Heckbert (editor), Academic Press, 1994, ISBN: 242 | 0123361559. Specifically it comes from EulerAngles.c by Ken 243 | Shoemake, and deals with the case where cos(y) is close to zero: 244 | 245 | See: http://www.graphicsgems.org/ 246 | 247 | The code appears to be licensed (from the website) as "can be used 248 | without restrictions". 249 | ''' 250 | M = np.asarray(M) 251 | if cy_thresh is None: 252 | try: 253 | cy_thresh = np.finfo(M.dtype).eps * 4 254 | except ValueError: 255 | cy_thresh = _FLOAT_EPS_4 256 | r11, r12, r13, r21, r22, r23, r31, r32, r33 = M.flat 257 | # cy: sqrt((cos(y)*cos(z))**2 + (cos(x)*cos(y))**2) 258 | cy = math.sqrt(r33*r33 + r23*r23) 259 | if cy > cy_thresh: # cos(y) not close to zero, standard form 260 | z = math.atan2(-r12, r11) # atan2(cos(y)*sin(z), cos(y)*cos(z)) 261 | y = math.atan2(r13, cy) # atan2(sin(y), cy) 262 | x = math.atan2(-r23, r33) # atan2(cos(y)*sin(x), cos(x)*cos(y)) 263 | else: # cos(y) (close to) zero, so x -> 0.0 (see above) 264 | # so r21 -> sin(z), r22 -> cos(z) and 265 | z = math.atan2(r21, r22) 266 | y = math.atan2(r13, cy) # atan2(sin(y), cy) 267 | x = 0.0 268 | return z, y, x 269 | 270 | 271 | def euler2quat(z=0, y=0, x=0): 272 | ''' Return quaternion corresponding to these Euler angles 273 | 274 | Uses the z, then y, then x convention above 275 | 276 | Parameters 277 | ---------- 278 | z : scalar 279 | Rotation angle in radians around z-axis (performed first) 280 | y : scalar 281 | Rotation angle in radians around y-axis 282 | x : scalar 283 | Rotation angle in radians around x-axis (performed last) 284 | 285 | Returns 286 | ------- 287 | quat : array shape (4,) 288 | Quaternion in w, x, y z (real, then vector) format 289 | 290 | Notes 291 | ----- 292 | We can derive this formula in Sympy using: 293 | 294 | 1. Formula giving quaternion corresponding to rotation of theta radians 295 | about arbitrary axis: 296 | http://mathworld.wolfram.com/EulerParameters.html 297 | 2. Generated formulae from 1.) for quaternions corresponding to 298 | theta radians rotations about ``x, y, z`` axes 299 | 3. Apply quaternion multiplication formula - 300 | http://en.wikipedia.org/wiki/Quaternions#Hamilton_product - to 301 | formulae from 2.) to give formula for combined rotations. 302 | ''' 303 | z = z/2.0 304 | y = y/2.0 305 | x = x/2.0 306 | cz = math.cos(z) 307 | sz = math.sin(z) 308 | cy = math.cos(y) 309 | sy = math.sin(y) 310 | cx = math.cos(x) 311 | sx = math.sin(x) 312 | return np.array([ 313 | cx*cy*cz - sx*sy*sz, 314 | cx*sy*sz + cy*cz*sx, 315 | cx*cz*sy - sx*cy*sz, 316 | cx*cy*sz + sx*cz*sy]) 317 | 318 | 319 | def quat2euler(q): 320 | ''' Return Euler angles corresponding to quaternion `q` 321 | 322 | Parameters 323 | ---------- 324 | q : 4 element sequence 325 | w, x, y, z of quaternion 326 | 327 | Returns 328 | ------- 329 | z : scalar 330 | Rotation angle in radians around z-axis (performed first) 331 | y : scalar 332 | Rotation angle in radians around y-axis 333 | x : scalar 334 | Rotation angle in radians around x-axis (performed last) 335 | 336 | Notes 337 | ----- 338 | It's possible to reduce the amount of calculation a little, by 339 | combining parts of the ``quat2mat`` and ``mat2euler`` functions, but 340 | the reduction in computation is small, and the code repetition is 341 | large. 342 | ''' 343 | # delayed import to avoid cyclic dependencies 344 | import nibabel.quaternions as nq 345 | return mat2euler(nq.quat2mat(q)) 346 | 347 | 348 | def euler2angle_axis(z=0, y=0, x=0): 349 | ''' Return angle, axis corresponding to these Euler angles 350 | 351 | Uses the z, then y, then x convention above 352 | 353 | Parameters 354 | ---------- 355 | z : scalar 356 | Rotation angle in radians around z-axis (performed first) 357 | y : scalar 358 | Rotation angle in radians around y-axis 359 | x : scalar 360 | Rotation angle in radians around x-axis (performed last) 361 | 362 | Returns 363 | ------- 364 | theta : scalar 365 | angle of rotation 366 | vector : array shape (3,) 367 | axis around which rotation occurs 368 | 369 | Examples 370 | -------- 371 | >>> theta, vec = euler2angle_axis(0, 1.5, 0) 372 | >>> print(theta) 373 | 1.5 374 | >>> np.allclose(vec, [0, 1, 0]) 375 | True 376 | ''' 377 | # delayed import to avoid cyclic dependencies 378 | import nibabel.quaternions as nq 379 | return nq.quat2angle_axis(euler2quat(z, y, x)) 380 | 381 | 382 | def angle_axis2euler(theta, vector, is_normalized=False): 383 | ''' Convert angle, axis pair to Euler angles 384 | 385 | Parameters 386 | ---------- 387 | theta : scalar 388 | angle of rotation 389 | vector : 3 element sequence 390 | vector specifying axis for rotation. 391 | is_normalized : bool, optional 392 | True if vector is already normalized (has norm of 1). Default 393 | False 394 | 395 | Returns 396 | ------- 397 | z : scalar 398 | y : scalar 399 | x : scalar 400 | Rotations in radians around z, y, x axes, respectively 401 | 402 | Examples 403 | -------- 404 | >>> z, y, x = angle_axis2euler(0, [1, 0, 0]) 405 | >>> np.allclose((z, y, x), 0) 406 | True 407 | 408 | Notes 409 | ----- 410 | It's possible to reduce the amount of calculation a little, by 411 | combining parts of the ``angle_axis2mat`` and ``mat2euler`` 412 | functions, but the reduction in computation is small, and the code 413 | repetition is large. 414 | ''' 415 | # delayed import to avoid cyclic dependencies 416 | import nibabel.quaternions as nq 417 | M = nq.angle_axis2mat(theta, vector, is_normalized) 418 | return mat2euler(M) 419 | -------------------------------------------------------------------------------- /utils/pc_distance/tf_approxmatch.cpp: -------------------------------------------------------------------------------- 1 | #include "tensorflow/core/framework/op.h" 2 | #include "tensorflow/core/framework/op_kernel.h" 3 | #include 4 | #include 5 | #include 6 | using namespace tensorflow; 7 | REGISTER_OP("ApproxMatch") 8 | .Input("xyz1: float32") 9 | .Input("xyz2: float32") 10 | .Output("match: float32"); 11 | REGISTER_OP("MatchCost") 12 | .Input("xyz1: float32") 13 | .Input("xyz2: float32") 14 | .Input("match: float32") 15 | .Output("cost: float32"); 16 | REGISTER_OP("MatchCostGrad") 17 | .Input("xyz1: float32") 18 | .Input("xyz2: float32") 19 | .Input("match: float32") 20 | .Output("grad1: float32") 21 | .Output("grad2: float32"); 22 | 23 | void approxmatch_cpu(int b,int n,int m,const float * xyz1,const float * xyz2,float * match){ 24 | for (int i=0;i saturatedl(n,double(factorl)),saturatedr(m,double(factorr)); 28 | std::vector weight(n*m); 29 | for (int j=0;j=-2;j--){ 32 | //printf("i=%d j=%d\n",i,j); 33 | double level=-powf(4.0,j); 34 | if (j==-2) 35 | level=0; 36 | for (int k=0;k ss(m,1e-9); 48 | for (int k=0;k ss2(m,0); 65 | for (int k=0;kinput(0); 150 | OP_REQUIRES(context,xyz1_tensor.dims()==3 && xyz1_tensor.shape().dim_size(2)==3,errors::InvalidArgument("ApproxMatch expects (batch_size,num_points,3) xyz1 shape")); 151 | auto xyz1_flat=xyz1_tensor.flat(); 152 | const float * xyz1=&(xyz1_flat(0)); 153 | int b=xyz1_tensor.shape().dim_size(0); 154 | int n=xyz1_tensor.shape().dim_size(1); 155 | //OP_REQUIRES(context,n<=4096,errors::InvalidArgument("ApproxMatch handles at most 4096 dataset points")); 156 | 157 | const Tensor& xyz2_tensor=context->input(1); 158 | OP_REQUIRES(context,xyz2_tensor.dims()==3 && xyz2_tensor.shape().dim_size(2)==3 && xyz2_tensor.shape().dim_size(0)==b,errors::InvalidArgument("ApproxMatch expects (batch_size,num_points,3) xyz2 shape, and batch_size must match")); 159 | int m=xyz2_tensor.shape().dim_size(1); 160 | //OP_REQUIRES(context,m<=1024,errors::InvalidArgument("ApproxMatch handles at most 1024 query points")); 161 | auto xyz2_flat=xyz2_tensor.flat(); 162 | const float * xyz2=&(xyz2_flat(0)); 163 | Tensor * match_tensor=NULL; 164 | OP_REQUIRES_OK(context,context->allocate_output(0,TensorShape{b,m,n},&match_tensor)); 165 | auto match_flat=match_tensor->flat(); 166 | float * match=&(match_flat(0)); 167 | Tensor temp_tensor; 168 | OP_REQUIRES_OK(context,context->allocate_temp(DataTypeToEnum::value,TensorShape{b,(n+m)*2},&temp_tensor)); 169 | auto temp_flat=temp_tensor.flat(); 170 | float * temp=&(temp_flat(0)); 171 | approxmatchLauncher(b,n,m,xyz1,xyz2,match,temp); 172 | } 173 | }; 174 | REGISTER_KERNEL_BUILDER(Name("ApproxMatch").Device(DEVICE_GPU), ApproxMatchGpuOp); 175 | class ApproxMatchOp: public OpKernel{ 176 | public: 177 | explicit ApproxMatchOp(OpKernelConstruction* context):OpKernel(context){} 178 | void Compute(OpKernelContext * context)override{ 179 | const Tensor& xyz1_tensor=context->input(0); 180 | OP_REQUIRES(context,xyz1_tensor.dims()==3 && xyz1_tensor.shape().dim_size(2)==3,errors::InvalidArgument("ApproxMatch expects (batch_size,num_points,3) xyz1 shape")); 181 | auto xyz1_flat=xyz1_tensor.flat(); 182 | const float * xyz1=&(xyz1_flat(0)); 183 | int b=xyz1_tensor.shape().dim_size(0); 184 | int n=xyz1_tensor.shape().dim_size(1); 185 | //OP_REQUIRES(context,n<=4096,errors::InvalidArgument("ApproxMatch handles at most 4096 dataset points")); 186 | 187 | const Tensor& xyz2_tensor=context->input(1); 188 | OP_REQUIRES(context,xyz2_tensor.dims()==3 && xyz2_tensor.shape().dim_size(2)==3 && xyz2_tensor.shape().dim_size(0)==b,errors::InvalidArgument("ApproxMatch expects (batch_size,num_points,3) xyz2 shape, and batch_size must match")); 189 | int m=xyz2_tensor.shape().dim_size(1); 190 | //OP_REQUIRES(context,m<=1024,errors::InvalidArgument("ApproxMatch handles at most 1024 query points")); 191 | auto xyz2_flat=xyz2_tensor.flat(); 192 | const float * xyz2=&(xyz2_flat(0)); 193 | Tensor * match_tensor=NULL; 194 | OP_REQUIRES_OK(context,context->allocate_output(0,TensorShape{b,m,n},&match_tensor)); 195 | auto match_flat=match_tensor->flat(); 196 | float * match=&(match_flat(0)); 197 | approxmatch_cpu(b,n,m,xyz1,xyz2,match); 198 | } 199 | }; 200 | REGISTER_KERNEL_BUILDER(Name("ApproxMatch").Device(DEVICE_CPU), ApproxMatchOp); 201 | class MatchCostGpuOp: public OpKernel{ 202 | public: 203 | explicit MatchCostGpuOp(OpKernelConstruction* context):OpKernel(context){} 204 | void Compute(OpKernelContext * context)override{ 205 | const Tensor& xyz1_tensor=context->input(0); 206 | OP_REQUIRES(context,xyz1_tensor.dims()==3 && xyz1_tensor.shape().dim_size(2)==3,errors::InvalidArgument("MatchCost expects (batch_size,num_points,3) xyz1 shape")); 207 | auto xyz1_flat=xyz1_tensor.flat(); 208 | const float * xyz1=&(xyz1_flat(0)); 209 | int b=xyz1_tensor.shape().dim_size(0); 210 | int n=xyz1_tensor.shape().dim_size(1); 211 | 212 | const Tensor& xyz2_tensor=context->input(1); 213 | OP_REQUIRES(context,xyz2_tensor.dims()==3 && xyz2_tensor.shape().dim_size(2)==3 && xyz2_tensor.shape().dim_size(0)==b,errors::InvalidArgument("MatchCost expects (batch_size,num_points,3) xyz2 shape, and batch_size must match")); 214 | int m=xyz2_tensor.shape().dim_size(1); 215 | auto xyz2_flat=xyz2_tensor.flat(); 216 | const float * xyz2=&(xyz2_flat(0)); 217 | 218 | const Tensor& match_tensor=context->input(2); 219 | OP_REQUIRES(context,match_tensor.dims()==3 && match_tensor.shape().dim_size(0)==b && match_tensor.shape().dim_size(1)==m && match_tensor.shape().dim_size(2)==n,errors::InvalidArgument("MatchCost expects (batch_size,#query,#dataset) match shape")); 220 | auto match_flat=match_tensor.flat(); 221 | const float * match=&(match_flat(0)); 222 | 223 | Tensor * cost_tensor=NULL; 224 | OP_REQUIRES_OK(context,context->allocate_output(0,TensorShape{b},&cost_tensor)); 225 | auto cost_flat=cost_tensor->flat(); 226 | float * cost=&(cost_flat(0)); 227 | matchcostLauncher(b,n,m,xyz1,xyz2,match,cost); 228 | } 229 | }; 230 | REGISTER_KERNEL_BUILDER(Name("MatchCost").Device(DEVICE_GPU), MatchCostGpuOp); 231 | class MatchCostOp: public OpKernel{ 232 | public: 233 | explicit MatchCostOp(OpKernelConstruction* context):OpKernel(context){} 234 | void Compute(OpKernelContext * context)override{ 235 | const Tensor& xyz1_tensor=context->input(0); 236 | OP_REQUIRES(context,xyz1_tensor.dims()==3 && xyz1_tensor.shape().dim_size(2)==3,errors::InvalidArgument("MatchCost expects (batch_size,num_points,3) xyz1 shape")); 237 | auto xyz1_flat=xyz1_tensor.flat(); 238 | const float * xyz1=&(xyz1_flat(0)); 239 | int b=xyz1_tensor.shape().dim_size(0); 240 | int n=xyz1_tensor.shape().dim_size(1); 241 | 242 | const Tensor& xyz2_tensor=context->input(1); 243 | OP_REQUIRES(context,xyz2_tensor.dims()==3 && xyz2_tensor.shape().dim_size(2)==3 && xyz2_tensor.shape().dim_size(0)==b,errors::InvalidArgument("MatchCost expects (batch_size,num_points,3) xyz2 shape, and batch_size must match")); 244 | int m=xyz2_tensor.shape().dim_size(1); 245 | auto xyz2_flat=xyz2_tensor.flat(); 246 | const float * xyz2=&(xyz2_flat(0)); 247 | 248 | const Tensor& match_tensor=context->input(2); 249 | OP_REQUIRES(context,match_tensor.dims()==3 && match_tensor.shape().dim_size(0)==b && match_tensor.shape().dim_size(1)==m && match_tensor.shape().dim_size(2)==n,errors::InvalidArgument("MatchCost expects (batch_size,#query,#dataset) match shape")); 250 | auto match_flat=match_tensor.flat(); 251 | const float * match=&(match_flat(0)); 252 | 253 | Tensor * cost_tensor=NULL; 254 | OP_REQUIRES_OK(context,context->allocate_output(0,TensorShape{b},&cost_tensor)); 255 | auto cost_flat=cost_tensor->flat(); 256 | float * cost=&(cost_flat(0)); 257 | matchcost_cpu(b,n,m,xyz1,xyz2,match,cost); 258 | } 259 | }; 260 | REGISTER_KERNEL_BUILDER(Name("MatchCost").Device(DEVICE_CPU), MatchCostOp); 261 | 262 | class MatchCostGradGpuOp: public OpKernel{ 263 | public: 264 | explicit MatchCostGradGpuOp(OpKernelConstruction* context):OpKernel(context){} 265 | void Compute(OpKernelContext * context)override{ 266 | const Tensor& xyz1_tensor=context->input(0); 267 | OP_REQUIRES(context,xyz1_tensor.dims()==3 && xyz1_tensor.shape().dim_size(2)==3,errors::InvalidArgument("MatchCostGrad expects (batch_size,num_points,3) xyz1 shape")); 268 | auto xyz1_flat=xyz1_tensor.flat(); 269 | const float * xyz1=&(xyz1_flat(0)); 270 | int b=xyz1_tensor.shape().dim_size(0); 271 | int n=xyz1_tensor.shape().dim_size(1); 272 | 273 | const Tensor& xyz2_tensor=context->input(1); 274 | OP_REQUIRES(context,xyz2_tensor.dims()==3 && xyz2_tensor.shape().dim_size(2)==3 && xyz2_tensor.shape().dim_size(0)==b,errors::InvalidArgument("MatchCostGrad expects (batch_size,num_points,3) xyz2 shape, and batch_size must match")); 275 | int m=xyz2_tensor.shape().dim_size(1); 276 | auto xyz2_flat=xyz2_tensor.flat(); 277 | const float * xyz2=&(xyz2_flat(0)); 278 | 279 | const Tensor& match_tensor=context->input(2); 280 | OP_REQUIRES(context,match_tensor.dims()==3 && match_tensor.shape().dim_size(0)==b && match_tensor.shape().dim_size(1)==m && match_tensor.shape().dim_size(2)==n,errors::InvalidArgument("MatchCost expects (batch_size,#query,#dataset) match shape")); 281 | auto match_flat=match_tensor.flat(); 282 | const float * match=&(match_flat(0)); 283 | 284 | Tensor * grad1_tensor=NULL; 285 | OP_REQUIRES_OK(context,context->allocate_output(0,TensorShape{b,n,3},&grad1_tensor)); 286 | auto grad1_flat=grad1_tensor->flat(); 287 | float * grad1=&(grad1_flat(0)); 288 | Tensor * grad2_tensor=NULL; 289 | OP_REQUIRES_OK(context,context->allocate_output(1,TensorShape{b,m,3},&grad2_tensor)); 290 | auto grad2_flat=grad2_tensor->flat(); 291 | float * grad2=&(grad2_flat(0)); 292 | matchcostgradLauncher(b,n,m,xyz1,xyz2,match,grad1,grad2); 293 | } 294 | }; 295 | REGISTER_KERNEL_BUILDER(Name("MatchCostGrad").Device(DEVICE_GPU), MatchCostGradGpuOp); 296 | class MatchCostGradOp: public OpKernel{ 297 | public: 298 | explicit MatchCostGradOp(OpKernelConstruction* context):OpKernel(context){} 299 | void Compute(OpKernelContext * context)override{ 300 | const Tensor& xyz1_tensor=context->input(0); 301 | OP_REQUIRES(context,xyz1_tensor.dims()==3 && xyz1_tensor.shape().dim_size(2)==3,errors::InvalidArgument("MatchCost expects (batch_size,num_points,3) xyz1 shape")); 302 | auto xyz1_flat=xyz1_tensor.flat(); 303 | const float * xyz1=&(xyz1_flat(0)); 304 | int b=xyz1_tensor.shape().dim_size(0); 305 | int n=xyz1_tensor.shape().dim_size(1); 306 | 307 | const Tensor& xyz2_tensor=context->input(1); 308 | OP_REQUIRES(context,xyz2_tensor.dims()==3 && xyz2_tensor.shape().dim_size(2)==3 && xyz2_tensor.shape().dim_size(0)==b,errors::InvalidArgument("MatchCost expects (batch_size,num_points,3) xyz2 shape, and batch_size must match")); 309 | int m=xyz2_tensor.shape().dim_size(1); 310 | auto xyz2_flat=xyz2_tensor.flat(); 311 | const float * xyz2=&(xyz2_flat(0)); 312 | 313 | const Tensor& match_tensor=context->input(2); 314 | OP_REQUIRES(context,match_tensor.dims()==3 && match_tensor.shape().dim_size(0)==b && match_tensor.shape().dim_size(1)==m && match_tensor.shape().dim_size(2)==n,errors::InvalidArgument("MatchCost expects (batch_size,#query,#dataset) match shape")); 315 | auto match_flat=match_tensor.flat(); 316 | const float * match=&(match_flat(0)); 317 | 318 | Tensor * grad1_tensor=NULL; 319 | OP_REQUIRES_OK(context,context->allocate_output(0,TensorShape{b,n,3},&grad1_tensor)); 320 | auto grad1_flat=grad1_tensor->flat(); 321 | float * grad1=&(grad1_flat(0)); 322 | Tensor * grad2_tensor=NULL; 323 | OP_REQUIRES_OK(context,context->allocate_output(1,TensorShape{b,m,3},&grad2_tensor)); 324 | auto grad2_flat=grad2_tensor->flat(); 325 | float * grad2=&(grad2_flat(0)); 326 | matchcostgrad_cpu(b,n,m,xyz1,xyz2,match,grad1,grad2); 327 | } 328 | }; 329 | REGISTER_KERNEL_BUILDER(Name("MatchCostGrad").Device(DEVICE_CPU), MatchCostGradOp); 330 | -------------------------------------------------------------------------------- /pointnet_autoencoder_train.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import math 3 | import h5py 4 | import numpy as np 5 | import tensorflow as tf 6 | import socket 7 | import importlib 8 | import os 9 | import sys 10 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 11 | sys.path.append(BASE_DIR) 12 | sys.path.append(os.path.join(BASE_DIR, 'models')) 13 | sys.path.append(os.path.join(BASE_DIR, 'utils')) 14 | import tf_util 15 | import helper 16 | import transforms3d.euler as t3d 17 | import provider 18 | 19 | parser = argparse.ArgumentParser() 20 | parser.add_argument('--mode', type=str, default='no_mode', help='mode: train or test') 21 | parser.add_argument('--gpu', type=int, default=0, help='GPU to use [default: GPU 0]') 22 | parser.add_argument('--model', default='pointnet_pose', help='Model name: pointnet_cls or pointnet_cls_basic [default: pointnet_cls]') 23 | parser.add_argument('--log_dir', default='log_trial3', help='Log dir [default: log]') 24 | parser.add_argument('--num_point', type=int, default=1024, help='Point Number [256/512/1024/2048] [default: 1024]') 25 | parser.add_argument('--max_epoch', type=int, default=250, help='Epoch to run [default: 250]') 26 | parser.add_argument('--batch_size', type=int, default=10, help='Batch Size during training [default: 32]') 27 | parser.add_argument('--learning_rate', type=float,default=0.001, help='Initial learning rate [default: 0.001]') 28 | parser.add_argument('--momentum', type=float, default=0.9, help='Initial learning rate [default: 0.9]') 29 | parser.add_argument('--optimizer', default='adam', help='adam or momentum [default: adam]') 30 | parser.add_argument('--decay_step', type=int, default=200000, help='Decay step for lr decay [default: 200000]') 31 | parser.add_argument('--decay_rate', type=float, default=0.7, help='Decay rate for lr decay [default: 0.8]') 32 | parser.add_argument('--model_path', type=str, default='log_trial2/model200.ckpt', help='Path of the weights (.ckpt file) to be used for test') 33 | parser.add_argument('--centroid_sub', type=bool, default=False, help='Centroid Subtraction from Source and Template before Pose Prediction.') 34 | parser.add_argument('--use_pretrained_model', type=bool, default=False, help='Use a pretrained model of airplane to initialize the training.') 35 | parser.add_argument('--use_random_poses', type=bool, default=False, help='Use of random poses to train the model in each batch') 36 | parser.add_argument('--data_dict', type=str, default='templates',help='Data used to train templates or multi_model_templates') 37 | parser.add_argument('--train_poses', type=str, default='itr_net_train_data.csv', help='Poses for training') 38 | parser.add_argument('--eval_poses', type=str, default='itr_net_eval_data.csv', help='Poses for evaluation') 39 | parser.add_argument('--feature_size', type=int, default=1024, help='Size of features extracted from PointNet') 40 | FLAGS = parser.parse_args() 41 | 42 | 43 | # ModelNet40 official train/test split 44 | TRAIN_FILES = provider.getDataFiles( \ 45 | os.path.join(BASE_DIR, 'data/modelnet40_ply_hdf5_2048/train_files.txt')) 46 | TEST_FILES = provider.getDataFiles(\ 47 | os.path.join(BASE_DIR, 'data/modelnet40_ply_hdf5_2048/test_files.txt')) 48 | 49 | TRAIN_POSES = FLAGS.train_poses 50 | EVAL_POSES = FLAGS.eval_poses 51 | 52 | # Change batch size during test mode. 53 | if FLAGS.mode == 'test': 54 | BATCH_SIZE = 1 55 | FLAGS.model = 'pointnet_pose_test' 56 | else: 57 | BATCH_SIZE = FLAGS.batch_size 58 | 59 | # Parameters for data 60 | NUM_POINT = FLAGS.num_point 61 | MAX_NUM_POINT = 2048 62 | NUM_CLASSES = 40 63 | centroid_subtraction_switch = FLAGS.centroid_sub 64 | 65 | # Network hyperparameters 66 | MAX_EPOCH = FLAGS.max_epoch 67 | BASE_LEARNING_RATE = FLAGS.learning_rate 68 | GPU_INDEX = FLAGS.gpu 69 | MOMENTUM = FLAGS.momentum 70 | OPTIMIZER = FLAGS.optimizer 71 | DECAY_STEP = FLAGS.decay_step 72 | DECAY_RATE = FLAGS.decay_rate 73 | BN_INIT_DECAY = 0.5 74 | BN_DECAY_DECAY_RATE = 0.5 75 | BN_DECAY_DECAY_STEP = float(DECAY_STEP) 76 | BN_DECAY_CLIP = 0.99 77 | 78 | # Model Import 79 | MODEL = importlib.import_module(FLAGS.model) # import network module 80 | MODEL_FILE = os.path.join(BASE_DIR, 'models', FLAGS.model+'.py') 81 | LOG_DIR = FLAGS.log_dir 82 | 83 | # Take backup of all files used to train the network with all the parameters. 84 | if FLAGS.mode == 'train': 85 | if not os.path.exists(LOG_DIR): os.mkdir(LOG_DIR) # Create Log_dir to store the log. 86 | os.system('cp %s %s' % (MODEL_FILE, LOG_DIR)) # bkp of model def 87 | os.system('cp pointnet_autoencoder_train.py %s' % (LOG_DIR)) # bkp of train procedure 88 | os.system('cp -a utils/ %s/'%(LOG_DIR)) # Store the utils code. 89 | os.system('cp helper.py %s'%(LOG_DIR)) 90 | LOG_FOUT = open(os.path.join(LOG_DIR, 'log_train.txt'), 'w')# Create a text file to store the loss function data. 91 | LOG_FOUT.write(str(FLAGS)+'\n') 92 | 93 | # Write all the data of loss function during training. 94 | def log_string(out_str): 95 | LOG_FOUT.write(out_str+'\n') 96 | LOG_FOUT.flush() 97 | print(out_str) 98 | 99 | # Calculate Learning Rate during training. 100 | def get_learning_rate(batch): 101 | learning_rate = tf.train.exponential_decay( 102 | BASE_LEARNING_RATE, # Base learning rate. 103 | batch * BATCH_SIZE, # Current index into the dataset. 104 | DECAY_STEP, # Decay step. 105 | DECAY_RATE, # Decay rate. 106 | staircase=True) 107 | learning_rate = tf.maximum(learning_rate, 0.00001) # CLIP THE LEARNING RATE! 108 | return learning_rate 109 | 110 | # Get Batch Normalization decay. 111 | def get_bn_decay(batch): 112 | bn_momentum = tf.train.exponential_decay( 113 | BN_INIT_DECAY, 114 | batch*BATCH_SIZE, 115 | BN_DECAY_DECAY_STEP, 116 | BN_DECAY_DECAY_RATE, 117 | staircase=True) 118 | bn_decay = tf.minimum(BN_DECAY_CLIP, 1 - bn_momentum) 119 | return bn_decay 120 | 121 | def train(): 122 | with tf.Graph().as_default(): 123 | with tf.device('/cpu:0'): 124 | batch = tf.Variable(0) # That tells the optimizer to helpfully increment the 'batch' parameter for you every time it trains. 125 | 126 | with tf.device('/gpu:'+str(GPU_INDEX)): 127 | is_training_pl = tf.placeholder(tf.bool, shape=()) # Flag for dropouts. 128 | bn_decay = get_bn_decay(batch) # Calculate BN decay. 129 | learning_rate = get_learning_rate(batch) # Calculate Learning Rate at each step. 130 | # Define a network to backpropagate the using final pose prediction. 131 | with tf.variable_scope('Network') as _: 132 | # Object of network class. 133 | network = MODEL.Network() 134 | # Get the placeholders. 135 | source_pointclouds_pl = network.placeholder_inputs(BATCH_SIZE, NUM_POINT) 136 | # Extract Features. 137 | source_global_feature = network.get_model(source_pointclouds_pl, FLAGS.feature_size, is_training_pl, bn_decay=bn_decay) 138 | # Find the predicted transformation. 139 | predicted_pointclouds_pl = network.decode_data(source_global_feature, is_training_pl, bn_decay=bn_decay) 140 | # Find the loss using source and transformed template point cloud. 141 | loss = network.get_loss_b(predicted_pointclouds_pl, source_pointclouds_pl) 142 | 143 | # Get training optimization algorithm. 144 | if OPTIMIZER == 'momentum': 145 | optimizer = tf.train.MomentumOptimizer(learning_rate, momentum=MOMENTUM) 146 | elif OPTIMIZER == 'adam': 147 | optimizer = tf.train.AdamOptimizer(learning_rate) 148 | 149 | # Update Network_L. 150 | train_op = optimizer.minimize(loss, global_step=batch) 151 | 152 | with tf.device('/cpu:0'): 153 | # Add the loss in tensorboard. 154 | tf.summary.scalar('learning_rate', learning_rate) 155 | tf.summary.scalar('bn_decay', bn_decay) # Write BN decay in summary. 156 | saver = tf.train.Saver() 157 | tf.summary.scalar('loss', loss) 158 | 159 | # Create a session 160 | config = tf.ConfigProto() 161 | config.gpu_options.allow_growth = True 162 | config.allow_soft_placement = True 163 | config.log_device_placement = False 164 | sess = tf.Session(config=config) 165 | 166 | # Add summary writers 167 | merged = tf.summary.merge_all() 168 | if FLAGS.mode == 'train': # Create summary writers only for train mode. 169 | train_writer = tf.summary.FileWriter(os.path.join(LOG_DIR, 'train'), 170 | sess.graph) 171 | eval_writer = tf.summary.FileWriter(os.path.join(LOG_DIR, 'eval')) 172 | 173 | # Init variables 174 | init = tf.global_variables_initializer() 175 | sess.run(init, {is_training_pl: True}) 176 | 177 | # Just to initialize weights with pretrained model. 178 | if FLAGS.use_pretrained_model: 179 | saver.restore(sess,os.path.join('log_trial1','model2000.ckpt')) 180 | 181 | # Create a dictionary to pass the tensors and placeholders in train and eval function for Network_L. 182 | ops = {'source_pointclouds_pl': source_pointclouds_pl, 183 | 'is_training_pl': is_training_pl, 184 | 'predicted_pointclouds_pl': predicted_pointclouds_pl, 185 | 'loss': loss, 186 | 'train_op': train_op, 187 | 'merged': merged, 188 | 'step': batch} 189 | 190 | if FLAGS.mode == 'train': 191 | # For actual training. 192 | for epoch in range(MAX_EPOCH): 193 | log_string('**** EPOCH %03d ****' % (epoch)) 194 | sys.stdout.flush() 195 | # Train for all triaining poses. 196 | train_one_epoch(sess, ops, train_writer) 197 | # Save the variables to disk. 198 | saver.save(sess, os.path.join(LOG_DIR, "model.ckpt")) 199 | if epoch % 50 == 0: 200 | # Evaluate the trained network after 50 epochs. 201 | eval_one_epoch(sess, ops, eval_writer) 202 | # Store the Trained weights in log directory. 203 | #if epoch % 50 == 0: 204 | save_path = saver.save(sess, os.path.join(LOG_DIR, "model"+str(epoch)+".ckpt")) 205 | log_string("Model saved in file: %s" % save_path) 206 | 207 | if FLAGS.mode == 'test': 208 | # Just to test the results 209 | test_one_epoch(sess, ops, saver, FLAGS.model_path) 210 | 211 | 212 | # Train the Network_L and copy weights from Network_L to Network19 to find the poses between source and template. 213 | def train_one_epoch(sess, ops, train_writer): 214 | # Arguments: 215 | # sess: Tensorflow session to handle tensors. 216 | # ops: Dictionary for tensors of Network_L 217 | # ops19: Dictionary for tensors of Network19 218 | # templates: Training Point Cloud data. 219 | # poses: Training pose data. 220 | 221 | is_training = True 222 | display_ptClouds = False 223 | 224 | train_file_idxs = np.arange(0, len(TRAIN_FILES)) 225 | np.random.shuffle(train_file_idxs) 226 | 227 | for fn in range(len(TRAIN_FILES)): 228 | log_string('----' + str(fn) + '-----') 229 | current_data, current_label = provider.loadDataFile(TRAIN_FILES[train_file_idxs[fn]]) 230 | current_data = current_data[:,0:NUM_POINT,:] 231 | current_data, _, _ = provider.shuffle_data(current_data, np.squeeze(current_label)) 232 | 233 | file_size = current_data.shape[0] 234 | num_batches = file_size // BATCH_SIZE 235 | loss_sum = 0 236 | 237 | # Training for each batch. 238 | for fn in range(num_batches): 239 | start_idx = fn*BATCH_SIZE # Start index of poses. 240 | end_idx = (fn+1)*BATCH_SIZE # End index of poses. 241 | 242 | template_data = current_data[start_idx:end_idx,:,:] 243 | 244 | template_data = provider.rotate_point_cloud(template_data) 245 | template_data = provider.jitter_point_cloud(template_data) 246 | 247 | # To visualize the source and point clouds: 248 | if display_ptClouds: 249 | helper.display_clouds_data(template_data[0]) 250 | 251 | # Feed the placeholders of Network_L with source data and template data obtained from N-Iterations. 252 | feed_dict = {ops['source_pointclouds_pl']: template_data, 253 | ops['is_training_pl']: is_training} 254 | 255 | # Ask the network to predict transformation, calculate loss using distance between actual points, calculate & apply gradients for Network_L and copy the weights to Network19. 256 | summary, step, _, loss_val = sess.run([ops['merged'], ops['step'], ops['train_op'], ops['loss']], feed_dict=feed_dict) 257 | train_writer.add_summary(summary, step) # Add all the summary to the tensorboard. 258 | 259 | # Display Loss Value. 260 | print("Batch: {} & Loss: {}\r".format(fn,loss_val)), 261 | sys.stdout.flush() 262 | 263 | # Add loss for each batch. 264 | loss_sum += loss_val 265 | 266 | log_string('Train Mean loss: %f' % (loss_sum/num_batches)) # Store and display mean loss of epoch. 267 | 268 | def eval_one_epoch(sess, ops, eval_writer): 269 | # Arguments: 270 | # sess: Tensorflow session to handle tensors. 271 | # ops: Dictionary for tensors of Network_L 272 | # ops19: Dictionary for tensors of Network19 273 | # templates: Training Point Cloud data. 274 | # poses: Training pose data. 275 | 276 | is_training = False 277 | display_ptClouds = False 278 | 279 | test_file_idxs = np.arange(0, len(TEST_FILES)) 280 | np.random.shuffle(test_file_idxs) 281 | 282 | for fn in range(len(TEST_FILES)): 283 | log_string('----' + str(fn) + '-----') 284 | current_data, current_label = provider.loadDataFile(TEST_FILES[test_file_idxs[fn]]) 285 | current_data = current_data[:,0:NUM_POINT,:] 286 | current_data, _, _ = provider.shuffle_data(current_data, np.squeeze(current_label)) 287 | 288 | file_size = current_data.shape[0] 289 | num_batches = file_size // BATCH_SIZE 290 | loss_sum = 0 # Total Loss in each batch. 291 | 292 | for fn in range(num_batches): 293 | start_idx = fn*BATCH_SIZE # Start index of poses. 294 | end_idx = (fn+1)*BATCH_SIZE # End index of poses. 295 | 296 | template_data = current_data[start_idx:end_idx,:,:] 297 | 298 | # To visualize the source and point clouds: 299 | if display_ptClouds: 300 | helper.display_clouds_data(template_data[0]) 301 | 302 | # Feed the placeholders of Network_L with source data and template data obtained from N-Iterations. 303 | feed_dict = {ops['source_pointclouds_pl']: template_data, 304 | ops['is_training_pl']: is_training} 305 | 306 | # Ask the network to predict transformation, calculate loss using distance between actual points. 307 | summary, step, loss_val = sess.run([ops['merged'], ops['step'], ops['loss']], feed_dict=feed_dict) 308 | eval_writer.add_summary(summary, step) # Add all the summary to the tensorboard. 309 | 310 | # Display Loss Value. 311 | print("Batch: {} & Loss: {}\r".format(fn,loss_val)), 312 | sys.stdout.flush() 313 | 314 | # Add loss for each batch. 315 | loss_sum += loss_val 316 | 317 | log_string('Eval Mean loss: %f' % (loss_sum/num_batches)) # Store and display mean loss of epoch. 318 | 319 | def test_one_epoch(sess, ops, saver, model_path): 320 | # Arguments: 321 | # sess: Tensorflow session to handle tensors. 322 | # ops: Dictionary for tensors of Network_L 323 | # ops19: Dictionary for tensors of Network19 324 | # templates: Training Point Cloud data. 325 | # poses: Training pose data. 326 | # saver: To restore the weights. 327 | # model_path: Path of log directory. 328 | 329 | saver.restore(sess, model_path) # Restore the weights of trained network. 330 | current_data, current_label = provider.loadDataFile(TEST_FILES[0]) 331 | 332 | is_training = False 333 | display_ptClouds = False 334 | 335 | template_idx = 11 336 | current_data = current_data[template_idx*BATCH_SIZE:(template_idx+1)*BATCH_SIZE] 337 | print(current_data.shape) 338 | 339 | template_data = current_data[:,0:NUM_POINT,:] 340 | batch_euler_pose = np.array([[0,0,0,90*(np.pi/180), 0*(np.pi/180), 0*(np.pi/180)]]) 341 | # template_data = helper.apply_transformation(template_data, batch_euler_pose) 342 | batch_euler_pose = np.array([[0,0,0,0*(np.pi/180), 0*(np.pi/180), 180*(np.pi/180)]]) 343 | # template_data = helper.apply_transformation(template_data, batch_euler_pose) 344 | 345 | # To visualize the source and point clouds: 346 | if display_ptClouds: 347 | helper.display_clouds_data(template_data[0]) 348 | 349 | # Feed the placeholders of Network_L with source data and template data obtained from N-Iterations. 350 | feed_dict = {ops['source_pointclouds_pl']: template_data, 351 | ops['is_training_pl']: is_training} 352 | 353 | # Ask the network to predict transformation, calculate loss using distance between actual points. 354 | import time 355 | start = time.time() 356 | step, predicted_pointclouds_pl = sess.run([ ops['step'], ops['predicted_pointclouds_pl']], feed_dict=feed_dict) 357 | end = time.time() 358 | print('Time Elapsed: {}'.format(end-start)) 359 | predicted_pointclouds_pl[0,:,0]=predicted_pointclouds_pl[0,:,0]+1 360 | helper.display_two_clouds(template_data[0], predicted_pointclouds_pl[0]) 361 | 362 | 363 | if __name__ == "__main__": 364 | if FLAGS.mode == 'no_mode': 365 | print('Specity a mode argument: train or test') 366 | elif FLAGS.mode == 'train': 367 | train() 368 | LOG_FOUT.close() 369 | else: 370 | train() -------------------------------------------------------------------------------- /utils/tf_util.py: -------------------------------------------------------------------------------- 1 | """ Wrapper functions for TensorFlow layers. 2 | 3 | Author: Charles R. Qi 4 | Date: November 2016 5 | """ 6 | 7 | import numpy as np 8 | import tensorflow as tf 9 | 10 | 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 | with tf.device('/cpu:0'): 20 | dtype = tf.float16 if use_fp16 else tf.float32 21 | var = tf.get_variable(name, shape, initializer=initializer, dtype=dtype) 22 | return var 23 | 24 | def _variable_with_weight_decay(name, shape, stddev, wd, use_xavier=True): 25 | """Helper to create an initialized Variable with weight decay. 26 | 27 | Note that the Variable is initialized with a truncated normal distribution. 28 | A weight decay is added only if one is specified. 29 | 30 | Args: 31 | name: name of the variable 32 | shape: list of ints 33 | stddev: standard deviation of a truncated Gaussian 34 | wd: add L2Loss weight decay multiplied by this float. If None, weight 35 | decay is not added for this Variable. 36 | use_xavier: bool, whether to use xavier initializer 37 | 38 | Returns: 39 | Variable Tensor 40 | """ 41 | if use_xavier: 42 | initializer = tf.contrib.layers.xavier_initializer() 43 | else: 44 | initializer = tf.truncated_normal_initializer(stddev=stddev) 45 | var = _variable_on_cpu(name, shape, initializer) 46 | if wd is not None: 47 | weight_decay = tf.multiply(tf.nn.l2_loss(var), wd, name='weight_loss') 48 | tf.add_to_collection('losses', weight_decay) 49 | return var 50 | 51 | 52 | def conv1d(inputs, 53 | num_output_channels, 54 | kernel_size, 55 | scope, 56 | stride=1, 57 | padding='SAME', 58 | use_xavier=True, 59 | stddev=1e-3, 60 | weight_decay=0.0, 61 | activation_fn=tf.nn.relu, 62 | bn=False, 63 | bn_decay=None, 64 | is_training=None): 65 | """ 1D convolution with non-linear operation. 66 | 67 | Args: 68 | inputs: 3-D tensor variable BxLxC 69 | num_output_channels: int 70 | kernel_size: int 71 | scope: string 72 | stride: int 73 | padding: 'SAME' or 'VALID' 74 | use_xavier: bool, use xavier_initializer if true 75 | stddev: float, stddev for truncated_normal init 76 | weight_decay: float 77 | activation_fn: function 78 | bn: bool, whether to use batch norm 79 | bn_decay: float or float tensor variable in [0,1] 80 | is_training: bool Tensor variable 81 | 82 | Returns: 83 | Variable tensor 84 | """ 85 | with tf.variable_scope(scope) as sc: 86 | num_in_channels = inputs.get_shape()[-1].value 87 | kernel_shape = [kernel_size, 88 | num_in_channels, num_output_channels] 89 | kernel = _variable_with_weight_decay('weights', 90 | shape=kernel_shape, 91 | use_xavier=use_xavier, 92 | stddev=stddev, 93 | wd=weight_decay) 94 | outputs = tf.nn.conv1d(inputs, kernel, 95 | stride=stride, 96 | padding=padding) 97 | biases = _variable_on_cpu('biases', [num_output_channels], 98 | tf.constant_initializer(0.0)) 99 | outputs = tf.nn.bias_add(outputs, biases) 100 | 101 | if bn: 102 | outputs = batch_norm_for_conv1d(outputs, is_training, 103 | bn_decay=bn_decay, scope='bn') 104 | 105 | if activation_fn is not None: 106 | outputs = activation_fn(outputs) 107 | return outputs 108 | 109 | 110 | 111 | 112 | def conv2d(inputs, 113 | num_output_channels, 114 | kernel_size, 115 | scope, 116 | stride=[1, 1], 117 | padding='SAME', 118 | use_xavier=True, 119 | stddev=1e-3, 120 | weight_decay=0.0, 121 | activation_fn=tf.nn.relu, 122 | bn=False, 123 | bn_decay=None, 124 | is_training=None): 125 | """ 2D convolution with non-linear operation. 126 | 127 | Args: 128 | inputs: 4-D tensor variable BxHxWxC 129 | num_output_channels: int 130 | kernel_size: a list of 2 ints 131 | scope: string 132 | stride: a list of 2 ints 133 | padding: 'SAME' or 'VALID' 134 | use_xavier: bool, use xavier_initializer if true 135 | stddev: float, stddev for truncated_normal init 136 | weight_decay: float 137 | activation_fn: function 138 | bn: bool, whether to use batch norm 139 | bn_decay: float or float tensor variable in [0,1] 140 | is_training: bool Tensor variable 141 | 142 | Returns: 143 | Variable tensor 144 | """ 145 | with tf.variable_scope(scope) as sc: 146 | kernel_h, kernel_w = kernel_size 147 | num_in_channels = inputs.get_shape()[-1].value 148 | kernel_shape = [kernel_h, kernel_w, 149 | num_in_channels, num_output_channels] 150 | kernel = _variable_with_weight_decay('weights', 151 | shape=kernel_shape, 152 | use_xavier=use_xavier, 153 | stddev=stddev, 154 | wd=weight_decay) 155 | stride_h, stride_w = stride 156 | outputs = tf.nn.conv2d(inputs, kernel, 157 | [1, stride_h, stride_w, 1], 158 | padding=padding) 159 | biases = _variable_on_cpu('biases', [num_output_channels], 160 | tf.constant_initializer(0.0)) 161 | outputs = tf.nn.bias_add(outputs, biases) 162 | 163 | if bn: 164 | pass 165 | # outputs = batch_norm_for_conv2d(outputs, is_training, 166 | # bn_decay=bn_decay, scope='bn') 167 | # outputs = batch_norm_for_conv2d(outputs, is_training, 168 | # bn_decay=bn_decay, scope=scope) 169 | 170 | if activation_fn is not None: 171 | outputs = activation_fn(outputs) 172 | return outputs 173 | 174 | 175 | def conv2d_transpose(inputs, 176 | num_output_channels, 177 | kernel_size, 178 | scope, 179 | stride=[1, 1], 180 | padding='SAME', 181 | use_xavier=True, 182 | stddev=1e-3, 183 | weight_decay=0.0, 184 | activation_fn=tf.nn.relu, 185 | bn=False, 186 | bn_decay=None, 187 | is_training=None): 188 | """ 2D convolution transpose with non-linear operation. 189 | 190 | Args: 191 | inputs: 4-D tensor variable BxHxWxC 192 | num_output_channels: int 193 | kernel_size: a list of 2 ints 194 | scope: string 195 | stride: a list of 2 ints 196 | padding: 'SAME' or 'VALID' 197 | use_xavier: bool, use xavier_initializer if true 198 | stddev: float, stddev for truncated_normal init 199 | weight_decay: float 200 | activation_fn: function 201 | bn: bool, whether to use batch norm 202 | bn_decay: float or float tensor variable in [0,1] 203 | is_training: bool Tensor variable 204 | 205 | Returns: 206 | Variable tensor 207 | 208 | Note: conv2d(conv2d_transpose(a, num_out, ksize, stride), a.shape[-1], ksize, stride) == a 209 | """ 210 | with tf.variable_scope(scope) as sc: 211 | kernel_h, kernel_w = kernel_size 212 | num_in_channels = inputs.get_shape()[-1].value 213 | kernel_shape = [kernel_h, kernel_w, 214 | num_output_channels, num_in_channels] # reversed to conv2d 215 | kernel = _variable_with_weight_decay('weights', 216 | shape=kernel_shape, 217 | use_xavier=use_xavier, 218 | stddev=stddev, 219 | wd=weight_decay) 220 | stride_h, stride_w = stride 221 | 222 | # from slim.convolution2d_transpose 223 | def get_deconv_dim(dim_size, stride_size, kernel_size, padding): 224 | dim_size *= stride_size 225 | 226 | if padding == 'VALID' and dim_size is not None: 227 | dim_size += max(kernel_size - stride_size, 0) 228 | return dim_size 229 | 230 | # caculate output shape 231 | batch_size = inputs.get_shape()[0].value 232 | height = inputs.get_shape()[1].value 233 | width = inputs.get_shape()[2].value 234 | out_height = get_deconv_dim(height, stride_h, kernel_h, padding) 235 | out_width = get_deconv_dim(width, stride_w, kernel_w, padding) 236 | output_shape = [batch_size, out_height, out_width, num_output_channels] 237 | 238 | outputs = tf.nn.conv2d_transpose(inputs, kernel, output_shape, 239 | [1, stride_h, stride_w, 1], 240 | padding=padding) 241 | biases = _variable_on_cpu('biases', [num_output_channels], 242 | tf.constant_initializer(0.0)) 243 | outputs = tf.nn.bias_add(outputs, biases) 244 | 245 | if bn: 246 | outputs = batch_norm_for_conv2d(outputs, is_training, 247 | bn_decay=bn_decay, scope='bn') 248 | 249 | if activation_fn is not None: 250 | outputs = activation_fn(outputs) 251 | return outputs 252 | 253 | 254 | 255 | def conv3d(inputs, 256 | num_output_channels, 257 | kernel_size, 258 | scope, 259 | stride=[1, 1, 1], 260 | padding='SAME', 261 | use_xavier=True, 262 | stddev=1e-3, 263 | weight_decay=0.0, 264 | activation_fn=tf.nn.relu, 265 | bn=False, 266 | bn_decay=None, 267 | is_training=None): 268 | """ 3D convolution with non-linear operation. 269 | 270 | Args: 271 | inputs: 5-D tensor variable BxDxHxWxC 272 | num_output_channels: int 273 | kernel_size: a list of 3 ints 274 | scope: string 275 | stride: a list of 3 ints 276 | padding: 'SAME' or 'VALID' 277 | use_xavier: bool, use xavier_initializer if true 278 | stddev: float, stddev for truncated_normal init 279 | weight_decay: float 280 | activation_fn: function 281 | bn: bool, whether to use batch norm 282 | bn_decay: float or float tensor variable in [0,1] 283 | is_training: bool Tensor variable 284 | 285 | Returns: 286 | Variable tensor 287 | """ 288 | with tf.variable_scope(scope) as sc: 289 | kernel_d, kernel_h, kernel_w = kernel_size 290 | num_in_channels = inputs.get_shape()[-1].value 291 | kernel_shape = [kernel_d, kernel_h, kernel_w, 292 | num_in_channels, num_output_channels] 293 | kernel = _variable_with_weight_decay('weights', 294 | shape=kernel_shape, 295 | use_xavier=use_xavier, 296 | stddev=stddev, 297 | wd=weight_decay) 298 | stride_d, stride_h, stride_w = stride 299 | outputs = tf.nn.conv3d(inputs, kernel, 300 | [1, stride_d, stride_h, stride_w, 1], 301 | padding=padding) 302 | biases = _variable_on_cpu('biases', [num_output_channels], 303 | tf.constant_initializer(0.0)) 304 | outputs = tf.nn.bias_add(outputs, biases) 305 | 306 | if bn: 307 | outputs = batch_norm_for_conv3d(outputs, is_training, 308 | bn_decay=bn_decay, scope='bn') 309 | 310 | if activation_fn is not None: 311 | outputs = activation_fn(outputs) 312 | return outputs 313 | 314 | def fully_connected(inputs, 315 | num_outputs, 316 | scope, 317 | use_xavier=True, 318 | stddev=1e-3, 319 | weight_decay=0.0, 320 | activation_fn=tf.nn.relu, 321 | bn=False, 322 | bn_decay=None, 323 | is_training=None): 324 | """ Fully connected layer with non-linear operation. 325 | 326 | Args: 327 | inputs: 2-D tensor BxN 328 | num_outputs: int 329 | 330 | Returns: 331 | Variable tensor of size B x num_outputs. 332 | """ 333 | with tf.variable_scope(scope ) as sc: 334 | num_input_units = inputs.get_shape()[-1].value 335 | weights = _variable_with_weight_decay('weights', 336 | shape=[num_input_units, num_outputs], 337 | use_xavier=use_xavier, 338 | stddev=stddev, 339 | wd=weight_decay) 340 | outputs = tf.matmul(inputs, weights) 341 | biases = _variable_on_cpu('biases', [num_outputs], 342 | tf.constant_initializer(0.0)) 343 | outputs = tf.nn.bias_add(outputs, biases) 344 | 345 | if bn: 346 | pass 347 | # outputs = batch_norm_for_fc(outputs, is_training, bn_decay, 'bn') 348 | # outputs = batch_norm_for_fc(outputs, is_training, bn_decay, scope) 349 | 350 | if activation_fn is not None: 351 | outputs = activation_fn(outputs) 352 | return outputs 353 | 354 | 355 | def max_pool2d(inputs, 356 | kernel_size, 357 | scope, 358 | stride=[2, 2], 359 | padding='VALID'): 360 | """ 2D max pooling. 361 | 362 | Args: 363 | inputs: 4-D tensor BxHxWxC 364 | kernel_size: a list of 2 ints 365 | stride: a list of 2 ints 366 | 367 | Returns: 368 | Variable tensor 369 | """ 370 | with tf.variable_scope(scope ) as sc: 371 | kernel_h, kernel_w = kernel_size 372 | stride_h, stride_w = stride 373 | outputs = tf.nn.max_pool(inputs, 374 | ksize=[1, kernel_h, kernel_w, 1], 375 | strides=[1, stride_h, stride_w, 1], 376 | padding=padding, 377 | name=sc.name) 378 | return outputs 379 | 380 | def avg_pool2d(inputs, 381 | kernel_size, 382 | scope, 383 | stride=[2, 2], 384 | padding='VALID'): 385 | """ 2D avg pooling. 386 | 387 | Args: 388 | inputs: 4-D tensor BxHxWxC 389 | kernel_size: a list of 2 ints 390 | stride: a list of 2 ints 391 | 392 | Returns: 393 | Variable tensor 394 | """ 395 | with tf.variable_scope(scope) as sc: 396 | kernel_h, kernel_w = kernel_size 397 | stride_h, stride_w = stride 398 | outputs = tf.nn.avg_pool(inputs, 399 | ksize=[1, kernel_h, kernel_w, 1], 400 | strides=[1, stride_h, stride_w, 1], 401 | padding=padding, 402 | name=sc.name) 403 | return outputs 404 | 405 | 406 | def max_pool3d(inputs, 407 | kernel_size, 408 | scope, 409 | stride=[2, 2, 2], 410 | padding='VALID'): 411 | """ 3D max pooling. 412 | 413 | Args: 414 | inputs: 5-D tensor BxDxHxWxC 415 | kernel_size: a list of 3 ints 416 | stride: a list of 3 ints 417 | 418 | Returns: 419 | Variable tensor 420 | """ 421 | with tf.variable_scope(scope) as sc: 422 | kernel_d, kernel_h, kernel_w = kernel_size 423 | stride_d, stride_h, stride_w = stride 424 | outputs = tf.nn.max_pool3d(inputs, 425 | ksize=[1, kernel_d, kernel_h, kernel_w, 1], 426 | strides=[1, stride_d, stride_h, stride_w, 1], 427 | padding=padding, 428 | name=sc.name) 429 | return outputs 430 | 431 | def avg_pool3d(inputs, 432 | kernel_size, 433 | scope, 434 | stride=[2, 2, 2], 435 | padding='VALID'): 436 | """ 3D avg pooling. 437 | 438 | Args: 439 | inputs: 5-D tensor BxDxHxWxC 440 | kernel_size: a list of 3 ints 441 | stride: a list of 3 ints 442 | 443 | Returns: 444 | Variable tensor 445 | """ 446 | with tf.variable_scope(scope) as sc: 447 | kernel_d, kernel_h, kernel_w = kernel_size 448 | stride_d, stride_h, stride_w = stride 449 | outputs = tf.nn.avg_pool3d(inputs, 450 | ksize=[1, kernel_d, kernel_h, kernel_w, 1], 451 | strides=[1, stride_d, stride_h, stride_w, 1], 452 | padding=padding, 453 | name=sc.name) 454 | return outputs 455 | 456 | 457 | 458 | 459 | 460 | def batch_norm_template(inputs, is_training, scope, moments_dims, bn_decay): 461 | """ Batch normalization on convolutional maps and beyond... 462 | Ref.: http://stackoverflow.com/questions/33949786/how-could-i-use-batch-normalization-in-tensorflow 463 | 464 | Args: 465 | inputs: Tensor, k-D input ... x C could be BC or BHWC or BDHWC 466 | is_training: boolean tf.Varialbe, true indicates training phase 467 | scope: string, variable scope 468 | moments_dims: a list of ints, indicating dimensions for moments calculation 469 | bn_decay: float or float tensor variable, controling moving average weight 470 | Return: 471 | normed: batch-normalized maps 472 | """ 473 | with tf.variable_scope(scope ) as sc: 474 | num_channels = inputs.get_shape()[-1].value 475 | beta = tf.get_variable('beta',initializer=tf.constant(0.0, shape=[num_channels]), 476 | trainable=True) 477 | gamma = tf.get_variable('gamma',initializer=tf.constant(1.0, shape=[num_channels]), 478 | trainable=True) 479 | # beta = tf.constant(0.0, shape=[num_channels]) 480 | # gamma = tf.constant(1.0, shape=[num_channels]) 481 | batch_mean, batch_var = tf.nn.moments(inputs, moments_dims, name='moments') 482 | decay = bn_decay if bn_decay is not None else 0.9 483 | ema = tf.train.ExponentialMovingAverage(decay=decay) 484 | # Operator that maintains moving averages of variables. 485 | ema_apply_op = tf.cond(is_training, 486 | lambda: ema.apply([batch_mean, batch_var]), 487 | lambda: tf.no_op()) 488 | 489 | # Update moving average and return current batch's avg and var. 490 | def mean_var_with_update(): 491 | with tf.control_dependencies([ema_apply_op]): 492 | return tf.identity(batch_mean), tf.identity(batch_var) 493 | 494 | # ema.average returns the Variable holding the average of var. 495 | mean, var = tf.cond(is_training, 496 | mean_var_with_update, 497 | lambda: (ema.average(batch_mean), ema.average(batch_var))) 498 | normed = tf.nn.batch_normalization(inputs, mean, var, beta, gamma, 1e-3) 499 | return normed 500 | 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 | -------------------------------------------------------------------------------- /helper.py: -------------------------------------------------------------------------------- 1 | import csv 2 | import matplotlib.pyplot as plt 3 | from mpl_toolkits.mplot3d import Axes3D 4 | import os 5 | import numpy as np 6 | import transforms3d.euler as t3d 7 | import helper 8 | import tensorflow as tf 9 | 10 | 11 | ###################### Data Handling Operations ######################### 12 | 13 | # Read the templates from a given file. 14 | def read_templates(file_name,templates_dict): 15 | with open(os.path.join('data',templates_dict,file_name),'r') as csvfile: 16 | csvreader = csv.reader(csvfile) 17 | data = [] 18 | for row in csvreader: 19 | row = [float(i) for i in row] 20 | data.append(row) 21 | return data # n2 x 2048 x 3 22 | 23 | # Read the file names having templates. 24 | def template_files(templates_dict): 25 | with open(os.path.join('data',templates_dict,'template_filenames.txt'),'r') as file: 26 | files = file.readlines() 27 | files = [x.strip() for x in files] 28 | print('Templates used to train data: ') 29 | print(files) 30 | return files # 1 x n1 31 | 32 | # Read the templates from each file. 33 | def templates_data(templates_dict): 34 | files = template_files(templates_dict) # Read the available file names. 35 | data = [] 36 | for i in range(len(files)): 37 | temp = read_templates(files[i],templates_dict) 38 | for i in temp: 39 | data.append(i) 40 | return np.asarray(data) # (n1 x n2 x 2048 x 3) & n = n1 x n2 41 | 42 | # Preprocess the templates and rearrange them. 43 | def process_templates(templates_dict): 44 | data = templates_data(templates_dict) # Read all the templates. 45 | print('No. of Total Templates: {}'.format(data.shape[0]/2048)) 46 | templates = [] 47 | for i in range(data.shape[0]/2048): 48 | start_idx = i*2048 49 | end_idx = (i+1)*2048 50 | templates.append(data[start_idx:end_idx,:]) 51 | return np.asarray(templates) # Return all the templates (n x 2048 x 3) 52 | 53 | # Read poses from given file. 54 | def read_poses(templates_dict, filename): 55 | # Arguments: 56 | # filename: Read data from a given file (string) 57 | # Output: 58 | # poses: Return array of all the poses in the file (n x 6) 59 | 60 | with open(os.path.join('data',templates_dict,filename),'r') as csvfile: 61 | csvreader = csv.reader(csvfile) 62 | poses = [] 63 | for row in csvreader: 64 | row = [float(i) for i in row] 65 | poses.append(row) 66 | return np.asarray(poses) 67 | 68 | # To Store the features obtained from network. 69 | def store_features(filename,feature): 70 | # Arguments: 71 | # filename: name of file to store features. 72 | # feature: Array of feature to store in the file. 73 | with open(filename,'w') as file: 74 | for i in range(1024): 75 | file.write(str(feature[i])) 76 | file.write(',') 77 | 78 | 79 | ###################### Transformation Operations ######################### 80 | 81 | def rotate_point_cloud_by_angle_y(batch_data, rotation_angle): 82 | """ Rotate the point cloud along up direction with certain angle. 83 | Input: 84 | BxNx3 array, original batch of point clouds 85 | Return: 86 | BxNx3 array, rotated batch of point clouds 87 | """ 88 | rotated_data = np.zeros(batch_data.shape, dtype=np.float32) 89 | for k in range(batch_data.shape[0]): 90 | #rotation_angle = np.random.uniform() * 2 * np.pi 91 | cosval = np.cos(rotation_angle) 92 | sinval = np.sin(rotation_angle) 93 | rotation_matrix = np.array([[cosval, 0, sinval], 94 | [0, 1, 0], 95 | [-sinval, 0, cosval]]) 96 | shape_pc = batch_data[k, ...] 97 | # rotated_data[k, ...] = np.dot(shape_pc.reshape((-1, 3)), rotation_matrix) 98 | rotated_data[k, ...] = np.dot(rotation_matrix, shape_pc.reshape((-1, 3)).T).T # Pre-Multiplication (changes done) 99 | return rotated_data 100 | 101 | def rotate_point_cloud_by_angle_x(batch_data, rotation_angle): 102 | """ Rotate the point cloud along up direction with certain angle. 103 | Input: 104 | BxNx3 array, original batch of point clouds 105 | Return: 106 | BxNx3 array, rotated batch of point clouds 107 | """ 108 | rotated_data = np.zeros(batch_data.shape, dtype=np.float32) 109 | for k in range(batch_data.shape[0]): 110 | #rotation_angle = np.random.uniform() * 2 * np.pi 111 | cosval = np.cos(rotation_angle) 112 | sinval = np.sin(rotation_angle) 113 | rotation_matrix = np.array([[1, 0, 0], 114 | [0, cosval, -sinval], 115 | [0, sinval, cosval]]) 116 | shape_pc = batch_data[k, ...] 117 | # rotated_data[k, ...] = np.dot(shape_pc.reshape((-1, 3)), rotation_matrix) 118 | rotated_data[k, ...] = np.dot(rotation_matrix, shape_pc.reshape((-1, 3)).T).T # Pre-Multiplication (changes done) 119 | return rotated_data 120 | 121 | def rotate_point_cloud_by_angle_z(batch_data, rotation_angle): 122 | """ Rotate the point cloud along up direction with certain angle. 123 | Input: 124 | BxNx3 array, original batch of point clouds 125 | Return: 126 | BxNx3 array, rotated batch of point clouds 127 | """ 128 | rotated_data = np.zeros(batch_data.shape, dtype=np.float32) 129 | for k in range(batch_data.shape[0]): 130 | #rotation_angle = np.random.uniform() * 2 * np.pi 131 | cosval = np.cos(rotation_angle) 132 | sinval = np.sin(rotation_angle) 133 | rotation_matrix = np.array([[cosval, -sinval, 0], 134 | [sinval, cosval, 0], 135 | [0, 0, 1]]) 136 | shape_pc = batch_data[k, ...] 137 | # rotated_data[k, ...] = np.dot(shape_pc.reshape((-1, 3)), rotation_matrix) 138 | rotated_data[k, ...] = np.dot(rotation_matrix, shape_pc.reshape((-1, 3)).T).T # Pre-Multiplication (changes done) 139 | return rotated_data 140 | 141 | # Translate the data as per given translation vector. 142 | def translate(data,shift): 143 | # Arguments: 144 | # data: Point Cloud data (1 x num_points x 3) 145 | # shift: Translation vector (1 x 3) 146 | 147 | try: 148 | data = np.asarray(data) 149 | except: 150 | pass 151 | return data+shift 152 | 153 | # Apply the given transformation to given point cloud data. 154 | def apply_transformation(datas,poses): # Transformation function for (2 & 4c, loss 8b) 155 | # Arguments: 156 | # datas: Point Clouds (batch_size x num_points x 3) 157 | # poses: translation+euler (Batch_size x 6) 158 | # Output: 159 | # transformed_data: Transformed Point Clouds by given poses (batch_size x num_points x 3) 160 | transformed_data = np.copy(datas) 161 | for i in range(datas.shape[0]): 162 | transformed_data[i,:,:] = rotate_point_cloud_by_angle_z(transformed_data[i,:,:],poses[i,5]) 163 | transformed_data[i,:,:] = rotate_point_cloud_by_angle_y(transformed_data[i,:,:],poses[i,4]) 164 | transformed_data[i,:,:] = rotate_point_cloud_by_angle_x(transformed_data[i,:,:],poses[i,3]) 165 | transformed_data[i,:,:] = translate(transformed_data[i,:,:],[poses[i,0],poses[i,1],poses[i,2]]) 166 | return transformed_data 167 | 168 | # Convert poses from 6D to 7D. # For loss function ( 8a ) 169 | def poses_euler2quat(poses): 170 | # Arguments: 171 | # poses: 6D pose (translation + euler) (batch_size x 6) 172 | # Output: 173 | # new_poses: 7D pose (translation + quaternions) (batch_size x 7) 174 | 175 | new_poses = [] # Store 7D poses 176 | for i in range(poses.shape[0]): 177 | temp = t3d.euler2quat(poses[i,3],poses[i,4],poses[i,5]) # Convert from euler to quaternion. (1x4) 178 | temp1 = [poses[i,0],poses[i,1],poses[i,2],temp[0],temp[1],temp[2],temp[3]] # Add translation & Quaternion (1x7) 179 | new_poses.append(temp1) 180 | return np.asarray(new_poses) 181 | 182 | # Geenerate random poses equal to batch_size. 183 | def generate_poses(batch_size): 184 | # Arguments: 185 | # batch_size: No of 6D poses required. 186 | # Output: 187 | # poses: Array of poses with translation and rotation (euler angles in radians) (batch_size x 6) 188 | 189 | poses = [] # List to store the 6D poses. 190 | for i in range(batch_size): 191 | # Generate random translations. 192 | x = np.round(2*np.random.random_sample()-1,2) 193 | y = np.round(2*np.random.random_sample()-1,2) 194 | z = np.round(2*np.random.random_sample()-1,2) 195 | # Generate random rotations. 196 | x_rot = np.round(np.pi*np.random.random_sample()-(np.pi/2),3) 197 | y_rot = np.round(np.pi*np.random.random_sample()-(np.pi/2),3) 198 | z_rot = np.round(np.pi*np.random.random_sample()-(np.pi/2),3) 199 | poses.append([x,y,z,x_rot,y_rot,z_rot]) 200 | return np.array(poses).reshape((batch_size,6)) 201 | 202 | # Convert 6D poses to transformation matrix. # (for 4b) 203 | def transformation(poses): 204 | # Arguments: 205 | # poses: 6D (x,y,z,euler_x,euler_y,euler_z) (in radians) 206 | # Output 207 | # transformation_matrix: batch_size x 4 x 4 208 | 209 | transformation_matrix = np.zeros((poses.shape[0],4,4)) 210 | transformation_matrix[:,3,3] = 1 211 | for i in range(poses.shape[0]): 212 | rot = t3d.euler2mat(poses[i,5],poses[i,4],poses[i,3],'szyx') # Calculate rotation matrix using transforms3d 213 | transformation_matrix[i,0:3,0:3]=rot # Store rotation matrix in transformation matrix. 214 | transformation_matrix[i,0:3,3]=poses[i,0:3] # Store translations in transformation matrix. 215 | return transformation_matrix 216 | 217 | # Convert poses (quaternions) to transformation matrix and apply on point cloud. 218 | def transformation_quat2mat(poses,TRANSFORMATIONS,templates_data): # (for 4b) 219 | # Arguments: 220 | # poses: 7D (x,y,z,quat_q0,quat_q1,quat_q2,quat_q3) (in radians) (batch_size x 7) 221 | # TRANSFORMATIONS: Overall tranformation matrix. 222 | # template_data: Point Cloud (batch_size x num_points x 3) 223 | # Output 224 | # TRANSFORMATIONS: Batch_size x 4 x 4 225 | # templates_data: Transformed template data (batch_size x num_points x 3) 226 | 227 | poses = np.array(poses) # Convert poses to array. 228 | poses = poses.reshape(poses.shape[-2],poses.shape[-1]) 229 | for i in range(poses.shape[0]): 230 | transformation_matrix = np.zeros((4,4)) 231 | transformation_matrix[3,3] = 1 232 | rot = t3d.quat2mat([poses[i,3],poses[i,4],poses[i,5],poses[i,6]]) # Calculate rotation matrix using transforms3d (library handles the normalization part of quaternion) 233 | transformation_matrix[0:3,0:3]=rot # Store rotation matrix in transformation matrix. 234 | transformation_matrix[0:3,3]=poses[i,0:3] # Store translations in transformation matrix. 235 | TRANSFORMATIONS[i,:,:] = np.dot(transformation_matrix,TRANSFORMATIONS[i,:,:]) # 4b (Multiply tranfromation matrix to Initial Transfromation Matrix) 236 | templates_data[i,:,:]=np.dot(rot,templates_data[i,:,:].T).T # Apply Rotation to Template Data 237 | templates_data[i,:,:]=templates_data[i,:,:]+poses[i,0:3] # Apply translation to template data 238 | return TRANSFORMATIONS,templates_data 239 | 240 | # Convert the Final Transformation Matrix to Translation + Orientation (Euler Angles in Degrees) 241 | def find_final_pose(TRANSFORMATIONS): 242 | # Arguments: 243 | # TRANSFORMATIONS: transformation matrix (batch_size x 4 x 4) 244 | # Output: 245 | # final_pose: final pose predicted by network (batch_size x 6) 246 | 247 | final_pose = np.zeros((TRANSFORMATIONS.shape[0],6)) # Array to store the poses. 248 | for i in range(TRANSFORMATIONS.shape[0]): 249 | rot = TRANSFORMATIONS[i,0:3,0:3] # Extract rotation matrix. 250 | euler = t3d.mat2euler(rot,'szyx') # Convert rotation matrix to euler angles. (Pre-multiplication) 251 | final_pose[i,3:6]=[euler[2],euler[1],euler[0]] # Store the translation 252 | final_pose[i,0:3]=TRANSFORMATIONS[i,0:3,3].T # Store the euler angles. 253 | return final_pose 254 | 255 | # Subtract the centroids from source and template (Like ICP) and then find the pose. 256 | def centroid_subtraction(source_data, template_data): 257 | # Arguments: 258 | # source_data: Source Point Clouds (batch_size x num_points x 3) 259 | # template_data: Template Point Clouds (batch_size x num_points x 3) 260 | # Output: 261 | # source_data: Centroid subtracted from source point cloud (batch_size x num_points x 3) 262 | # template_data: Centroid subtracted from template point cloud (batch_size x num_points x 3) 263 | # centroid_translation_pose: Apply this pose after final iteration. (batch_size x 7) 264 | 265 | centroid_translation_pose = np.zeros((source_data.shape[0],7)) 266 | for i in range(source_data.shape[0]): 267 | source_centroid = np.mean(source_data[i],axis=0) 268 | template_centroid = np.mean(template_data[i],axis=0) 269 | source_data[i] = source_data[i] - source_centroid 270 | template_data[i] = template_data[i] - template_centroid 271 | centroid_translation = source_centroid - template_centroid 272 | centroid_translation_pose[i] = np.array([centroid_translation[0],centroid_translation[1],centroid_translation[2],1,0,0,0]) 273 | return source_data, template_data, centroid_translation_pose 274 | 275 | 276 | ###################### Shuffling Operations ######################### 277 | 278 | # Randomly shuffle given array of poses for training procedure. 279 | def shuffle_templates(templates): 280 | # Arguments: 281 | # templates: Input array of templates to get randomly shuffled (batch_size x num_points x 3) 282 | # Output: 283 | # shuffled_templates: Randomly ordered poses (batch_size x num_points x 3) 284 | 285 | shuffled_templates = np.zeros(templates.shape) # Array to store shuffled templates. 286 | templates_idxs = np.arange(0,templates.shape[0]) 287 | np.random.shuffle(templates_idxs) # Randomly shuffle template indices. 288 | for i in range(templates.shape[0]): 289 | shuffled_templates[i,:,:]=templates[templates_idxs[i],:,:] # Rearrange them as per shuffled indices. 290 | return shuffled_templates 291 | 292 | # Randomly shuffle given array of poses for training procedure. 293 | def shuffle_poses(poses): 294 | # Arguments: 295 | # poses: Input array of poses to get randomly shuffled (batch_size x n) 296 | # Output: 297 | # shuffled_poses: Randomly ordered poses (batch_size x n) 298 | 299 | shuffled_poses = np.zeros(poses.shape) # Array to store shuffled poses. 300 | poses_idxs = np.arange(0,poses.shape[0]) 301 | np.random.shuffle(poses_idxs) # Shuffle the indexes of poses. 302 | for i in range(poses.shape[0]): 303 | shuffled_poses[i,:]=poses[poses_idxs[i],:] # Rearrange them as per shuffled indexes. 304 | return shuffled_poses 305 | 306 | # Generate random transformation/pose for data augmentation. 307 | def random_trans(): 308 | # Output: 309 | # 6D pose with first 3 translation values and last 3 euler angles in radian about x,y,z-axes. (1x6) 310 | 311 | # Generate random translations. 312 | x_trans, y_trans, z_trans = 0.4*np.random.uniform()-0.2, 0.4*np.random.uniform()-0.2, 0.4*np.random.uniform()-0.2 313 | # Generate random rotation angles. 314 | x_rot, y_rot, z_rot = (np.pi/9)*np.random.uniform()-(np.pi/18), (np.pi/9)*np.random.uniform()-(np.pi/18), (np.pi/9)*np.random.uniform()-(np.pi/18) 315 | return [x_trans,y_trans,z_trans,x_rot,y_rot,z_rot] 316 | 317 | # Generate random poses for each batch to train the network. 318 | def generate_random_poses(batch_size): 319 | # Arguments: 320 | # Batch_size: No of poses in the output 321 | # Output: 322 | # poses: Randomly generated poses (batch_size x 6) 323 | 324 | poses = [] 325 | for i in range(batch_size): 326 | x_trans, y_trans, z_trans = 2*np.random.uniform()-1, 2*np.random.uniform()-1, 2*np.random.uniform()-1 # Generate random translation 327 | x_rot, y_rot, z_rot = (np.pi)*np.random.uniform()-(np.pi/2), (np.pi)*np.random.uniform()-(np.pi/2), (np.pi)*np.random.uniform()-(np.pi/2) # Generate random orientation 328 | poses.append([np.round(x_trans,4), np.round(y_trans,4), np.round(z_trans,4), np.round(x_rot,4), np.round(y_rot,4), np.round(z_rot,4)]) # round upto 4 decimal digits 329 | return np.array(poses) 330 | 331 | ###################### Tensor Operations ######################### 332 | 333 | def rotate_point_cloud_by_angle_y_tensor(data, rotation_angle): 334 | """ Rotate the point cloud along up direction with certain angle. 335 | Input: 336 | Nx3 array, original batch of point clouds 337 | Return: 338 | Nx3 array, rotated batch of point clouds 339 | """ 340 | cosval = tf.cos(rotation_angle) 341 | sinval = tf.sin(rotation_angle) 342 | rotation_matrix = tf.reshape([[cosval, 0, sinval],[0, 1, 0],[-sinval, 0, cosval]], [3,3]) 343 | data = tf.reshape(data, [-1, 3]) 344 | rotated_data = tf.transpose(tf.tensordot(rotation_matrix, tf.transpose(data), [1,0])) 345 | return rotated_data 346 | 347 | def rotate_point_cloud_by_angle_x_tensor(data, rotation_angle): 348 | """ Rotate the point cloud along up direction with certain angle. 349 | Input: 350 | Nx3 array, original batch of point clouds 351 | Return: 352 | Nx3 array, rotated batch of point clouds 353 | """ 354 | cosval = tf.cos(rotation_angle) 355 | sinval = tf.sin(rotation_angle) 356 | rotation_matrix = tf.reshape([[1, 0, 0],[0, cosval, -sinval],[0, sinval, cosval]], [3,3]) 357 | data = tf.reshape(data, [-1, 3]) 358 | rotated_data = tf.transpose(tf.tensordot(rotation_matrix, tf.transpose(data), [1,0])) 359 | return rotated_data 360 | 361 | def rotate_point_cloud_by_angle_z_tensor(data, rotation_angle): 362 | """ Rotate the point cloud along up direction with certain angle. 363 | Input: 364 | Nx3 array, original batch of point clouds 365 | Return: 366 | Nx3 array, rotated batch of point clouds 367 | """ 368 | cosval = tf.cos(rotation_angle) 369 | sinval = tf.sin(rotation_angle) 370 | rotation_matrix = tf.reshape([[cosval, -sinval, 0],[sinval, cosval, 0],[0, 0, 1]], [3,3]) 371 | data = tf.reshape(data, [-1, 3]) 372 | rotated_data = tf.transpose(tf.tensordot(rotation_matrix, tf.transpose(data), [1,0])) 373 | return rotated_data 374 | 375 | def translate_tensor(data,shift): 376 | # Add the translation vector to given tensor. (num_point x 3) 377 | return tf.add(data,shift) 378 | 379 | # Tranform the data as per given poses with orientation as euler in degrees. 380 | def transformation_tensor(datas,poses): 381 | # Arguments: 382 | # datas: Tensor of Point Cloud (batch_size x num_points x 3) 383 | # poses: Tensor of Poses (translation + euler angles in degrees) (batch_size x num_points x 3) 384 | # Ouput: 385 | # transformed_data: Tensor of transformed point cloud (batch_size x num_points x 3) 386 | 387 | transformed_data = tf.zeros([datas.shape[1], datas.shape[2]]) # Tensor to store the transformed point clouds as tensor. 388 | for i in range(datas.shape[0]): 389 | transformed_data_t = rotate_point_cloud_by_angle_x_tensor(datas[i,...],poses[i,3]) # Rotate about x-axis 390 | transformed_data_t = rotate_point_cloud_by_angle_y_tensor(transformed_data_t,poses[i,4]) # Rotate about y-axis 391 | transformed_data_t = rotate_point_cloud_by_angle_z_tensor(transformed_data_t,poses[i,5]) # Rotate about z-axis 392 | transformed_data_t = translate_tensor(transformed_data_t,[poses[i,0],poses[i,1],poses[i,2]]) # Translate by given vector. 393 | transformed_data = tf.concat([transformed_data, transformed_data_t], 0) # Append the transformed tensor point cloud. 394 | transformed_data = tf.reshape(transformed_data, [-1, datas.shape[1], datas.shape[2]])[1:] # Reshape the transformed tensor and remove first one. (batch_size x num_point x 3) 395 | return transformed_data 396 | 397 | # Tranform the data as per given poses with orientation as quaternion. 398 | def transformation_quat_tensor(data,quat,translation): 399 | # Arguments: 400 | # data: Tensor of Point Cloud. (batch_size x num_point x 3) 401 | # quat: Quaternion tensor to generate rotation matrix. (batch_size x 4) 402 | # translation: Translation tensor to translate the point cloud. (batch_size x 3) 403 | # Outputs: 404 | # transformed_data: Tensor of Rotated and Translated Point Cloud Data. (batch_size x num_points x 3) 405 | 406 | transformed_data = tf.zeros([data.shape[1],3]) # Tensor to store transformed data. 407 | for i in range(quat.shape[0]): 408 | # Seperate each quaternion value. 409 | q0 = tf.slice(quat,[i,0],[1,1]) 410 | q1 = tf.slice(quat,[i,1],[1,1]) 411 | q2 = tf.slice(quat,[i,2],[1,1]) 412 | q3 = tf.slice(quat,[i,3],[1,1]) 413 | # Convert quaternion to rotation matrix. 414 | # Ref: http://www-evasion.inrialpes.fr/people/Franck.Hetroy/Teaching/ProjetsImage/2007/Bib/besl_mckay-pami1992.pdf 415 | # A method for Registration of 3D shapes paper by Paul J. Besl and Neil D McKay. 416 | R = [[q0*q0+q1*q1-q2*q2-q3*q3, 2*(q1*q2-q0*q3), 2*(q1*q3+q0*q2)], 417 | [2*(q1*q2+q0*q3), q0*q0+q2*q2-q1*q1-q3*q3, 2*(q2*q3-q0*q1)], 418 | [2*(q1*q3-q0*q2), 2*(q2*q3+q0*q1), q0*q0+q3*q3-q1*q1-q2*q2]] 419 | 420 | R = tf.reshape(R,[3,3]) # Convert R into a single tensor of shape 3x3. 421 | # tf.tensordot: Arg: tensor1, tensor2, axes 422 | # axes defined for tensor1 & tensor2 should be of same size. 423 | # axis 1 of R is of size 3 and axis 0 of data (3xnum_points) is of size 3. 424 | 425 | temp_rotated_data = tf.transpose(tf.tensordot(R, tf.transpose(data[i,...]), [1,0])) # Rotate the data. (num_points x 3) 426 | temp_rotated_data = tf.add(temp_rotated_data,translation[i,...]) # Add the translation (num_points x 3) 427 | transformed_data = tf.concat([transformed_data, temp_rotated_data],0) # Append data (batch_size x num_points x 3) 428 | transformed_data = tf.reshape(transformed_data, [-1,data.shape[1],3])[1:] # Reshape data and remove first point cloud. (batch_size x num_point x 3) 429 | return transformed_data 430 | 431 | 432 | 433 | ###################### Display Operations ######################### 434 | 435 | # Display data inside ModelNet files. 436 | def display_clouds(filename,model_no): 437 | # Arguments: 438 | # filename: Name of file to read the data from. (string) 439 | # model_no: Number to choose the model inside that file. (int) 440 | 441 | data = [] 442 | # Read the entire data from that file. 443 | with open(os.path.join('data','templates',filename),'r') as csvfile: 444 | csvreader = csv.reader(csvfile) 445 | for row in csvreader: 446 | row = [float(x) for x in row] 447 | data.append(row) 448 | fig = plt.figure() 449 | ax = fig.add_subplot(111,projection='3d') 450 | data = np.asarray(data) 451 | 452 | start_idx = model_no*2048 453 | end_idx = (model_no+1)*2048 454 | data = data[start_idx:end_idx,:] # Choose specific data related to the given model number. 455 | 456 | X,Y,Z = [],[],[] 457 | for row in data: 458 | X.append(row[0]) 459 | Y.append(row[1]) 460 | Z.append(row[2]) 461 | ax.scatter(X,Y,Z) 462 | plt.show() 463 | 464 | # Display given Point Cloud Data in blue color (default). 465 | def display_clouds_data(data): 466 | # Arguments: 467 | # data: array of point clouds (num_points x 3) 468 | 469 | fig = plt.figure() 470 | ax = fig.add_subplot(111,projection='3d') 471 | try: 472 | data = data.tolist() 473 | except: 474 | pass 475 | X,Y,Z = [],[],[] 476 | for row in data: 477 | X.append(row[0]) 478 | Y.append(row[1]) 479 | Z.append(row[2]) 480 | ax.scatter(X,Y,Z) 481 | plt.show() 482 | 483 | # Display given template, source and predicted point cloud data. 484 | def display_three_clouds(data1,data2,data3,title): 485 | # Arguments: 486 | # data1 Template Data (num_points x 3) (Red) 487 | # data2 Source Data (num_points x 3) (Green) 488 | # data3 Predicted Data (num_points x 3) (Blue) 489 | 490 | fig = plt.figure() 491 | ax = fig.add_subplot(111,projection='3d') 492 | try: 493 | data1 = data1.tolist() 494 | data2 = data2.tolist() 495 | data3 = data3.tolist() 496 | except: 497 | pass 498 | # Add Template Data in Plot 499 | X,Y,Z = [],[],[] 500 | for row in data1: 501 | X.append(row[0]) 502 | Y.append(row[1]) 503 | Z.append(row[2]) 504 | l1 = ax.scatter(X,Y,Z,c=[1,0,0,1]) 505 | # Add Source Data in Plot 506 | X,Y,Z = [],[],[] 507 | for row in data2: 508 | X.append(row[0]) 509 | Y.append(row[1]) 510 | Z.append(row[2]) 511 | l2 = ax.scatter(X,Y,Z,c=[0,1,0,1]) 512 | # Add Predicted Data in Plot 513 | X,Y,Z = [],[],[] 514 | for row in data3: 515 | X.append(row[0]) 516 | Y.append(row[1]) 517 | Z.append(row[2]) 518 | l3 = ax.scatter(X,Y,Z,c=[0,0,1,1]) 519 | 520 | # Add details to Plot. 521 | plt.legend((l1,l2,l3),('Template Data','Source Data','Predicted Data'),prop={'size':15},markerscale=4) 522 | ax.tick_params(labelsize=10) 523 | ax.set_xlabel('X-axis',fontsize=15) 524 | ax.set_ylabel('Y-axis',fontsize=15) 525 | ax.set_zlabel('Z-axis',fontsize=15) 526 | plt.title(title,fontdict={'fontsize':25}) 527 | ax.xaxis.set_tick_params(labelsize=15) 528 | ax.yaxis.set_tick_params(labelsize=15) 529 | ax.zaxis.set_tick_params(labelsize=15) 530 | plt.show() 531 | 532 | # Display given template, source and predicted point cloud data. 533 | def display_two_clouds(data1,data2): 534 | # Arguments: 535 | # data1 Template Data (num_points x 3) (Red) 536 | # data2 Source Data (num_points x 3) (Green) 537 | # data3 Predicted Data (num_points x 3) (Blue) 538 | 539 | fig = plt.figure() 540 | ax = fig.add_subplot(111,projection='3d') 541 | try: 542 | data1 = data1.tolist() 543 | data2 = data2.tolist() 544 | data3 = data3.tolist() 545 | except: 546 | pass 547 | # Add Template Data in Plot 548 | X,Y,Z = [],[],[] 549 | for row in data1: 550 | X.append(row[0]) 551 | Y.append(row[1]) 552 | Z.append(row[2]) 553 | l1 = ax.scatter(X,Y,Z,c=[1,0,0,1]) 554 | # Add Source Data in Plot 555 | X,Y,Z = [],[],[] 556 | for row in data2: 557 | X.append(row[0]) 558 | Y.append(row[1]) 559 | Z.append(row[2]) 560 | l2 = ax.scatter(X,Y,Z,c=[0,0,1,1]) 561 | 562 | # Add details to Plot. 563 | plt.legend((l1,l2),('Input Data','Output Data'),prop={'size':15},markerscale=4) 564 | ax.tick_params(labelsize=10) 565 | ax.set_xlabel('X-axis',fontsize=15) 566 | ax.set_ylabel('Y-axis',fontsize=15) 567 | ax.set_zlabel('Z-axis',fontsize=15) 568 | # plt.title(title,fontdict={'fontsize':25}) 569 | ax.xaxis.set_tick_params(labelsize=15) 570 | ax.yaxis.set_tick_params(labelsize=15) 571 | ax.zaxis.set_tick_params(labelsize=15) 572 | ax.set_axis_off() 573 | plt.show() 574 | 575 | # Display template, source, predicted point cloud data with results after each iteration. 576 | def display_itr_clouds(data1,data2,data3,ITR,title): 577 | # Arguments: 578 | # data1 Template Data (num_points x 3) (Red) 579 | # data2 Source Data (num_points x 3) (Green) 580 | # data3 Predicted Data (num_points x 3) (Blue) 581 | # ITR Point Clouds obtained after each iteration (iterations x batch_size x num of points x 3) (Yellow) 582 | 583 | fig = plt.figure() 584 | ax = fig.add_subplot(111,projection='3d') 585 | print(ITR.shape) # Display Number of Point Clouds in ITR. 586 | try: 587 | data1 = data1.tolist() 588 | data2 = data2.tolist() 589 | data3 = data3.tolist() 590 | except: 591 | pass 592 | # Add Template Data in Plot 593 | X,Y,Z = [],[],[] 594 | for row in data1: 595 | X.append(row[0]) 596 | Y.append(row[1]) 597 | Z.append(row[2]) 598 | l1 = ax.scatter(X,Y,Z,c=[1,0,0,1]) 599 | # Add Source Data in Plot 600 | X,Y,Z = [],[],[] 601 | for row in data2: 602 | X.append(row[0]) 603 | Y.append(row[1]) 604 | Z.append(row[2]) 605 | l2 = ax.scatter(X,Y,Z,c=[0,1,0,1]) 606 | # Add Predicted Data in Plot 607 | X,Y,Z = [],[],[] 608 | for row in data3: 609 | X.append(row[0]) 610 | Y.append(row[1]) 611 | Z.append(row[2]) 612 | l3 = ax.scatter(X,Y,Z,c=[0,0,1,1]) 613 | # Add point clouds after each iteration in Plot. 614 | for itr_data in ITR: 615 | X,Y,Z = [],[],[] 616 | for row in itr_data[0]: 617 | X.append(row[0]) 618 | Y.append(row[1]) 619 | Z.append(row[2]) 620 | ax.scatter(X,Y,Z,c=[1,1,0,0.5]) 621 | 622 | # Add details to Plot. 623 | plt.legend((l1,l2,l3),('Template Data','Source Data','Predicted Data'),prop={'size':15},markerscale=4) 624 | ax.tick_params(labelsize=10) 625 | ax.set_xlabel('X-axis',fontsize=15) 626 | ax.set_ylabel('Y-axis',fontsize=15) 627 | ax.set_zlabel('Z-axis',fontsize=15) 628 | plt.title(title,fontdict={'fontsize':25}) 629 | ax.xaxis.set_tick_params(labelsize=15) 630 | ax.yaxis.set_tick_params(labelsize=15) 631 | ax.zaxis.set_tick_params(labelsize=15) 632 | plt.show() 633 | 634 | 635 | 636 | 637 | 638 | if __name__=='__main__': 639 | # a = np.array([[0,0,0,0,0,0],[0,0,0,90,0,0]]) 640 | # print a.shape 641 | # a = poses_euler2quat(a) 642 | # print(a[1,3]*a[1,3]+a[1,4]*a[1,4]+a[1,5]*a[1,5]+a[1,6]*a[1,6]) 643 | # print(a[0,3]*a[0,3]+a[0,4]*a[0,4]+a[0,5]*a[0,5]+a[0,6]*a[0,6]) 644 | # print a.shape 645 | # display_clouds('airplane_templates.csv',0) 646 | 647 | templates = helper.process_templates('multi_model_templates') 648 | # templates = helper.process_templates('templates') 649 | # airplane = templates[0,:,:] 650 | idx = 199 651 | fig = plt.figure() 652 | ax = fig.add_subplot(111,projection='3d') 653 | # start = idx*2048 654 | # end = (idx+1)*2048 655 | ax.scatter(templates[idx,:,0],templates[idx,:,1],templates[idx,:,2]) 656 | plt.show() 657 | print(templates.shape) 658 | 659 | -------------------------------------------------------------------------------- /utils/plyfile.py: -------------------------------------------------------------------------------- 1 | # Copyright 2014 Darsh Ranjan 2 | # 3 | # This file is part of python-plyfile. 4 | # 5 | # python-plyfile is free software: you can redistribute it and/or 6 | # modify it under the terms of the GNU General Public License as 7 | # published by the Free Software Foundation, either version 3 of the 8 | # License, or (at your option) any later version. 9 | # 10 | # python-plyfile is distributed in the hope that it will be useful, 11 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 12 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU 13 | # General Public License for more details. 14 | # 15 | # You should have received a copy of the GNU General Public License 16 | # along with python-plyfile. If not, see 17 | # . 18 | 19 | from itertools import islice as _islice 20 | 21 | import numpy as _np 22 | from sys import byteorder as _byteorder 23 | 24 | 25 | try: 26 | _range = xrange 27 | except NameError: 28 | _range = range 29 | 30 | 31 | # Many-many relation 32 | _data_type_relation = [ 33 | ('int8', 'i1'), 34 | ('char', 'i1'), 35 | ('uint8', 'u1'), 36 | ('uchar', 'b1'), 37 | ('uchar', 'u1'), 38 | ('int16', 'i2'), 39 | ('short', 'i2'), 40 | ('uint16', 'u2'), 41 | ('ushort', 'u2'), 42 | ('int32', 'i4'), 43 | ('int', 'i4'), 44 | ('uint32', 'u4'), 45 | ('uint', 'u4'), 46 | ('float32', 'f4'), 47 | ('float', 'f4'), 48 | ('float64', 'f8'), 49 | ('double', 'f8') 50 | ] 51 | 52 | _data_types = dict(_data_type_relation) 53 | _data_type_reverse = dict((b, a) for (a, b) in _data_type_relation) 54 | 55 | _types_list = [] 56 | _types_set = set() 57 | for (_a, _b) in _data_type_relation: 58 | if _a not in _types_set: 59 | _types_list.append(_a) 60 | _types_set.add(_a) 61 | if _b not in _types_set: 62 | _types_list.append(_b) 63 | _types_set.add(_b) 64 | 65 | 66 | _byte_order_map = { 67 | 'ascii': '=', 68 | 'binary_little_endian': '<', 69 | 'binary_big_endian': '>' 70 | } 71 | 72 | _byte_order_reverse = { 73 | '<': 'binary_little_endian', 74 | '>': 'binary_big_endian' 75 | } 76 | 77 | _native_byte_order = {'little': '<', 'big': '>'}[_byteorder] 78 | 79 | 80 | def _lookup_type(type_str): 81 | if type_str not in _data_type_reverse: 82 | try: 83 | type_str = _data_types[type_str] 84 | except KeyError: 85 | raise ValueError("field type %r not in %r" % 86 | (type_str, _types_list)) 87 | 88 | return _data_type_reverse[type_str] 89 | 90 | 91 | def _split_line(line, n): 92 | fields = line.split(None, n) 93 | if len(fields) == n: 94 | fields.append('') 95 | 96 | assert len(fields) == n + 1 97 | 98 | return fields 99 | 100 | 101 | def make2d(array, cols=None, dtype=None): 102 | ''' 103 | Make a 2D array from an array of arrays. The `cols' and `dtype' 104 | arguments can be omitted if the array is not empty. 105 | 106 | ''' 107 | if (cols is None or dtype is None) and not len(array): 108 | raise RuntimeError("cols and dtype must be specified for empty " 109 | "array") 110 | 111 | if cols is None: 112 | cols = len(array[0]) 113 | 114 | if dtype is None: 115 | dtype = array[0].dtype 116 | 117 | return _np.fromiter(array, [('_', dtype, (cols,))], 118 | count=len(array))['_'] 119 | 120 | 121 | class PlyParseError(Exception): 122 | 123 | ''' 124 | Raised when a PLY file cannot be parsed. 125 | 126 | The attributes `element', `row', `property', and `message' give 127 | additional information. 128 | 129 | ''' 130 | 131 | def __init__(self, message, element=None, row=None, prop=None): 132 | self.message = message 133 | self.element = element 134 | self.row = row 135 | self.prop = prop 136 | 137 | s = '' 138 | if self.element: 139 | s += 'element %r: ' % self.element.name 140 | if self.row is not None: 141 | s += 'row %d: ' % self.row 142 | if self.prop: 143 | s += 'property %r: ' % self.prop.name 144 | s += self.message 145 | 146 | Exception.__init__(self, s) 147 | 148 | def __repr__(self): 149 | return ('PlyParseError(%r, element=%r, row=%r, prop=%r)' % 150 | self.message, self.element, self.row, self.prop) 151 | 152 | 153 | class PlyData(object): 154 | 155 | ''' 156 | PLY file header and data. 157 | 158 | A PlyData instance is created in one of two ways: by the static 159 | method PlyData.read (to read a PLY file), or directly from __init__ 160 | given a sequence of elements (which can then be written to a PLY 161 | file). 162 | 163 | ''' 164 | 165 | def __init__(self, elements=[], text=False, byte_order='=', 166 | comments=[], obj_info=[]): 167 | ''' 168 | elements: sequence of PlyElement instances. 169 | 170 | text: whether the resulting PLY file will be text (True) or 171 | binary (False). 172 | 173 | byte_order: '<' for little-endian, '>' for big-endian, or '=' 174 | for native. This is only relevant if `text' is False. 175 | 176 | comments: sequence of strings that will be placed in the header 177 | between the 'ply' and 'format ...' lines. 178 | 179 | obj_info: like comments, but will be placed in the header with 180 | "obj_info ..." instead of "comment ...". 181 | 182 | ''' 183 | if byte_order == '=' and not text: 184 | byte_order = _native_byte_order 185 | 186 | self.byte_order = byte_order 187 | self.text = text 188 | 189 | self.comments = list(comments) 190 | self.obj_info = list(obj_info) 191 | self.elements = elements 192 | 193 | def _get_elements(self): 194 | return self._elements 195 | 196 | def _set_elements(self, elements): 197 | self._elements = tuple(elements) 198 | self._index() 199 | 200 | elements = property(_get_elements, _set_elements) 201 | 202 | def _get_byte_order(self): 203 | return self._byte_order 204 | 205 | def _set_byte_order(self, byte_order): 206 | if byte_order not in ['<', '>', '=']: 207 | raise ValueError("byte order must be '<', '>', or '='") 208 | 209 | self._byte_order = byte_order 210 | 211 | byte_order = property(_get_byte_order, _set_byte_order) 212 | 213 | def _index(self): 214 | self._element_lookup = dict((elt.name, elt) for elt in 215 | self._elements) 216 | if len(self._element_lookup) != len(self._elements): 217 | raise ValueError("two elements with same name") 218 | 219 | @staticmethod 220 | def _parse_header(stream): 221 | ''' 222 | Parse a PLY header from a readable file-like stream. 223 | 224 | ''' 225 | lines = [] 226 | comments = {'comment': [], 'obj_info': []} 227 | while True: 228 | line = stream.readline().decode('ascii').strip() 229 | fields = _split_line(line, 1) 230 | 231 | if fields[0] == 'end_header': 232 | break 233 | 234 | elif fields[0] in comments.keys(): 235 | lines.append(fields) 236 | else: 237 | lines.append(line.split()) 238 | 239 | a = 0 240 | if lines[a] != ['ply']: 241 | raise PlyParseError("expected 'ply'") 242 | 243 | a += 1 244 | while lines[a][0] in comments.keys(): 245 | comments[lines[a][0]].append(lines[a][1]) 246 | a += 1 247 | 248 | if lines[a][0] != 'format': 249 | raise PlyParseError("expected 'format'") 250 | 251 | if lines[a][2] != '1.0': 252 | raise PlyParseError("expected version '1.0'") 253 | 254 | if len(lines[a]) != 3: 255 | raise PlyParseError("too many fields after 'format'") 256 | 257 | fmt = lines[a][1] 258 | 259 | if fmt not in _byte_order_map: 260 | raise PlyParseError("don't understand format %r" % fmt) 261 | 262 | byte_order = _byte_order_map[fmt] 263 | text = fmt == 'ascii' 264 | 265 | a += 1 266 | while a < len(lines) and lines[a][0] in comments.keys(): 267 | comments[lines[a][0]].append(lines[a][1]) 268 | a += 1 269 | 270 | return PlyData(PlyElement._parse_multi(lines[a:]), 271 | text, byte_order, 272 | comments['comment'], comments['obj_info']) 273 | 274 | @staticmethod 275 | def read(stream): 276 | ''' 277 | Read PLY data from a readable file-like object or filename. 278 | 279 | ''' 280 | (must_close, stream) = _open_stream(stream, 'read') 281 | try: 282 | data = PlyData._parse_header(stream) 283 | for elt in data: 284 | elt._read(stream, data.text, data.byte_order) 285 | finally: 286 | if must_close: 287 | stream.close() 288 | 289 | return data 290 | 291 | def write(self, stream): 292 | ''' 293 | Write PLY data to a writeable file-like object or filename. 294 | 295 | ''' 296 | (must_close, stream) = _open_stream(stream, 'write') 297 | try: 298 | stream.write(self.header.encode('ascii')) 299 | stream.write(b'\r\n') 300 | for elt in self: 301 | elt._write(stream, self.text, self.byte_order) 302 | finally: 303 | if must_close: 304 | stream.close() 305 | 306 | @property 307 | def header(self): 308 | ''' 309 | Provide PLY-formatted metadata for the instance. 310 | 311 | ''' 312 | lines = ['ply'] 313 | 314 | if self.text: 315 | lines.append('format ascii 1.0') 316 | else: 317 | lines.append('format ' + 318 | _byte_order_reverse[self.byte_order] + 319 | ' 1.0') 320 | 321 | # Some information is lost here, since all comments are placed 322 | # between the 'format' line and the first element. 323 | for c in self.comments: 324 | lines.append('comment ' + c) 325 | 326 | for c in self.obj_info: 327 | lines.append('obj_info ' + c) 328 | 329 | lines.extend(elt.header for elt in self.elements) 330 | lines.append('end_header') 331 | return '\r\n'.join(lines) 332 | 333 | def __iter__(self): 334 | return iter(self.elements) 335 | 336 | def __len__(self): 337 | return len(self.elements) 338 | 339 | def __contains__(self, name): 340 | return name in self._element_lookup 341 | 342 | def __getitem__(self, name): 343 | return self._element_lookup[name] 344 | 345 | def __str__(self): 346 | return self.header 347 | 348 | def __repr__(self): 349 | return ('PlyData(%r, text=%r, byte_order=%r, ' 350 | 'comments=%r, obj_info=%r)' % 351 | (self.elements, self.text, self.byte_order, 352 | self.comments, self.obj_info)) 353 | 354 | 355 | def _open_stream(stream, read_or_write): 356 | if hasattr(stream, read_or_write): 357 | return (False, stream) 358 | try: 359 | return (True, open(stream, read_or_write[0] + 'b')) 360 | except TypeError: 361 | raise RuntimeError("expected open file or filename") 362 | 363 | 364 | class PlyElement(object): 365 | 366 | ''' 367 | PLY file element. 368 | 369 | A client of this library doesn't normally need to instantiate this 370 | directly, so the following is only for the sake of documenting the 371 | internals. 372 | 373 | Creating a PlyElement instance is generally done in one of two ways: 374 | as a byproduct of PlyData.read (when reading a PLY file) and by 375 | PlyElement.describe (before writing a PLY file). 376 | 377 | ''' 378 | 379 | def __init__(self, name, properties, count, comments=[]): 380 | ''' 381 | This is not part of the public interface. The preferred methods 382 | of obtaining PlyElement instances are PlyData.read (to read from 383 | a file) and PlyElement.describe (to construct from a numpy 384 | array). 385 | 386 | ''' 387 | self._name = str(name) 388 | self._check_name() 389 | self._count = count 390 | 391 | self._properties = tuple(properties) 392 | self._index() 393 | 394 | self.comments = list(comments) 395 | 396 | self._have_list = any(isinstance(p, PlyListProperty) 397 | for p in self.properties) 398 | 399 | @property 400 | def count(self): 401 | return self._count 402 | 403 | def _get_data(self): 404 | return self._data 405 | 406 | def _set_data(self, data): 407 | self._data = data 408 | self._count = len(data) 409 | self._check_sanity() 410 | 411 | data = property(_get_data, _set_data) 412 | 413 | def _check_sanity(self): 414 | for prop in self.properties: 415 | if prop.name not in self._data.dtype.fields: 416 | raise ValueError("dangling property %r" % prop.name) 417 | 418 | def _get_properties(self): 419 | return self._properties 420 | 421 | def _set_properties(self, properties): 422 | self._properties = tuple(properties) 423 | self._check_sanity() 424 | self._index() 425 | 426 | properties = property(_get_properties, _set_properties) 427 | 428 | def _index(self): 429 | self._property_lookup = dict((prop.name, prop) 430 | for prop in self._properties) 431 | if len(self._property_lookup) != len(self._properties): 432 | raise ValueError("two properties with same name") 433 | 434 | def ply_property(self, name): 435 | return self._property_lookup[name] 436 | 437 | @property 438 | def name(self): 439 | return self._name 440 | 441 | def _check_name(self): 442 | if any(c.isspace() for c in self._name): 443 | msg = "element name %r contains spaces" % self._name 444 | raise ValueError(msg) 445 | 446 | def dtype(self, byte_order='='): 447 | ''' 448 | Return the numpy dtype of the in-memory representation of the 449 | data. (If there are no list properties, and the PLY format is 450 | binary, then this also accurately describes the on-disk 451 | representation of the element.) 452 | 453 | ''' 454 | return [(prop.name, prop.dtype(byte_order)) 455 | for prop in self.properties] 456 | 457 | @staticmethod 458 | def _parse_multi(header_lines): 459 | ''' 460 | Parse a list of PLY element definitions. 461 | 462 | ''' 463 | elements = [] 464 | while header_lines: 465 | (elt, header_lines) = PlyElement._parse_one(header_lines) 466 | elements.append(elt) 467 | 468 | return elements 469 | 470 | @staticmethod 471 | def _parse_one(lines): 472 | ''' 473 | Consume one element definition. The unconsumed input is 474 | returned along with a PlyElement instance. 475 | 476 | ''' 477 | a = 0 478 | line = lines[a] 479 | 480 | if line[0] != 'element': 481 | raise PlyParseError("expected 'element'") 482 | if len(line) > 3: 483 | raise PlyParseError("too many fields after 'element'") 484 | if len(line) < 3: 485 | raise PlyParseError("too few fields after 'element'") 486 | 487 | (name, count) = (line[1], int(line[2])) 488 | 489 | comments = [] 490 | properties = [] 491 | while True: 492 | a += 1 493 | if a >= len(lines): 494 | break 495 | 496 | if lines[a][0] == 'comment': 497 | comments.append(lines[a][1]) 498 | elif lines[a][0] == 'property': 499 | properties.append(PlyProperty._parse_one(lines[a])) 500 | else: 501 | break 502 | 503 | return (PlyElement(name, properties, count, comments), 504 | lines[a:]) 505 | 506 | @staticmethod 507 | def describe(data, name, len_types={}, val_types={}, 508 | comments=[]): 509 | ''' 510 | Construct a PlyElement from an array's metadata. 511 | 512 | len_types and val_types can be given as mappings from list 513 | property names to type strings (like 'u1', 'f4', etc., or 514 | 'int8', 'float32', etc.). These can be used to define the length 515 | and value types of list properties. List property lengths 516 | always default to type 'u1' (8-bit unsigned integer), and value 517 | types default to 'i4' (32-bit integer). 518 | 519 | ''' 520 | if not isinstance(data, _np.ndarray): 521 | raise TypeError("only numpy arrays are supported") 522 | 523 | if len(data.shape) != 1: 524 | raise ValueError("only one-dimensional arrays are " 525 | "supported") 526 | 527 | count = len(data) 528 | 529 | properties = [] 530 | descr = data.dtype.descr 531 | 532 | for t in descr: 533 | if not isinstance(t[1], str): 534 | raise ValueError("nested records not supported") 535 | 536 | if not t[0]: 537 | raise ValueError("field with empty name") 538 | 539 | if len(t) != 2 or t[1][1] == 'O': 540 | # non-scalar field, which corresponds to a list 541 | # property in PLY. 542 | 543 | if t[1][1] == 'O': 544 | if len(t) != 2: 545 | raise ValueError("non-scalar object fields not " 546 | "supported") 547 | 548 | len_str = _data_type_reverse[len_types.get(t[0], 'u1')] 549 | if t[1][1] == 'O': 550 | val_type = val_types.get(t[0], 'i4') 551 | val_str = _lookup_type(val_type) 552 | else: 553 | val_str = _lookup_type(t[1][1:]) 554 | 555 | prop = PlyListProperty(t[0], len_str, val_str) 556 | else: 557 | val_str = _lookup_type(t[1][1:]) 558 | prop = PlyProperty(t[0], val_str) 559 | 560 | properties.append(prop) 561 | 562 | elt = PlyElement(name, properties, count, comments) 563 | elt.data = data 564 | 565 | return elt 566 | 567 | def _read(self, stream, text, byte_order): 568 | ''' 569 | Read the actual data from a PLY file. 570 | 571 | ''' 572 | if text: 573 | self._read_txt(stream) 574 | else: 575 | if self._have_list: 576 | # There are list properties, so a simple load is 577 | # impossible. 578 | self._read_bin(stream, byte_order) 579 | else: 580 | # There are no list properties, so loading the data is 581 | # much more straightforward. 582 | self._data = _np.fromfile(stream, 583 | self.dtype(byte_order), 584 | self.count) 585 | 586 | if len(self._data) < self.count: 587 | k = len(self._data) 588 | del self._data 589 | raise PlyParseError("early end-of-file", self, k) 590 | 591 | self._check_sanity() 592 | 593 | def _write(self, stream, text, byte_order): 594 | ''' 595 | Write the data to a PLY file. 596 | 597 | ''' 598 | if text: 599 | self._write_txt(stream) 600 | else: 601 | if self._have_list: 602 | # There are list properties, so serialization is 603 | # slightly complicated. 604 | self._write_bin(stream, byte_order) 605 | else: 606 | # no list properties, so serialization is 607 | # straightforward. 608 | self.data.astype(self.dtype(byte_order), 609 | copy=False).tofile(stream) 610 | 611 | def _read_txt(self, stream): 612 | ''' 613 | Load a PLY element from an ASCII-format PLY file. The element 614 | may contain list properties. 615 | 616 | ''' 617 | self._data = _np.empty(self.count, dtype=self.dtype()) 618 | 619 | k = 0 620 | for line in _islice(iter(stream.readline, b''), self.count): 621 | fields = iter(line.strip().split()) 622 | for prop in self.properties: 623 | try: 624 | self._data[prop.name][k] = prop._from_fields(fields) 625 | except StopIteration: 626 | raise PlyParseError("early end-of-line", 627 | self, k, prop) 628 | except ValueError: 629 | raise PlyParseError("malformed input", 630 | self, k, prop) 631 | try: 632 | next(fields) 633 | except StopIteration: 634 | pass 635 | else: 636 | raise PlyParseError("expected end-of-line", self, k) 637 | k += 1 638 | 639 | if k < self.count: 640 | del self._data 641 | raise PlyParseError("early end-of-file", self, k) 642 | 643 | def _write_txt(self, stream): 644 | ''' 645 | Save a PLY element to an ASCII-format PLY file. The element may 646 | contain list properties. 647 | 648 | ''' 649 | for rec in self.data: 650 | fields = [] 651 | for prop in self.properties: 652 | fields.extend(prop._to_fields(rec[prop.name])) 653 | 654 | _np.savetxt(stream, [fields], '%.18g', newline='\r\n') 655 | 656 | def _read_bin(self, stream, byte_order): 657 | ''' 658 | Load a PLY element from a binary PLY file. The element may 659 | contain list properties. 660 | 661 | ''' 662 | self._data = _np.empty(self.count, dtype=self.dtype(byte_order)) 663 | 664 | for k in _range(self.count): 665 | for prop in self.properties: 666 | try: 667 | self._data[prop.name][k] = \ 668 | prop._read_bin(stream, byte_order) 669 | except StopIteration: 670 | raise PlyParseError("early end-of-file", 671 | self, k, prop) 672 | 673 | def _write_bin(self, stream, byte_order): 674 | ''' 675 | Save a PLY element to a binary PLY file. The element may 676 | contain list properties. 677 | 678 | ''' 679 | for rec in self.data: 680 | for prop in self.properties: 681 | prop._write_bin(rec[prop.name], stream, byte_order) 682 | 683 | @property 684 | def header(self): 685 | ''' 686 | Format this element's metadata as it would appear in a PLY 687 | header. 688 | 689 | ''' 690 | lines = ['element %s %d' % (self.name, self.count)] 691 | 692 | # Some information is lost here, since all comments are placed 693 | # between the 'element' line and the first property definition. 694 | for c in self.comments: 695 | lines.append('comment ' + c) 696 | 697 | lines.extend(list(map(str, self.properties))) 698 | 699 | return '\r\n'.join(lines) 700 | 701 | def __getitem__(self, key): 702 | return self.data[key] 703 | 704 | def __setitem__(self, key, value): 705 | self.data[key] = value 706 | 707 | def __str__(self): 708 | return self.header 709 | 710 | def __repr__(self): 711 | return ('PlyElement(%r, %r, count=%d, comments=%r)' % 712 | (self.name, self.properties, self.count, 713 | self.comments)) 714 | 715 | 716 | class PlyProperty(object): 717 | 718 | ''' 719 | PLY property description. This class is pure metadata; the data 720 | itself is contained in PlyElement instances. 721 | 722 | ''' 723 | 724 | def __init__(self, name, val_dtype): 725 | self._name = str(name) 726 | self._check_name() 727 | self.val_dtype = val_dtype 728 | 729 | def _get_val_dtype(self): 730 | return self._val_dtype 731 | 732 | def _set_val_dtype(self, val_dtype): 733 | self._val_dtype = _data_types[_lookup_type(val_dtype)] 734 | 735 | val_dtype = property(_get_val_dtype, _set_val_dtype) 736 | 737 | @property 738 | def name(self): 739 | return self._name 740 | 741 | def _check_name(self): 742 | if any(c.isspace() for c in self._name): 743 | msg = "Error: property name %r contains spaces" % self._name 744 | raise RuntimeError(msg) 745 | 746 | @staticmethod 747 | def _parse_one(line): 748 | assert line[0] == 'property' 749 | 750 | if line[1] == 'list': 751 | if len(line) > 5: 752 | raise PlyParseError("too many fields after " 753 | "'property list'") 754 | if len(line) < 5: 755 | raise PlyParseError("too few fields after " 756 | "'property list'") 757 | 758 | return PlyListProperty(line[4], line[2], line[3]) 759 | 760 | else: 761 | if len(line) > 3: 762 | raise PlyParseError("too many fields after " 763 | "'property'") 764 | if len(line) < 3: 765 | raise PlyParseError("too few fields after " 766 | "'property'") 767 | 768 | return PlyProperty(line[2], line[1]) 769 | 770 | def dtype(self, byte_order='='): 771 | ''' 772 | Return the numpy dtype description for this property (as a tuple 773 | of strings). 774 | 775 | ''' 776 | return byte_order + self.val_dtype 777 | 778 | def _from_fields(self, fields): 779 | ''' 780 | Parse from generator. Raise StopIteration if the property could 781 | not be read. 782 | 783 | ''' 784 | return _np.dtype(self.dtype()).type(next(fields)) 785 | 786 | def _to_fields(self, data): 787 | ''' 788 | Return generator over one item. 789 | 790 | ''' 791 | yield _np.dtype(self.dtype()).type(data) 792 | 793 | def _read_bin(self, stream, byte_order): 794 | ''' 795 | Read data from a binary stream. Raise StopIteration if the 796 | property could not be read. 797 | 798 | ''' 799 | try: 800 | return _np.fromfile(stream, self.dtype(byte_order), 1)[0] 801 | except IndexError: 802 | raise StopIteration 803 | 804 | def _write_bin(self, data, stream, byte_order): 805 | ''' 806 | Write data to a binary stream. 807 | 808 | ''' 809 | _np.dtype(self.dtype(byte_order)).type(data).tofile(stream) 810 | 811 | def __str__(self): 812 | val_str = _data_type_reverse[self.val_dtype] 813 | return 'property %s %s' % (val_str, self.name) 814 | 815 | def __repr__(self): 816 | return 'PlyProperty(%r, %r)' % (self.name, 817 | _lookup_type(self.val_dtype)) 818 | 819 | 820 | class PlyListProperty(PlyProperty): 821 | 822 | ''' 823 | PLY list property description. 824 | 825 | ''' 826 | 827 | def __init__(self, name, len_dtype, val_dtype): 828 | PlyProperty.__init__(self, name, val_dtype) 829 | 830 | self.len_dtype = len_dtype 831 | 832 | def _get_len_dtype(self): 833 | return self._len_dtype 834 | 835 | def _set_len_dtype(self, len_dtype): 836 | self._len_dtype = _data_types[_lookup_type(len_dtype)] 837 | 838 | len_dtype = property(_get_len_dtype, _set_len_dtype) 839 | 840 | def dtype(self, byte_order='='): 841 | ''' 842 | List properties always have a numpy dtype of "object". 843 | 844 | ''' 845 | return '|O' 846 | 847 | def list_dtype(self, byte_order='='): 848 | ''' 849 | Return the pair (len_dtype, val_dtype) (both numpy-friendly 850 | strings). 851 | 852 | ''' 853 | return (byte_order + self.len_dtype, 854 | byte_order + self.val_dtype) 855 | 856 | def _from_fields(self, fields): 857 | (len_t, val_t) = self.list_dtype() 858 | 859 | n = int(_np.dtype(len_t).type(next(fields))) 860 | 861 | data = _np.loadtxt(list(_islice(fields, n)), val_t, ndmin=1) 862 | if len(data) < n: 863 | raise StopIteration 864 | 865 | return data 866 | 867 | def _to_fields(self, data): 868 | ''' 869 | Return generator over the (numerical) PLY representation of the 870 | list data (length followed by actual data). 871 | 872 | ''' 873 | (len_t, val_t) = self.list_dtype() 874 | 875 | data = _np.asarray(data, dtype=val_t).ravel() 876 | 877 | yield _np.dtype(len_t).type(data.size) 878 | for x in data: 879 | yield x 880 | 881 | def _read_bin(self, stream, byte_order): 882 | (len_t, val_t) = self.list_dtype(byte_order) 883 | 884 | try: 885 | n = _np.fromfile(stream, len_t, 1)[0] 886 | except IndexError: 887 | raise StopIteration 888 | 889 | data = _np.fromfile(stream, val_t, n) 890 | if len(data) < n: 891 | raise StopIteration 892 | 893 | return data 894 | 895 | def _write_bin(self, data, stream, byte_order): 896 | ''' 897 | Write data to a binary stream. 898 | 899 | ''' 900 | (len_t, val_t) = self.list_dtype(byte_order) 901 | 902 | data = _np.asarray(data, dtype=val_t).ravel() 903 | 904 | _np.array(data.size, dtype=len_t).tofile(stream) 905 | data.tofile(stream) 906 | 907 | def __str__(self): 908 | len_str = _data_type_reverse[self.len_dtype] 909 | val_str = _data_type_reverse[self.val_dtype] 910 | return 'property list %s %s %s' % (len_str, val_str, self.name) 911 | 912 | def __repr__(self): 913 | return ('PlyListProperty(%r, %r, %r)' % 914 | (self.name, 915 | _lookup_type(self.len_dtype), 916 | _lookup_type(self.val_dtype))) 917 | --------------------------------------------------------------------------------