def deconv(pts, fts, qrs, qrs_fts, C, tag, is_training):
"""
pts: points of previous layer, e.g.,(B, N/2, 3)
fts: point features of previous layer, e.g.,(B, N/2, C)
qrs: selected representative points of this layer, e.g.,(B, N, 3)
qrs_fts: e.g., (B, N, C)
"""
dist, idx = three_nn(qrs, pts)
dist = tf.maximum(dist, 1e-10)
norm = tf.reduce_sum((1.0 / dist), axis=2, keep_dims=True)
norm = tf.tile(norm, [1, 1, 3])
weight = (1.0 / dist) / norm
interpolated_fts = three_interpolate(fts, idx, weight) # (B, N, C)
if qrs_fts is not None:
interpolated_fts = tf.concat([interpolated_fts, qrs_fts], axis=2)
interpolated_fts = tf.expand_dims(interpolated_fts, 2)
new_features = pf.conv2d(interpolated_fts,
C,
tag,
is_training=is_training,
kernel_size=[1, 1])
new_features = tf.squeeze(new_features, [2])
return new_features
def input_transform_net(point_cloud, is_training,N, PP,KK=3):
""" Input (XYZ) Transform Net, input is BxNx3 gray image
Return:
Transformation matrix of size 3xK """
input_image = tf.expand_dims(point_cloud, -1)
print(input_image)
net = pf.conv2d(input_image, 64, 'tconv1', is_training,(1,3))
net = pf.dense(net, 128, 'tconv2', is_training)
net = pf.dense(net, 1024, 'tconv3', is_training)
print(net)
net = pf.max_pool2d(net,[1024,1],'tmaxpool')
net = pf.dense(net, 512, 'tfc1',is_training)
net = pf.dense(net, 256, 'tfc2',is_training)
net = tf.reshape(net, (N, -1))
with tf.variable_scope('transform_XYZ') as sc:
assert(KK==3)
weights = tf.get_variable('weights', (256, 3*KK),
initializer=tf.constant_initializer(0.0),
dtype=tf.float32)
biases = tf.get_variable('biases', (3*KK),
initializer=tf.constant_initializer(0.0),
dtype=tf.float32)
biases += tf.constant([1,0,0,0,1,0,0,0,1], dtype=tf.float32)
transform = tf.matmul(net, weights)
transform = tf.nn.bias_add(transform, biases)
transform = tf.reshape(transform, (N, 3, KK))
return transform
def point_conv_transpose_fc_no_skip(pts_in,
fts_in,
indices,
np_out,
K,
C,
is_training,
activation=tf.nn.relu,
bn=True,
tag='point_deconv'):
'''
NOTE: no skipped link of output points and features
pts_in: xyz points, Tensor, (batch_size, npoint, 3)
fts_in: features, Tensor, (batch_size, npoint, channel_fts)
np_out: number of points that is expected
K: neighbor size
C: output feature channel number
'''
with tf.variable_scope(tag):
batch_size = fts_in.shape[0]
n_pts_in = fts_in.shape[1]
if pts_in is None:
nn_fts_input = fts_in # (N, P, C_fts_in)
#print('Error: fts_in is None.')
else:
nn_fts_input = tf.concat(
[pts_in, fts_in], axis=-1,
name='_nn_fts_input') # (N, P, C_fts_in+3)
nn_fts_input = tf.expand_dims(nn_fts_input,
axis=2) # (N, P, 1, C_fts_in)
# using fc
fts_fc = pf.dense(nn_fts_input, K * C, 'fc_fts_prop',
is_training) # (N, P, 1, K*C)
fts_fc = tf.reshape(fts_fc,
(batch_size, n_pts_in, K, C)) # (N, P, K, C)
fts_pts_conv = pf.conv2d(
fts_all_concat,
C,
tag + '_fts_pts_conv1D',
is_training,
kernel_size=(1, 1),
strides=(1, 1),
with_bn=bn,
activation=activation) # (N, P, K, C+3_ch_skip?) -> (N, P, K, C)
# summing up feature for each point
fts_dense = tf.scatter_nd(
tf.reshape(indices, (-1, 2)), tf.reshape(fts_pts_conv, (-1, C)),
(batch_size, n_pts_out, C)) # (N, N_PTS_OUT, c)
return fts_dense
def point_conv(pts,
fts,
qrs,
K,
D,
C,
is_training,
activation=tf.nn.relu,
bn=True,
sorting_method='cxyz',
center_patch=True,
tag='point_conv'):
'''
pts: xyz points, Tensor, (batch_size, npoint, 3)
fts: features, Tensor, (batch_size, npoint, channel_fts)
qrs: xyz points as queries, Tensor, (batch_size, n_qrs, 3)
K: number of neighbors for one query, scalar, int
D: dilation rate
C: output feature channel number
'''
with tf.variable_scope(tag):
_, indices_dilated = pf.knn_indices_general(qrs, pts, K * D, True)
indices = indices_dilated[:, :, ::D, :]
if sorting_method is not None:
indices = pf.sort_points(pts, indices, sorting_method)
nn_pts = tf.gather_nd(pts, indices, name='nn_pts') # (N, P, K, 3)
if center_patch:
nn_pts_center = tf.expand_dims(
qrs, axis=2, name='nn_pts_center') # (N, P, 1, 3)
nn_pts_local = tf.subtract(
nn_pts, nn_pts_center,
name='nn_pts_local') # (N, P, K, 3), move to local frame
else:
nn_pts_local = nn_pts
if fts is None:
nn_fts_input = nn_pts_local
else:
nn_fts_from_prev = tf.gather_nd(fts, indices, 'nn_fts_from_prev')
nn_fts_input = tf.concat(
[nn_pts_local, nn_fts_from_prev], axis=-1,
name='nn_fts_input') # (N, P, K, 3+C_prev_fts)
fts_conv = pf.conv2d(nn_fts_input,
C,
'fts_conv',
is_training, (1, K),
with_bn=bn,
activation=activation) # (N, P, 1, C)
fts_conv_3d = tf.squeeze(fts_conv, axis=2,
name='fts_conv_3d') # (N, P, C)
return fts_conv_3d, indices
def pointnet1(point_cloud, is_training, N, PP):
transform = input_transform_net(point_cloud, is_training, N, PP, KK=3)
point_cloud_transformed = tf.matmul(point_cloud, transform)
input_image = tf.expand_dims(point_cloud_transformed, -1) # (N,PP,3,1)
net = pf.conv2d(input_image, 64, 'p1_conv1', is_training,(1,3)) # (N,PP,1,64)
net = pf.dense(net, 64, 'p1_mlp1', is_training) # (N,PP,1,64)
net = pf.dense(net, 64, 'p1_mlp2', is_training) # (N,PP,1,64)
net = pf.dense(net, 128, 'p1_mlp3', is_training) # (N,PP,1,128)
net = pf.dense(net, 1024, 'p1_mlp4', is_training) # (N,PP,1,1024)
net1 = net
net = pf.max_pool2d(net,[PP,1],'p1_maxpool') # (N,1,1,1024)
net = tf.squeeze(net, axis=[2])
net1 = tf.squeeze(net1, axis=[2])
return net,net1
def xconv(pts, fts, qrs, tag, N, K, D, P, C, C_pts_fts, is_training, with_X_transformation, depth_multiplier,
sorting_method=None, with_global=False):
if D == 1:
_, indices = pf.knn_indices_general(qrs, pts, K, True)
else:
_, indices_dilated = pf.knn_indices_general(qrs, pts, K * D, True)
indices = indices_dilated[:, :, ::D, :]
if sorting_method is not None:
indices = pf.sort_points(pts, indices, sorting_method)
nn_pts = tf.gather_nd(pts, indices, name=tag + 'nn_pts') # (N, P, K, 3)
nn_pts_center = tf.expand_dims(qrs, axis=2, name=tag + 'nn_pts_center') # (N, P, 1, 3)
nn_pts_local = tf.subtract(nn_pts, nn_pts_center, name=tag + 'nn_pts_local') # (N, P, K, 3)
# Prepare features to be transformed
nn_fts_from_pts_0 = pf.dense(nn_pts_local, C_pts_fts, tag + 'nn_fts_from_pts_0', is_training)
nn_fts_from_pts = pf.dense(nn_fts_from_pts_0, C_pts_fts, tag + 'nn_fts_from_pts', is_training)
if fts is None:
nn_fts_input = nn_fts_from_pts
else:
nn_fts_from_prev = tf.gather_nd(fts, indices, name=tag + 'nn_fts_from_prev')
nn_fts_input = tf.concat([nn_fts_from_pts, nn_fts_from_prev], axis=-1, name=tag + 'nn_fts_input')
if with_X_transformation:
######################## X-transformation #########################
X_0 = pf.conv2d(nn_pts_local, K * K, tag + 'X_0', is_training, (1, K))
X_0_KK = tf.reshape(X_0, (N, P, K, K), name=tag + 'X_0_KK')
X_1 = pf.depthwise_conv2d(X_0_KK, K, tag + 'X_1', is_training, (1, K))
X_1_KK = tf.reshape(X_1, (N, P, K, K), name=tag + 'X_1_KK')
X_2 = pf.depthwise_conv2d(X_1_KK, K, tag + 'X_2', is_training, (1, K), activation=None)
X_2_KK = tf.reshape(X_2, (N, P, K, K), name=tag + 'X_2_KK')
fts_X = tf.matmul(X_2_KK, nn_fts_input, name=tag + 'fts_X')
###################################################################
else:
fts_X = nn_fts_input
fts_conv = pf.separable_conv2d(fts_X, C, tag + 'fts_conv', is_training, (1, K), depth_multiplier=depth_multiplier)
fts_conv_3d = tf.squeeze(fts_conv, axis=2, name=tag + 'fts_conv_3d')
if with_global:
fts_global_0 = pf.dense(qrs, C // 4, tag + 'fts_global_0', is_training)
fts_global = pf.dense(fts_global_0, C // 4, tag + 'fts_global_', is_training)
return tf.concat([fts_global, fts_conv_3d], axis=-1, name=tag + 'fts_conv_3d_with_global')
else:
return fts_conv_3d
def xconv(pts,
fts,
qrs,
tag,
N,
K,
D,
P,
C,
depth_multiplier,
is_training,
C_pts_fts,
with_x_transformation=True,
sorting_method=None,
with_global=False):
###xconv(pts,fts,qrs,tag,N,K,D,P,C,C_pts_fts,is_training,with_X_transformation,depth_multipiler,sorting_method,with_global)
# @params:
# N: Number of output points
# K: Number of nearest neighbot
# D: dilation rate
# C: Number of output channels
# P:the representative point number in the output, -1 means all input points are output representative points
# x_transformation: replace max_pooling in PointNet
# sorting_method:
# with_global
# pts:Input point cloud
# qrs:queries
# fts:features
_, indices_dilated = pf.knn_indices_general(qrs, pts, K * D, True)
indices = indices_dilated[:, :, ::D, :]
if sorting_method is not None:
indices = pf.sort_points(pts, indices, sorting_method)
nn_pts = tf.gather_nd(pts, indices, name=tag + 'nn_pts') #(N,P,K,3)
nn_pts_center = tf.expand_dims(qrs, axis=2,
name=tag + 'nn_pts_center') #(N,P,1,3)
nn_pts_local = tf.subtract(nn_pts,
nn_pts_center,
name=tag + 'nn_pts_local') #(N,P,K,3)
#Prepare features to be transformed
nn_fts_from_pts_0 = pf.dense(nn_pts_local, C_pts_fts,
tag + 'nn_fts_from_pts_0', is_training)
nn_fts_from_pts = pf.dense(nn_fts_from_pts_0, C_pts_fts,
tag + 'nn_fts_from_pts', is_training)
if fts is None:
nn_fts_input = nn_fts_from_pts
else:
nn_fts_from_prev = tf.gather_nd(fts,
indices,
name=tag + 'nn_fts_from_prev')
nn_fts_input = tf.concat([nn_fts_from_pts, nn_fts_from_prev],
axis=-1,
name=tag + 'nn_fts_input')
#X_transformation
if with_x_transformation:
############################X_transformation#######################################################
X_0 = pf.conv2d(nn_pts_local, K * K, tag + 'X_0', is_training, (1, K))
X_0_KK = tf.reshape(X_0, (N, P, K, K), name=tag + 'X_0_KK')
X_1 = pf.depthwise_conv2d(X_0_KK, K, tag + 'X_1', is_training, (1, K))
X_1_KK = tf.reshape(X_1, (N, P, K, K), name=tag + 'X_1_KK')
X_2 = pf.depthwise_conv2d(X_1_KK,
K,
tag + 'X_2',
is_training, (1, K),
activation=None)
X_2_KK = tf.reshape(X_2, (N, P, K, K), name=tag + 'X_2_KK')
fts_X = tf.matmul(X_2_KK, nn_fts_input, name=tag + 'fts_X')
#####################################################################################################
else:
fts_X = nn_fts_input
fts_conv = pf.separable_conv2d(fts_X,
C,
tag + 'fts_conv',
is_training, (1, K),
depth_multiplier=depth_multiplier)
fts_conv_3d = tf.squeeze(fts_conv, axis=2, name=tag + 'fts_conv_3d')
if with_global:
fts_global_0 = pf.dense(qrs, C // 4, tag + 'fts_global_0', is_training)
fts_global = pf.dense(fts_global_0, C // 4, tag + 'fts_global',
is_training)
return tf.concat([fts_global, fts_conv_3d],
axis=-1,
name=tag + 'fts_conv_3d_with_global')
else:
return fts_conv_3d
def xconv(pts,
fts,
qrs,
tag,
N,
K,
D,
P,
C,
C_pts_fts,
is_training,
with_X_transformation,
depth_multiplier,
sorting_method=None,
with_global=False):
_, indices_dilated = pf.knn_indices_general(
qrs, pts, K * D, True
) # 带孔K*D近邻； (N, P, 3) + (N, P, 3) => (N, P, K*D, 2) N=batch_num; P=point_num; 最后一维两个数：batch_indice + point_indice
indices = indices_dilated[:, :, ::
D, :] # (N, P, K, 2) 隔D取样，增加receptive field
if sorting_method is not None:
indices = pf.sort_points(pts, indices, sorting_method)
nn_pts = tf.gather_nd(
pts, indices, name=tag + 'nn_pts'
) # (N, P, 3) + (N, P, K, 2) => (N, P, K, 3) indices维度比pts高，最后一维做索引，batch_indice->N, point_indice->P, 最后剩下(N, P, 3)里的3是坐标，补进去返回形状是 (N, P, K, 3)
nn_pts_center = tf.expand_dims(qrs, axis=2,
name=tag + 'nn_pts_center') # (N, P, 1, 3)
nn_pts_local = tf.subtract(
nn_pts, nn_pts_center,
name=tag + 'nn_pts_local') # (N, P, K, 3)，减去中心点的坐标系，得到局部坐标系
# Prepare features to be transformed
# 2个全连接层 对点的相对坐标系升维，(N, P, K, 3) => (N, P, K, C_pts_fts)
# 升维的目的是为了和前面的特征拼接 (ct.)
nn_fts_from_pts_0 = pf.dense(nn_pts_local, C_pts_fts,
tag + 'nn_fts_from_pts_0', is_training)
nn_fts_from_pts = pf.dense(nn_fts_from_pts_0, C_pts_fts,
tag + 'nn_fts_from_pts', is_training)
# 如果第一次没有feature，就只使用从局部得到的点特征，如果有之前的特征，就拼接二者(ct.)
if fts is None:
nn_fts_input = nn_fts_from_pts
else:
nn_fts_from_prev = tf.gather_nd(fts,
indices,
name=tag + 'nn_fts_from_prev')
nn_fts_input = tf.concat([nn_fts_from_pts, nn_fts_from_prev],
axis=-1,
name=tag + 'nn_fts_input') # 拼接坐标特征和之前的特征
if with_X_transformation:
######################## X-transformation #########################
# (N, P, K, 3) > (N, P, 1, K * K) ； kernel size 是(1, K)
X_0 = pf.conv2d(nn_pts_local, K * K, tag + 'X_0', is_training, (1, K))
# (N, P, 1, K * K) > (N, P, K, K)
X_0_KK = tf.reshape(X_0, (N, P, K, K), name=tag + 'X_0_KK')
# (N, P, K, K) > (N, P, 1, K * K) 可能为了较少参数量用depthwise_conv2d换conv2d
X_1 = pf.depthwise_conv2d(X_0_KK, K, tag + 'X_1', is_training, (1, K))
# (N, P, 1, K * K) > (N, P, K, K)
X_1_KK = tf.reshape(X_1, (N, P, K, K), name=tag + 'X_1_KK')
# (N, P, K, K) > (N, P, 1, K * K)
X_2 = pf.depthwise_conv2d(X_1_KK,
K,
tag + 'X_2',
is_training, (1, K),
activation=None)
# (N, P, 1, K * K) > (N, P, K, K) 最后的x矩阵 （K*K）
X_2_KK = tf.reshape(X_2, (N, P, K, K), name=tag + 'X_2_KK')
# 用得到的 x 矩阵乘以之前的特征，所以这里其实是有2个分支，一个计算x矩阵,另外一个计算特征；(ct.)
# (N, P, K, K) * (N, P, K, C_pts_fts) = (N, P, K, C_pts_fts)
fts_X = tf.matmul(X_2_KK, nn_fts_input, name=tag + 'fts_X')
###################################################################
else:
fts_X = nn_fts_input
# 最后的分离卷积
fts_conv = pf.separable_conv2d(fts_X,
C,
tag + 'fts_conv',
is_training, (1, K),
depth_multiplier=depth_multiplier
) # (N, P, K, C_pts_fts) -> (N, P, 1, C)
fts_conv_3d = tf.squeeze(fts_conv, axis=2,
name=tag + 'fts_conv_3d') # (N, P, C)
# 用代表点全局位置信息,
if with_global:
fts_global_0 = pf.dense(qrs, C // 4, tag + 'fts_global_0', is_training)
fts_global = pf.dense(
fts_global_0, C // 4, tag + 'fts_global', is_training
) # (N, P, C//4) 最后一层时作者使receptive field < 1让model更会捕捉local feature; 从“subvolume supervision”得到启发 to further address the over-fitting problem.;
return tf.concat([fts_global, fts_conv_3d],
axis=-1,
name=tag + 'fts_conv_3d_with_global')
else:
return fts_conv_3d
def xconv_cov(cov,
pts,
fts,
qrs,
tag,
N,
K,
D,
P,
C,
C_pts_fts,
is_training,
with_X_transformation,
depth_multiplier,
sorting_method=None,
with_global=False):
"""
cov: point covariance (B, N, 9) of previous layer
pts: points of previous layer,
fts: point features of previous layer,
qrs: selected representative points of this layer,
N: batch_size,
K: neighbor_size,
D: dilation parameter,
P: the number of selected representative points,
C: output feature number per point,
C_pts_fts: feature number for each local point
"""
if D == 1:
_, indices = pf.knn_indices_general(qrs, pts, K, True)
else:
_, indices_dilated = pf.knn_indices_general(qrs, pts, K * D, True)
indices = indices_dilated[:, :, ::D, :]
if sorting_method is not None:
indices = pf.sort_points(pts, indices, sorting_method)
nn_covs = tf.gather_nd(cov, indices, name=tag + 'nn_pts') # (N, P, K, 9)
# nn_pts_center = tf.expand_dims(qrs, axis=2, name=tag + 'nn_pts_center') # (N, P, 1, 3)
# nn_pts_local = tf.subtract(nn_pts, nn_pts_center, name=tag + 'nn_pts_local') # (N, P, K, 3)
# Prepare features to be transformed
nn_fts_from_pts_0 = pf.dense(nn_covs,
C_pts_fts,
tag + 'nn_fts_from_pts_0',
is_training,
with_bn=False) # following paper
nn_fts_from_pts = pf.dense(nn_fts_from_pts_0,
C_pts_fts,
tag + 'nn_fts_from_pts',
is_training,
with_bn=False) # following paper
if fts is None:
nn_fts_input = nn_fts_from_pts
else:
nn_fts_from_prev = tf.gather_nd(fts,
indices,
name=tag + 'nn_fts_from_prev')
nn_fts_input = tf.concat([nn_fts_from_pts, nn_fts_from_prev],
axis=-1,
name=tag + 'nn_fts_input')
if with_X_transformation:
######################## X-transformation #########################
X_0 = pf.conv2d(nn_covs,
K * K,
tag + 'X_0',
is_training, (1, K),
with_bn=False) # following paper
X_0_KK = tf.reshape(X_0, (N, P, K, K), name=tag + 'X_0_KK')
X_1 = pf.depthwise_conv2d(X_0_KK,
K,
tag + 'X_1',
is_training, (1, K),
with_bn=False) # following paper
X_1_KK = tf.reshape(X_1, (N, P, K, K), name=tag + 'X_1_KK')
X_2 = pf.depthwise_conv2d(X_1_KK,
K,
tag + 'X_2',
is_training, (1, K),
with_bn=False,
activation=None) # following paper
X_2_KK = tf.reshape(X_2, (N, P, K, K), name=tag + 'X_2_KK')
fts_X = tf.matmul(X_2_KK, nn_fts_input, name=tag + 'fts_X')
###################################################################
else:
fts_X = nn_fts_input
fts_conv = pf.separable_conv2d(fts_X,
C,
tag + 'fts_conv',
is_training, (1, K),
depth_multiplier=depth_multiplier,
with_bn=True)
fts_conv_3d = tf.squeeze(fts_conv, axis=2, name=tag + 'fts_conv_3d')
if with_global:
fts_global_0 = pf.dense(qrs,
C // 4,
tag + 'fts_global_0',
is_training,
with_bn=True)
fts_global = pf.dense(fts_global_0,
C // 4,
tag + 'fts_global_',
is_training,
with_bn=True)
return tf.concat([fts_global, fts_conv_3d],
axis=-1,
name=tag + 'fts_conv_3d_with_global')
else:
return fts_conv_3d
def deconv_new(pts,
fts,
qrs,
tag,
N,
K,
radius,
P,
C,
C_pts_fts,
is_training,
with_X_transformation,
depth_multiplier,
D=1,
sorting_method=None,
with_global=False,
knn=False):
"""
pts: points of previous layer, e.g.,(B, N/2, 3)
fts: point features of previous layer, e.g.,(B, N/2, C)
qrs: selected representative points of this layer, e.g.,(B, N, 3)
N: batch_size,
K: neighbor_size,
D: dilation parameter,
P: the number of selected representative points,
C: output feature number per point,
C_pts_fts: feature number for each local point
radius: float32, the radius of query ball search
knn: True: knn; False: query ball
"""
dist, idx = three_nn(qrs, pts)
dist = tf.maximum(dist, 1e-10)
norm = tf.reduce_sum((1.0 / dist), axis=2, keep_dims=True)
norm = tf.tile(norm, [1, 1, 3])
weight = (1.0 / dist) / norm
interpolated_fts = three_interpolate(fts, idx, weight) # (B, N, C)
if knn:
_, indices_dilated = pf.knn_indices_general(qrs, qrs, K * D, True)
indices = indices_dilated[:, :, ::D, :]
nn_pts = tf.gather_nd(qrs, indices,
name=tag + 'nn_pts') # (B, N, K, 3)
nn_fts_from_prev = tf.gather_nd(interpolated_fts,
indices,
name=tag + 'nn_fts')
else:
indices, pts_cnt = query_ball_point(radius, K, qrs, qrs)
nn_pts = group_point(qrs, indices)
nn_fts_from_prev = group_point(interpolated_fts, indices)
nn_pts_center = tf.expand_dims(qrs, axis=2,
name=tag + 'nn_pts_center') # (B, N, 1, 3)
nn_pts_local = tf.subtract(nn_pts,
nn_pts_center,
name=tag + 'nn_pts_local') # (B, N, K, 3)
# Prepare features to be transformed
nn_fts_from_pts_0 = pf.dense(nn_pts_local, C_pts_fts,
tag + 'nn_fts_from_pts_0', is_training)
nn_fts_from_pts = pf.dense(nn_fts_from_pts_0, C_pts_fts,
tag + 'nn_fts_from_pts', is_training)
nn_fts_input = tf.concat([nn_fts_from_pts, nn_fts_from_prev],
axis=-1,
name=tag + 'nn_fts_input')
if with_X_transformation:
######################## X-transformation #########################
X_0 = pf.conv2d(nn_pts_local, K * K, tag + 'X_0', is_training,
(1, K)) # following paper
X_0_KK = tf.reshape(X_0, (N, P, K, K), name=tag + 'X_0_KK')
X_1 = pf.depthwise_conv2d(X_0_KK, K, tag + 'X_1', is_training,
(1, K)) # following paper
X_1_KK = tf.reshape(X_1, (N, P, K, K), name=tag + 'X_1_KK')
X_2 = pf.depthwise_conv2d(X_1_KK,
K,
tag + 'X_2',
is_training, (1, K),
activation=None) # following paper
X_2_KK = tf.reshape(X_2, (N, P, K, K), name=tag + 'X_2_KK')
fts_X = tf.matmul(X_2_KK, nn_fts_input, name=tag + 'fts_X')
###################################################################
else:
fts_X = nn_fts_input
fts_conv = pf.separable_conv2d(fts_X,
C,
tag + 'fts_conv',
is_training, (1, K),
depth_multiplier=depth_multiplier)
fts_conv_3d = tf.squeeze(fts_conv, axis=2, name=tag + 'fts_conv_3d')
if with_global:
fts_global_0 = pf.dense(qrs, C // 4, tag + 'fts_global_0', is_training)
fts_global = pf.dense(fts_global_0, C // 4, tag + 'fts_global_',
is_training)
return tf.concat([fts_global, fts_conv_3d],
axis=-1,
name=tag + 'fts_conv_3d_with_global')
else:
return fts_conv_3d
def ficonv(pts,
fts,
qrs,
tag,
N,
K1,
mm,
sigma,
scale,
K,
D,
P,
C,
C_pts_fts,
kernel_num,
is_training,
with_kernel_registering,
with_kernel_shape_comparison,
with_point_transformation,
with_feature_transformation,
with_learning_feature_transformation,
kenel_initialization_method,
depth_multiplier,
sorting_method=None,
with_global=False):
Dis, indices_dilated = pf.knn_indices_general(qrs, pts, K * D, True)
indices = indices_dilated[:, :, ::D, :]
if sorting_method is not None:
indices = pf.sort_points(pts, indices, sorting_method)
nn_pts = tf.gather_nd(pts, indices, name=tag + 'nn_pts') # (N, P, K, 3)
nn_pts_center = tf.expand_dims(qrs, axis=2,
name=tag + 'nn_pts_center') # (N, P, 1, 3)
nn_pts_local = tf.subtract(nn_pts,
nn_pts_center,
name=tag + 'nn_pts_local') # (N, P, K, 3)
if with_point_transformation or with_feature_transformation:
X_0 = pf.conv2d(nn_pts_local, K * K, tag + 'X_0', is_training, (1, K))
X_0_KK = tf.reshape(X_0, (N, P, K, K), name=tag + 'X_0_KK')
X_1 = pf.depthwise_conv2d(X_0_KK, K, tag + 'X_1', is_training, (1, K))
X_1_KK = tf.reshape(X_1, (N, P, K, K), name=tag + 'X_1_KK')
X_2 = pf.depthwise_conv2d(X_1_KK,
K,
tag + 'X_2',
is_training, (1, K),
activation=None)
X_2_KK = tf.reshape(X_2, (N, P, K, K), name=tag + 'X_2_KK')
if with_point_transformation:
if with_learning_feature_transformation:
nn_pts_local = tf.matmul(X_2_KK, nn_pts_local)
# Prepare features to be transformed
nn_fts_from_pts_0 = pf.dense(nn_pts_local, C_pts_fts,
tag + 'nn_fts_from_pts_0',
is_training)
nn_fts_from_pts = pf.dense(nn_fts_from_pts_0, C_pts_fts,
tag + 'nn_fts_from_pts', is_training)
else:
# Prepare features to be transformed
nn_fts_from_pts_0 = pf.dense(nn_pts_local, C_pts_fts,
tag + 'nn_fts_from_pts_0',
is_training)
nn_fts_from_pts = pf.dense(nn_fts_from_pts_0, C_pts_fts,
tag + 'nn_fts_from_pts', is_training)
nn_pts_local = tf.matmul(X_2_KK, nn_pts_local)
else:
if with_learning_feature_transformation:
nn_pts_local_ = tf.matmul(X_2_KK,
nn_pts_local,
name=tag + 'nn_pts_local_')
# Prepare features to be transformed
nn_fts_from_pts_0 = pf.dense(nn_pts_local_, C_pts_fts,
tag + 'nn_fts_from_pts_0',
is_training)
nn_fts_from_pts = pf.dense(nn_fts_from_pts_0, C_pts_fts,
tag + 'nn_fts_from_pts', is_training)
else:
nn_fts_from_pts_0 = pf.dense(nn_pts_local, C_pts_fts,
tag + 'nn_fts_from_pts_0',
is_training)
nn_fts_from_pts = pf.dense(nn_fts_from_pts_0, C_pts_fts,
tag + 'nn_fts_from_pts', is_training)
if fts is None:
nn_fts_input = nn_fts_from_pts
else:
nn_fts_from_prev = tf.gather_nd(fts,
indices,
name=tag + 'nn_fts_from_prev')
nn_fts_input = tf.concat([nn_fts_from_pts, nn_fts_from_prev],
axis=-1,
name=tag + 'nn_fts_input')
P1 = tf.shape(nn_pts_local)[1]
dim1 = 3
if with_kernel_registering:
######################## preparing #########################
if with_feature_transformation:
nn_fts_input = tf.matmul(X_2_KK, nn_fts_input)
r_data = tf.reduce_sum(nn_pts_local * nn_pts_local,
axis=3,
keep_dims=True,
name=tag + 'kernel_pow')
######################## kernel-registering #########################
shape_id = 0
if kenel_initialization_method == 'random':
kernel_shape = tf.Variable(tf.random_uniform([K1, dim1],
minval=-0.5,
maxval=0.5,
dtype=tf.float32),
name=tag + 'kernel_shape' +
str(shape_id))
else:
kernel_shape = tf.Variable(tf.random_normal([K1, dim1],
mean=0.0,
stddev=1.0,
dtype=tf.float32),
name=tag + 'kernel_shape' +
str(shape_id))
kernel_shape_dis = tf.sqrt(tf.reduce_sum(kernel_shape * kernel_shape,
axis=1),
name=tag + 'kernel_shape_dis' +
str(shape_id))
kernel_shape_normal = scale * tf.div(
kernel_shape,
tf.reduce_max(kernel_shape_dis),
name=tag + 'kernel_shape_normal' + str(shape_id))
r_kernel = tf.reduce_sum(kernel_shape_normal * kernel_shape_normal,
axis=1,
keep_dims=True,
name=tag + 'kernel_pow' + str(shape_id))
reshape_data = tf.reshape(nn_pts_local, [N * P1 * K, dim1],
name=tag + 'reshape_kernel' + str(shape_id))
m = tf.reshape(tf.matmul(reshape_data,
tf.transpose(kernel_shape_normal)),
[N, P1, K, K1],
name=tag + 'mm' + str(shape_id))
dis_matrix = tf.transpose(r_data - 2 * m + tf.transpose(r_kernel),
perm=[0, 1, 3, 2],
name=tag + 'dis_matrix' + str(shape_id))
coef_matrix = tf.exp(tf.div(-dis_matrix, sigma),
name=tag + 'coef_matrix' + str(shape_id))
#coef_matrix = tf.transpose(r_data-2*m+tf.transpose(r_kernel),perm=[0,1,3,2],name=tag+'coef_matrix'+str(shape_id))
if with_kernel_shape_comparison:
coef_global = tf.reduce_sum(
coef_matrix, axis=[2, 3], keep_dims=True) / K
coef_normal = coef_global * tf.div(
coef_matrix,
tf.reduce_sum(coef_matrix, axis=3, keep_dims=True),
name=tag + 'coef_normal' + str(shape_id))
else:
coef_normal = tf.div(coef_matrix,
tf.reduce_sum(coef_matrix,
axis=3,
keep_dims=True),
name=tag + 'coef_normal' + str(shape_id))
fts_X = tf.matmul(coef_normal,
nn_fts_input,
name=tag + 'fts_X' + str(shape_id))
###################################################################
fts_conv = pf.separable_conv2d(fts_X,
math.ceil(mm * C / kernel_num),
tag + 'fts_conv' + str(shape_id),
is_training, (1, K1),
depth_multiplier=depth_multiplier)
fts_conv_3d = tf.squeeze(fts_conv,
axis=2,
name=tag + 'fts_conv_3d' + str(shape_id))
for shape_id in range(kernel_num - 1):
shape_id = shape_id + 1
if kenel_initialization_method == 'random':
kernel_shape = tf.Variable(tf.random_uniform([K1, dim1],
minval=-0.5,
maxval=0.5,
dtype=tf.float32),
name=tag + 'kernel_shape' +
str(shape_id))
else:
kernel_shape = tf.Variable(tf.random_normal([K1, dim1],
mean=0.0,
stddev=1.0,
dtype=tf.float32),
name=tag + 'kernel_shape' +
str(shape_id))
kernel_shape_dis = tf.sqrt(
tf.reduce_sum(kernel_shape * kernel_shape, axis=1),
name=tag + 'kernel_shape_dis' + str(shape_id))
kernel_shape_normal = scale * tf.div(
kernel_shape,
tf.reduce_max(kernel_shape_dis),
name=tag + 'kernel_shape_normal' + str(shape_id))
r_kernel = tf.reduce_sum(kernel_shape_normal * kernel_shape_normal,
axis=1,
keep_dims=True,
name=tag + 'kernel_pow' + str(shape_id))
reshape_data = tf.reshape(nn_pts_local, [N * P1 * K, dim1],
name=tag + 'reshape_kernel' +
str(shape_id))
m = tf.reshape(tf.matmul(reshape_data,
tf.transpose(kernel_shape_normal)),
[N, P1, K, K1],
name=tag + 'mm' + str(shape_id))
dis_matrix = tf.transpose(r_data - 2 * m + tf.transpose(r_kernel),
perm=[0, 1, 3, 2],
name=tag + 'dis_matrix' + str(shape_id))
coef_matrix = tf.exp(tf.div(-dis_matrix, sigma),
name=tag + 'coef_matrix' + str(shape_id))
#coef_matrix = tf.transpose(r_data-2*m+tf.transpose(r_kernel),perm=[0,1,3,2],name=tag+'coef_matrix'+str(shape_id))
if with_kernel_shape_comparison:
coef_global = tf.reduce_sum(
coef_matrix, axis=[2, 3], keep_dims=True) / K
coef_normal = coef_global * tf.div(
coef_matrix,
tf.reduce_sum(coef_matrix, axis=3, keep_dims=True),
name=tag + 'coef_normal' + str(shape_id))
else:
coef_normal = tf.div(coef_matrix,
tf.reduce_sum(coef_matrix,
axis=3,
keep_dims=True),
name=tag + 'coef_normal' + str(shape_id))
fts_X = tf.matmul(coef_normal,
nn_fts_input,
name=tag + 'fts_X' + str(shape_id))
###################################################################
fts_conv = pf.separable_conv2d(fts_X,
math.ceil(mm * C / kernel_num),
tag + 'fts_conv' + str(shape_id),
is_training, (1, K1),
depth_multiplier=depth_multiplier)
fts_conv_3d = tf.concat(
[fts_conv_3d, tf.squeeze(fts_conv, axis=2)],
axis=-1,
name=tag + 'fts_conv_3d' + str(shape_id))
else:
fts_X = nn_fts_input
fts_conv = pf.separable_conv2d(fts_X,
C,
tag + 'fts_conv',
is_training, (1, K),
depth_multiplier=depth_multiplier)
fts_conv_3d = tf.squeeze(fts_conv, axis=2, name=tag + 'fts_conv_3d')
if with_global:
fts_global_0 = pf.dense(qrs, C // 4, tag + 'fts_global_0', is_training)
fts_global = pf.dense(fts_global_0, C // 4, tag + 'fts_global',
is_training)
return tf.concat([fts_global, fts_conv_3d],
axis=-1,
name=tag + 'fts_conv_3d_with_global')
else:
return fts_conv_3d
def point_conv_transpose(pts_in,
fts_in,
pts_out,
fts_out_skipped,
indices,
K,
D,
C,
is_training,
activation=tf.nn.relu,
bn=True,
sorting_method='cxyz',
center_patch=True,
tconv='fc',
tag='point_deconv'):
'''
pts_in: xyz points, Tensor, (batch_size, npoint, 3)
fts_in: features, Tensor, (batch_size, npoint, channel_fts)
pts_out: xyz points, Tensor, (batch_size, n_pts_out, 3), NOTE: pts_out is denser than pts.
K: number of neighbors for one query, scalar, int
D: dilation rate
C: output feature channel number
tconv: fc or t_conv
in the comments:
N - batch size
P - number of pts_in
K - number of neighbors
'''
with tf.variable_scope(tag):
batch_size = pts_out.shape[0]
n_pts_out = pts_out.shape[1]
n_pts_in = pts_in.shape[1]
if indices is None:
_, indices_dilated = pf.knn_indices_general(
pts_in, pts_out, K * D, True)
indices = indices_dilated[:, :, ::D, :]
if sorting_method is not None:
indices = pf.sort_points(pts_out, indices,
sorting_method) # (N, P, K, 2)
# move the gathered pts_out to local frame or not
nn_pts_out = tf.gather_nd(pts_out, indices,
name='nn_pts_out') # (N, P, K, 3)
if center_patch:
nn_pts_out_center = tf.expand_dims(
pts_in, axis=2, name='nn_pts_out_center') # (N, P, 1, 3)
nn_pts_out_local = tf.subtract(
nn_pts_out, nn_pts_out_center,
name='nn_pts_out_local') # (N, P, K, 3)
else:
nn_pts_out_local = nn_pts_out
if fts_in is None:
nn_fts_input = pts_in # (N, P, 3)
else:
# NOTE: do not concat pts_in with fts_in now, since fts_in already contains information of pts_in
nn_fts_input = fts_in # (N, P, C_fts_in)
nn_fts_input = tf.expand_dims(nn_fts_input,
axis=2) # (N, P, 1, C_fts_in)
# using fc
if tconv == 'fc':
fts_fc = pf.dense(nn_fts_input, K * C, 'fc_fts_prop',
is_training) # (N, P, 1, K*C)
fts_fc = tf.reshape(fts_fc,
(batch_size, n_pts_in, K, C)) # (N, P, K, C)
elif tconv == 't_conv':
fts_fc = pf.conv2d_transpose(nn_fts_input,
C,
tag + '_fts_deconv',
is_training,
kernel_size=(1, K),
strides=(1, 1),
with_bn=bn,
activation=activation) # (N, P, K, C)
if fts_out_skipped is not None:
nn_fts_out_skipped = tf.gather_nd(fts_out_skipped,
indices,
name='nn_fts_out_from_skip')
fts_all_concat = tf.concat(
[fts_fc, nn_fts_out_skipped, nn_pts_out_local],
axis=-1) # (N, P, K, C+3+channel_skipped)
else:
fts_all_concat = tf.concat([fts_fc, nn_pts_out_local],
axis=-1) # (N, P, K, C+3)
fts_pts_conv = pf.conv2d(
fts_all_concat,
C,
'fts_pts_conv1D',
is_training,
kernel_size=(1, 1),
strides=(1, 1),
with_bn=bn,
activation=activation) # (N, P, K, C+3+[ch_skip]) -> (N, P, K, C)
# summing up feature for each point
fts_dense = tf.scatter_nd(
tf.reshape(indices, (-1, 2)), tf.reshape(fts_pts_conv, (-1, C)),
(batch_size, n_pts_out, C)) # (N, N_PTS_OUT, c)
return fts_dense
def point_conv_transpose_2(pts_in,
fts_in,
pts_out,
fts_out_skipped,
indices,
K,
D,
C,
is_training,
activation=tf.nn.relu,
bn=True,
sorting_method='cxyz',
tag='point_conv'):
'''
pts_in: xyz points, Tensor, (batch_size, npoint, 3)
fts_in: features, Tensor, (batch_size, npoint, channel_fts)
pts_out: xyz points, Tensor, (batch_size, n_pts_out, 3), NOTE: pts_out is denser than pts.
K: number of neighbors for one query, scalar, int
D: dilation rate
C: output feature channel number
in the comments:
N - batch size
P - number of pts_in
K - number of neighbors
'''
batch_size = pts_out.shape[0]
n_pts_out = pts_out.shape[1]
n_pts_in = pts_in.shape[1]
if indices is None:
_, indices_dilated = pf.knn_indices_general(pts_in, pts_out, K * D,
True)
indices = indices_dilated[:, :, ::D, :]
if sorting_method is not None:
indices = pf.sort_points(pts_out, indices,
sorting_method) # (N, P, K, 2)
# move the gathered pts_out to local frame
nn_pts_out = tf.gather_nd(pts_out, indices,
name=tag + 'nn_pts_out') # (N, P, K, 3)
nn_pts_out_center = tf.expand_dims(pts_in,
axis=2,
name=tag +
'nn_pts_out_center') # (N, P, 1, 3)
nn_pts_out_local = tf.subtract(nn_pts_out,
nn_pts_out_center,
name=tag +
'nn_pts_out_local') # (N, P, K, 3)
if fts_in is None:
nn_fts_input = pts_in # (N, P, 3)
#print('Error: fts_in is None.')
else:
#nn_fts_input = tf.concat([pts_in, fts_in], axis=-1, name=tag + 'nn_fts_input') # (N, P, 3+C_fts_in)
nn_fts_input = fts_in # (N, P, C_fts_in)
nn_fts_input = tf.expand_dims(nn_fts_input, axis=2) # (N, P, 1, C_fts_in)
# TODO: maybe try other conv1d? fc? there
fts_conv = pf.conv2d_transpose(
nn_fts_input,
C,
tag + 'fts_deconv',
is_training,
kernel_size=(1, 1),
strides=(1, K),
with_bn=bn,
activation=activation) # (N, P, 1, C_fts_in) -> (N, P, K, C)
if fts_out_skipped is not None:
nn_fts_out_skipped = tf.gather_nd(fts_out_skipped,
indices,
name=tag + 'nn_fts_out_from_skip')
fts_all_concat = tf.concat(
[fts_conv, nn_fts_out_skipped, nn_pts_out_local],
axis=-1) # (N, P, K, C+3+channel_skipped)
else:
fts_all_concat = tf.concat([fts_conv, nn_pts_out_local],
axis=-1) # (N, P, K, C+3)
fts_pts_conv = pf.conv2d(
fts_all_concat,
C,
tag + 'fts_pts_conv1D',
is_training,
kernel_size=(1, 1),
strides=(1, 1),
with_bn=bn,
activation=activation) # (N, P, K, C+3_ch_skip?) -> (N, P, K, C)
# summing up feature for each point
fts_dense = tf.scatter_nd(tf.reshape(indices, (-1, 2)),
tf.reshape(fts_pts_conv, (-1, C)),
(batch_size, n_pts_out, C)) # (N, N_PTS_OUT, c)
return fts_dense
def point_conv_transpose_1(pts_in,
fts_in,
pts_out,
indices,
K,
D,
C,
is_training,
activation=tf.nn.relu,
bn=True,
sorting_method='cxyz',
tag='point_conv'):
'''
pts_in: xyz points, Tensor, (batch_size, npoint, 3)
fts_in: features, Tensor, (batch_size, npoint, channel_fts)
pts_out: xyz points, Tensor, (batch_size, n_pts_out, 3), NOTE: pts_out is denser than pts.
K: number of neighbors for one query, scalar, int
D: dilation rate
C: output feature channel number
in the comments:
N - batch size
P - number of pts_in
K - number of neighbors
'''
batch_size = pts_out.shape[0]
n_pts_out = pts_out.shape[1]
n_pts_in = pts_in.shape[1]
if indices == None:
_, indices_dilated = pf.knn_indices_general(pts_in, pts_out, K * D,
True)
indices = indices_dilated[:, :, ::D, :]
if sorting_method is not None:
indices = pf.sort_points(pts_out, indices,
sorting_method) # (N, P, K, 2)
# move the gathered pts_out to local frame
nn_pts = tf.gather_nd(pts_out, indices,
name=tag + '_nn_pts') # (N, P, K, 3)
nn_pts_center = tf.expand_dims(pts_in, axis=2,
name=tag + '_nn_pts_center') # (N, P, 1, 3)
nn_pts_local = tf.subtract(nn_pts,
nn_pts_center,
name=tag + '_nn_pts_local') # (N, P, K, 3)
if fts_in is None:
#nn_fts_input = pts_in # (N, P, 3)
print('Error: fts_in is None.')
else:
#nn_fts_input = tf.concat([pts_in, fts_in], axis=-1, name=tag + 'nn_fts_input') # (N, P, 3+C_fts_in)
nn_fts_input = fts_in # (N, P, C_fts_in)
nn_fts_input = tf.expand_dims(nn_fts_input, axis=2) # (N, P, 1, C_fts_in)
# TODO: maybe try other conv1d? fc? there
fts_conv = pf.conv2d_transpose(
nn_fts_input,
C,
tag + '_fts_deconv',
is_training,
kernel_size=(1, 1),
strides=(1, K),
with_bn=bn,
activation=activation) # (N, P, 1, C_fts_in) -> (N, P, K, C)
fts_conv_pts_out_concat = tf.concat([fts_conv, nn_pts_local],
axis=-1) # (N, P, K, C+3)
fts_pts_conv = pf.conv2d(
fts_conv_pts_out_concat,
C,
tag + '_fts_pts_conv1D',
is_training,
kernel_size=(1, 1),
strides=(1, 1),
with_bn=bn,
activation=activation) # (N, P, K, C+3) -> (N, P, K, C)
######## collect fts for points in the output #######################
##########################################################################
npoint_indices = tf.reshape(tf.range(n_pts_in), [-1, 1]) # (P, 1)
npoint_indices = tf.tile(npoint_indices, (1, K)) # (P, K)
npoint_indices = tf.expand_dims(npoint_indices, axis=0) # (1, P, K)
npoint_indices = tf.tile(npoint_indices, (batch_size, 1, 1)) # (N, P, K)
npoint_indices = tf.expand_dims(npoint_indices, axis=3) # (N, P, K, 1)
NPK_indices = tf.concat(
[
tf.expand_dims(indices[:, :, :, 0], axis=-1), npoint_indices,
tf.expand_dims(indices[:, :, :, 1], axis=-1)
],
axis=3
) # (N, P, K, 3) last dim is 3 for (batch_idx, point_idx, neighbor_idx)
NPK_indices = tf.expand_dims(NPK_indices, axis=3) # (N, P, K, 1, 3)
NPK_indices = tf.tile(NPK_indices, (1, 1, 1, C, 1)) # (N, P, K, C, 3)
channel_indices = tf.reshape(tf.range(C),
[1, 1, 1, -1, 1]) # (1, 1, 1, C, 1)
channel_indices = tf.tile(
channel_indices, (batch_size, n_pts_in, K, 1, 1)) # (N, P, K, C, 1)
final_indices = tf.concat(
[NPK_indices, channel_indices], axis=4
) # (N, P, K, C, 4) last dim: (batch_idx, rep_point_idx, neighbor_idx_in_pts_out, channel_idx)
######## reshape for constructing sparse tensor ################
final_indices = tf.reshape(final_indices, [-1, 4])
final_value = tf.reshape(fts_pts_conv, [-1])
# using sparse tensor and sparse ops
# TODO: in tf 1.4, the sparse_reduce_sum cannot infer the shape even if the input shape is known.
# so, try to use latest tf, a version of after Oct. 2018.
#fts_sparse = tf.SparseTensor(indices=tf.cast(final_indices, tf.int64), values=final_value, dense_shape=[batch_size, n_pts_in, n_pts_out, C]) # (N, P, N_PTS_OUT, C) in sparse representation
#fts_out = tf.sparse_reduce_sum(fts_sparse, axis=1) # (N, N_PTS_OUT, C)
#fts_dense = tf.sparse_tensor_to_dense(fts_sparse, default_value=0, validate_indices=False) # (N, P, N_PTS_OUT, C)
# using dense tensor, memory-consuming
fts_dense = tf.scatter_nd(
final_indices, final_value,
[batch_size, n_pts_in, n_pts_out, C]) # (N, P, N_PTS_OUT, C)
fts_out = tf.reduce_sum(fts_dense, axis=1) # (N, N_PTS_OUT, C)
return fts_out
def xconv(pts, fts, qrs, tag, N, K, D, P, C, C_pts_fts, is_training, with_X_transformation, depth_multiplier,
sorting_method=None, with_global=False):
'''
pts: batch points input
fts: batch feature inputs
qrs: points outputs
tag: tags
N: batch_size
K: k neighbors, like kernel size in normal conv2d
D: dilation rate
P: sample points size in one sample
C: input channel num
C_pts_fts: output channel num
'''
# get k*d near points'indices from pts according to qrs
_, indices_dilated = pf.knn_indices_general(qrs, pts, K * D, True)
# select k points' indices
indices = indices_dilated[:, :, ::D, :]
if sorting_method is not None:
indices = pf.sort_points(pts, indices, sorting_method)
# get the k points' data according to the indices
nn_pts = tf.gather_nd(pts, indices, name=tag + 'nn_pts') # (N, P, K, 3)
# expand qrs from (N, P, 3) to (N, P, 1, 3)
nn_pts_center = tf.expand_dims(qrs, axis=2, name=tag + 'nn_pts_center') # (N, P, 1, 3)
# remove qrs from pts
nn_pts_local = tf.subtract(nn_pts, nn_pts_center, name=tag + 'nn_pts_local') # (N, P, K, 3)
# Prepare features to be transformed
# step 2, MLP
nn_fts_from_pts_0 = pf.dense(nn_pts_local, C_pts_fts, tag + 'nn_fts_from_pts_0', is_training)
nn_fts_from_pts = pf.dense(nn_fts_from_pts_0, C_pts_fts, tag + 'nn_fts_from_pts', is_training)
# step 3, concat origin features with the new features
if fts is None:
nn_fts_input = nn_fts_from_pts
else:
nn_fts_from_prev = tf.gather_nd(fts, indices, name=tag + 'nn_fts_from_prev')
nn_fts_input = tf.concat([nn_fts_from_pts, nn_fts_from_prev], axis=-1, name=tag + 'nn_fts_input')
# step 4, according to the appendix, the second MLP is formes with two depthwise_conv layers
if with_X_transformation:
######################## X-transformation #########################
X_0 = pf.conv2d(nn_pts_local, K * K, tag + 'X_0', is_training, (1, K))
X_0_KK = tf.reshape(X_0, (N, P, K, K), name=tag + 'X_0_KK')
X_1 = pf.depthwise_conv2d(X_0_KK, K, tag + 'X_1', is_training, (1, K))
X_1_KK = tf.reshape(X_1, (N, P, K, K), name=tag + 'X_1_KK')
X_2 = pf.depthwise_conv2d(X_1_KK, K, tag + 'X_2', is_training, (1, K), activation=None)
X_2_KK = tf.reshape(X_2, (N, P, K, K), name=tag + 'X_2_KK')
# step 5
fts_X = tf.matmul(X_2_KK, nn_fts_input, name=tag + 'fts_X')
###################################################################
else:
fts_X = nn_fts_input
# step 6
fts_conv = pf.separable_conv2d(fts_X, C, tag + 'fts_conv', is_training, (1, K), depth_multiplier=depth_multiplier)
fts_conv_3d = tf.squeeze(fts_conv, axis=2, name=tag + 'fts_conv_3d')
if with_global:
fts_global_0 = pf.dense(qrs, C // 4, tag + 'fts_global_0', is_training, with_bn=False)
fts_global = pf.dense(fts_global_0, C // 4, tag + 'fts_global', is_training, with_bn=False)
#fts_global_0 = pf.dense(qrs, C // 4, tag + 'fts_global_0', is_training)
#fts_global = pf.dense(fts_global_0, C // 4, tag + 'fts_global', is_training)
return tf.concat([fts_global, fts_conv_3d], axis=-1, name=tag + 'fts_conv_3d_with_global')
else:
return fts_conv_3d
