# Minhyuk Sung (mhsung@cs.stanford.edu)
# April 2018
import os, sys
BASE_DIR = os.path.normpath(
os.path.join(os.path.dirname(os.path.abspath(__file__))))
sys.path.append(os.path.join(BASE_DIR, 'utils'))
from pointnet_util import pointnet_sa_module, pointnet_fp_module
from transform_nets import *
import tensorflow as tf
import tf_util
def get_batch_norm_decay(global_step, batch_size, bn_decay_step):
BN_INIT_DECAY = 0.5
BN_DECAY_RATE = 0.5
BN_DECAY_CLIP = 0.99
bn_momentum = tf.train.exponential_decay(
BN_INIT_DECAY,
global_step*batch_size,
bn_decay_step,
BN_DECAY_RATE,
staircase=True)
bn_decay = tf.minimum(BN_DECAY_CLIP, 1 - bn_momentum)
return bn_decay
def get_learning_rate(init_learning_rate, global_step, batch_size,
decay_step, decay_rate):
learning_rate = tf.train.exponential_decay(
init_learning_rate,
global_step*batch_size,
decay_step,
decay_rate,
staircase=True)
learing_rate = tf.maximum(learning_rate, 0.00001)
return learning_rate
def build_pointnet_seg(X, out_dim, is_training, bn_decay, scope, use_stn=True):
n_points = X.get_shape()[1].value
n_dim = X.get_shape()[2].value
if use_stn:
### STN 1 ###
with tf.variable_scope('stn_1') as sc:
transform_1 = input_transform_net(X, is_training, bn_decay, K=n_dim)
X = tf.matmul(X, transform_1)
###
X_expanded = tf.expand_dims(X, -1)
net = tf_util.conv2d(X_expanded, 64, [1,n_dim], padding='VALID',
stride=[1,1], bn=True, is_training=is_training,
bn_decay=bn_decay, scope=scope+'_conv1')
net = tf_util.conv2d(net, 64, [1,1], padding='VALID',
stride=[1,1], bn=True, is_training=is_training,
bn_decay=bn_decay, scope=scope+'_conv2')
if use_stn:
### STN 2 ###
with tf.variable_scope('stn_2') as sc:
transform_2 = feature_transform_net(net, is_training, bn_decay, K=64)
net = tf.matmul(tf.squeeze(net, axis=2), transform_2)
net = tf.expand_dims(net, axis=2)
###
net_per_points = net
net = tf_util.conv2d(net, 64, [1,1], padding='VALID',
stride=[1,1], bn=True, is_training=is_training,
bn_decay=bn_decay, scope=scope+'_conv3')
net = tf_util.conv2d(net, 128, [1,1], padding='VALID',
stride=[1,1], bn=True, is_training=is_training,
bn_decay=bn_decay, scope=scope+'_conv4')
net = tf_util.conv2d(net, 1024, [1,1], padding='VALID',
stride=[1,1], bn=True, is_training=is_training,
bn_decay=bn_decay, scope=scope+'_conv5')
net = tf_util.max_pool2d(net, [n_points, 1], padding='VALID',
scope=scope+'_maxpool')
net_global = tf.tile(net, [1, n_points, 1, 1])
net = tf.concat([net_per_points, net_global], 3)
net = tf_util.conv2d(net, 512, [1,1], padding='VALID',
stride=[1,1], bn=True, is_training=is_training,
bn_decay=bn_decay, scope=scope+'_conv6')
net = tf_util.conv2d(net, 256, [1,1], padding='VALID',
stride=[1,1], bn=True, is_training=is_training,
bn_decay=bn_decay, scope=scope+'_conv7')
net = tf_util.conv2d(net, 128, [1,1], padding='VALID',
stride=[1,1], bn=True, is_training=is_training,
bn_decay=bn_decay, scope=scope+'_conv8')
net = tf_util.conv2d(net, 128, [1,1], padding='VALID',
stride=[1,1], bn=True, is_training=is_training,
bn_decay=bn_decay, scope=scope+'_conv9')
net = tf_util.conv2d(net, out_dim, [1,1], padding='VALID',
stride=[1,1], activation_fn=None, scope=scope+'_net')
net = tf.squeeze(net)
# Transform loss
if use_stn:
stn_loss = transform_loss(transform_2)
else:
stn_loss = 0.
return net, stn_loss
def build_pointnet2_seg(X, out_dim, is_training, bn_decay, scope):
n_points = X.get_shape()[1].value
l0_xyz = tf.slice(X, [0,0,0], [-1,-1,3])
l0_points = tf.slice(X, [0,0,3], [-1,-1,0])
# Set Abstraction layers
l1_xyz, l1_points, l1_indices = pointnet_sa_module(l0_xyz, l0_points,
npoint=512, radius=0.2, nsample=64, mlp=[64,64,128],
mlp2=None, group_all=False, is_training=is_training,
bn_decay=bn_decay, scope='layer1')
l2_xyz, l2_points, l2_indices = pointnet_sa_module(l1_xyz, l1_points,
npoint=128, radius=0.4, nsample=64, mlp=[128,128,256],
mlp2=None, group_all=False, is_training=is_training,
bn_decay=bn_decay, scope='layer2')
l3_xyz, l3_points, l3_indices = pointnet_sa_module(l2_xyz, l2_points,
npoint=None, radius=None, nsample=None, mlp=[256,512,1024],
mlp2=None, group_all=True, is_training=is_training,
bn_decay=bn_decay, scope='layer3')
# Feature Propagation layers
l2_points = pointnet_fp_module(l2_xyz, l3_xyz, l2_points, l3_points,
[256,256], is_training, bn_decay, scope='fa_layer1')
l1_points = pointnet_fp_module(l1_xyz, l2_xyz, l1_points, l2_points,
[256,128], is_training, bn_decay, scope='fa_layer2')
l0_points = pointnet_fp_module(l0_xyz, l1_xyz,
tf.concat([l0_xyz,l0_points],axis=-1), l1_points, [128,128,128],
is_training, bn_decay, scope='fa_layer3')
# FC layers
net = tf_util.conv1d(l0_points, 128, 1, padding='VALID', bn=True,
is_training=is_training, scope='fc1', bn_decay=bn_decay)
net = tf_util.dropout(net, keep_prob=0.5, is_training=is_training,
scope='dp1')
net = tf_util.conv1d(net, out_dim, 1, padding='VALID', activation_fn=None,
scope='fc2')
return net, 0
def residual_block(X, out_dim, is_training, bn_decay, scope):
n_dim = X.get_shape()[2].value
X_expanded = tf.expand_dims(X, -2)
net = tf_util.conv2d(X_expanded, out_dim, [1,1], padding='VALID',
stride=[1,1], bn=True, is_training=is_training,
bn_decay=bn_decay, activation_fn=tf.nn.relu, scope=scope+'_conv1')
net = tf_util.conv2d(net, out_dim, [1,1], padding='VALID',
stride=[1,1], bn=True, is_training=is_training,
bn_decay=bn_decay, activation_fn=None, scope=scope+'_conv2')
if n_dim != out_dim:
X_expanded = tf_util.conv2d(X_expanded, out_dim, [1,1], padding='VALID',
stride=[1,1], activation_fn=None, scope=scope+'_conv3')
net = tf.nn.relu(net + X_expanded)
net = tf.squeeze(net)
return net
def build_resnet_seg(X, out_dim, is_training, bn_decay, scope):
n_dim = X.get_shape()[2].value
n_layers = 7
net = X
for i in range(n_layers):
dim = out_dim if (i + 1) == n_layers else n_dim
net = residual_block(net, dim, is_training, bn_decay,
'{}_{:d}'.format(scope, i))
return net, 0
