def multi_segmentation_head(features, object_labels, config, dropout_prob):
# Boolean of training
training = dropout_prob < 0.99
with tf.variable_scope('head_unary_conv'):
# Get a feature for each point and for each object_class
nC = len(config.num_classes)
w = weight_variable([int(features.shape[1]), nC * config.first_features_dim])
features = conv_ops.unary_convolution(features, w)
features = leaky_relu(batch_norm(features,
config.use_batch_norm,
config.batch_norm_momentum,
training))
features = tf.reshape(features, [-1, nC, config.first_features_dim])
features = tf.transpose(features, [1, 0, 2])
with tf.variable_scope('softmax'):
# Get a logit for each point and for each object_class
maxC = np.max(config.num_classes)
w = weight_variable([nC, config.first_features_dim, maxC])
b = bias_variable([maxC])
all_logits = conv_ops.unary_convolution(features, w) + b
# Pool according to real object class
nd_inds = tf.stack((object_labels, tf.range(tf.shape(object_labels)[0])))
nd_inds = tf.transpose(nd_inds)
logits = tf.gather_nd(all_logits, nd_inds)
return logits
def side_4_head(features, config, dropout_prob):
"""
Logits generation head for side feature 4 (semantic segmentation mask)
:param features: [Point_num, feature_channel]
:param config
:param dropout_prob
:return logits: [Point_num, num_classes] (not normalized)
"""
# Boolean of training
training = dropout_prob < 0.99
# Unary conv (equivalent to fully connected for each pixel)
with tf.variable_scope('head_unary_conv_side4'):
w = weight_variable([int(features.shape[1]), config.num_classes])
features = conv_ops.unary_convolution(features, w)
features = leaky_relu(
batch_norm(features, config.use_batch_norm,
config.batch_norm_momentum, training))
# Softmax
with tf.variable_scope('softmax_side4'):
w = weight_variable([config.num_classes, config.num_classes])
b = bias_variable([config.num_classes])
logits = conv_ops.unary_convolution(features, w) + b
return logits
def resnetb_block(layer_ind, inputs, features, radius, fdim, config, training):
"""
Block performing a resnet bottleneck convolution (1conv > KPconv > 1conv + shortcut)
"""
with tf.variable_scope('conv1'):
w = weight_variable([int(features.shape[1]), fdim // 2])
x = conv_ops.unary_convolution(features, w)
x = leaky_relu(batch_norm(x,
config.use_batch_norm,
config.batch_norm_momentum,
training))
with tf.variable_scope('conv2'):
w = weight_variable([config.num_kernel_points, int(x.shape[1]), fdim // 2])
for i in range(2):
if i == 0:
x1 = KPConv(inputs['points'][layer_ind],
inputs['points'][layer_ind],
inputs['neighbors'][layer_ind],
x,
w,
radius,
config)
x1 = KPConv(inputs['points'][layer_ind],
inputs['points'][layer_ind],
inputs['neighbors'][layer_ind],
x+x1,
w,
radius,
config)
x = leaky_relu(batch_norm(x1,
config.use_batch_norm,
config.batch_norm_momentum,
training))
with tf.variable_scope('conv3'):
w = weight_variable([int(x.shape[1]), 2 * fdim])
x = conv_ops.unary_convolution(x, w)
x = batch_norm(x,
config.use_batch_norm,
config.batch_norm_momentum,
training)
with tf.variable_scope('shortcut'):
if int(features.shape[1]) != 2 * fdim:
w = weight_variable([int(features.shape[1]), 2 * fdim])
shortcut = conv_ops.unary_convolution(features, w)
shortcut = batch_norm(shortcut,
config.use_batch_norm,
config.batch_norm_momentum,
training)
else:
shortcut = features
return leaky_relu(x + shortcut)
def resnetb_light_strided_block(layer_ind, inputs, features, radius, fdim, config, training):
"""
Block performing a strided resnet bottleneck convolution (shortcut is a maxpooling)
"""
with tf.variable_scope('conv1'):
if int(features.shape[1]) != fdim:
w = weight_variable([int(features.shape[1]), fdim])
x = conv_ops.unary_convolution(features, w)
x = batch_norm(x,
config.use_batch_norm,
config.batch_norm_momentum,
training)
else:
x = features
with tf.variable_scope('conv2'):
w = weight_variable([config.num_kernel_points, int(x.shape[1]), fdim])
x = KPConv(inputs['points'][layer_ind + 1],
inputs['points'][layer_ind],
inputs['pools'][layer_ind],
x,
w,
radius,
config)
x = leaky_relu(batch_norm(x,
config.use_batch_norm,
config.batch_norm_momentum,
training))
with tf.variable_scope('conv3'):
w = weight_variable([int(x.shape[1]), 2 * fdim])
x = conv_ops.unary_convolution(x, w)
x = batch_norm(x,
config.use_batch_norm,
config.batch_norm_momentum,
training)
with tf.variable_scope('shortcut'):
# Pool shortcuts to strided points TODO: max_pool or closest_pool ?
shortcut = ind_max_pool(features, inputs['pools'][layer_ind])
# shortcut = closest_pool(features, neighbors_indices)
# Regular upsample of the features if not the same dimension
if int(shortcut.shape[1]) != 2 * fdim:
w = weight_variable([int(shortcut.shape[1]), 2 * fdim])
shortcut = conv_ops.unary_convolution(shortcut, w)
shortcut = batch_norm(shortcut,
config.use_batch_norm,
config.batch_norm_momentum,
training)
return leaky_relu(x + shortcut)
def resnetb_upsample_block(layer_ind, inputs, features, radius, fdim, config, training):
"""
Block performing an upsampling resnet bottleneck convolution (shortcut is nearest interpolation)
"""
with tf.variable_scope('conv1'):
w = weight_variable([int(features.shape[1]), fdim // 2])
x = conv_ops.unary_convolution(features, w)
x = leaky_relu(batch_norm(x,
config.use_batch_norm,
config.batch_norm_momentum,
training))
with tf.variable_scope('conv2'):
w = weight_variable([config.num_kernel_points, int(x.shape[1]), fdim // 2])
x = KPConv(inputs['points'][layer_ind - 1],
inputs['points'][layer_ind],
inputs['upsamples'][layer_ind - 1],
x,
w,
radius,
config)
x = leaky_relu(batch_norm(x,
config.use_batch_norm,
config.batch_norm_momentum,
training))
with tf.variable_scope('conv3'):
w = weight_variable([int(x.shape[1]), 2 * fdim])
x = conv_ops.unary_convolution(x, w)
x = batch_norm(x,
config.use_batch_norm,
config.batch_norm_momentum,
training)
with tf.variable_scope('shortcut'):
# Pool shortcuts to strided points (nearest interpolation)
shortcut = closest_pool(features, inputs['upsamples'][layer_ind - 1])
# Regular upsample of the features if not the same dimension
if int(shortcut.shape[1]) != 2 * fdim:
w = weight_variable([int(shortcut.shape[1]), 2 * fdim])
shortcut = conv_ops.unary_convolution(shortcut, w)
shortcut = batch_norm(shortcut,
config.use_batch_norm,
config.batch_norm_momentum,
training)
return leaky_relu(x + shortcut)
def resnet_block(layer_ind, inputs, features, radius, fdim, config, training):
"""
Block performing a resnet double convolution (two convolution vgglike and a shortcut)
"""
with tf.variable_scope('conv1'):
w = weight_variable(
[config.num_kernel_points,
int(features.shape[1]), fdim])
x = KPConv(inputs['points'][layer_ind], inputs['points'][layer_ind],
inputs['neighbors'][layer_ind], features, w, radius, config)
x = leaky_relu(
batch_norm(x, config.use_batch_norm, config.batch_norm_momentum,
training))
with tf.variable_scope('conv2'):
w = weight_variable([config.num_kernel_points, int(x.shape[1]), fdim])
x = KPConv(inputs['points'][layer_ind], inputs['points'][layer_ind],
inputs['neighbors'][layer_ind], x, w, radius, config)
x = leaky_relu(
batch_norm(x, config.use_batch_norm, config.batch_norm_momentum,
training))
with tf.variable_scope('shortcut'):
if int(features.shape[1]) != fdim:
w = weight_variable([int(features.shape[1]), fdim])
shortcut = conv_ops.unary_convolution(features, w)
shortcut = batch_norm(shortcut, config.use_batch_norm,
config.batch_norm_momentum, training)
else:
shortcut = features
return leaky_relu(x + shortcut)
def unary_block(layer_ind, inputs, features, radius, fdim, config, training):
"""
Block performing a simple 1x1 convolution
"""
w = weight_variable([int(features.shape[1]), fdim])
x = conv_ops.unary_convolution(features, w)
x = leaky_relu(
batch_norm(x, config.use_batch_norm, config.batch_norm_momentum,
training))
return x
def segmentation_head(features, config, dropout_prob):
# Boolean of training
training = dropout_prob < 0.99
# Unary conv (equivalent to fully connected for each pixel)
with tf.variable_scope('head_unary_conv'):
w = weight_variable([int(features.shape[1]), config.first_features_dim])
features = conv_ops.unary_convolution(features, w)
features = leaky_relu(batch_norm(features,
config.use_batch_norm,
config.batch_norm_momentum,
training))
# Softmax
with tf.variable_scope('softmax'):
w = weight_variable([config.first_features_dim, config.num_classes])
b = bias_variable([config.num_classes])
logits = conv_ops.unary_convolution(features, w) + b
return logits
def forward(self, query_points, support_points, neighbors_indices,
features):
"""
This module performs a unary 1x1 convolution (same with MLP)
:param features: float32[n_points, in_fdim] - input features
:return: output_features float32[n_points, out_fdim]
"""
x = conv_ops.unary_convolution(features, self.weight)
if self.config.use_batch_norm:
x = self.relu(self.bn(x))
else:
x = self.relu(x)
return x
def assemble_decoder(inputs, config, dropout_prob, bottleneck_features, coarse,
double_fold):
"""
Assembles decoder architecture using folding operations
:param inputs:
:param config:
:param dropout_prob:
:param bottleneck_features:
:param coarse:
:return:
"""
# Boolean of training
training = dropout_prob < 0.99
with tf.variable_scope('folding'):
# Create tiled grid features
x = tf.linspace(-config.grid_scale, config.grid_scale,
config.grid_size)
y = tf.linspace(-config.grid_scale, config.grid_scale,
config.grid_size)
grid = tf.meshgrid(x, y)
grid = tf.expand_dims(tf.reshape(tf.stack(grid, axis=2), [-1, 2]), 0)
grid_feat = tf.tile(
grid, [tf.shape(bottleneck_features)[0], config.num_coarse, 1])
# Create tiled coarse point cloud features
point_feat = tf.tile(tf.expand_dims(coarse, 2),
[1, 1, config.grid_size**2, 1])
point_feat = tf.reshape(point_feat, [-1, config.num_fine, 3])
# Create tiled bottleneck features
global_feat = tf.tile(tf.expand_dims(bottleneck_features, 1),
[1, config.num_fine, 1])
feat = tf.concat([grid_feat, point_feat, global_feat], axis=2)
center = tf.tile(tf.expand_dims(coarse, 2),
[1, 1, config.grid_size**2, 1])
center = tf.reshape(center, [-1, config.num_fine, 3])
feat = tf.reshape(feat, [
-1,
grid_feat.shape[-1] + point_feat.shape[-1] + global_feat.shape[-1]
])
w = weight_variable([int(feat.shape[1]), 512])
x = conv_ops.unary_convolution(feat, w)
x = leaky_relu(
batch_norm(x, config.use_batch_norm, config.batch_norm_momentum,
training))
w = weight_variable([int(x.shape[1]), 512])
x = conv_ops.unary_convolution(x, w)
x = leaky_relu(
batch_norm(x, config.use_batch_norm, config.batch_norm_momentum,
training))
w = weight_variable([int(x.shape[1]), 3])
x = conv_ops.unary_convolution(x, w)
x = batch_norm(x, config.use_batch_norm, config.batch_norm_momentum,
training)
if double_fold:
x = tf.reshape(x, [-1, config.num_gt_points, 3])
x = tf.concat([x, global_feat], axis=2)
x = tf.reshape(x, [-1, x.shape[-1]])
w = weight_variable([int(x.shape[1]), 512])
x = conv_ops.unary_convolution(x, w)
x = leaky_relu(
batch_norm(x, config.use_batch_norm,
config.batch_norm_momentum, training))
w = weight_variable([int(x.shape[1]), 512])
x = conv_ops.unary_convolution(x, w)
x = leaky_relu(
batch_norm(x, config.use_batch_norm,
config.batch_norm_momentum, training))
w = weight_variable([int(x.shape[1]), 3])
x = conv_ops.unary_convolution(x, w)
x = batch_norm(x, config.use_batch_norm,
config.batch_norm_momentum, training)
x = tf.reshape(x, [-1, config.num_gt_points, 3])
fine = x + center
return fine
def inception_deformable_strided_block(layer_ind, inputs, features, radius, fdim, config, training):
"""
Block performing an inception style convolution combining rigid and deformable KPConv
(1conv > rigid) \
> CONCAT > 1conv + shortcut
(1conv > rigid > deform) /
"""
with tf.variable_scope('path1'):
with tf.variable_scope('unary'):
w = weight_variable([int(features.shape[1]), fdim // 2])
x1 = conv_ops.unary_convolution(features, w)
x1 = leaky_relu(batch_norm(x1,
config.use_batch_norm,
config.batch_norm_momentum,
training))
with tf.variable_scope('conv'):
w = weight_variable([config.num_kernel_points, int(x1.shape[1]), fdim // 2])
x1 = KPConv(inputs['points'][layer_ind+1],
inputs['points'][layer_ind],
inputs['pools'][layer_ind],
x1,
w,
radius,
config)
with tf.variable_scope('path2'):
with tf.variable_scope('unary'):
w = weight_variable([int(features.shape[1]), fdim // 2])
x2 = conv_ops.unary_convolution(features, w)
x2 = leaky_relu(batch_norm(x2,
config.use_batch_norm,
config.batch_norm_momentum,
training))
with tf.variable_scope('conv'):
w = weight_variable([config.num_kernel_points, int(x2.shape[1]), fdim // 2])
x2 = KPConv(inputs['points'][layer_ind+1],
inputs['points'][layer_ind],
inputs['pools'][layer_ind],
x2,
w,
radius,
config)
with tf.variable_scope('conv2_deform'):
w = weight_variable([config.num_kernel_points, int(x2.shape[1]), fdim // 2])
x2 = KPConv_deformable_v2(inputs['points'][layer_ind+1],
inputs['points'][layer_ind],
inputs['pools'][layer_ind],
x2,
w,
radius,
config)
with tf.variable_scope('concat'):
x = tf.concat([x1, x2], axis=1)
x = leaky_relu(batch_norm(x,
config.use_batch_norm,
config.batch_norm_momentum,
training))
with tf.variable_scope('unary'):
w = weight_variable([int(x.shape[1]), 2 * fdim])
x = conv_ops.unary_convolution(x, w)
x = batch_norm(x,
config.use_batch_norm,
config.batch_norm_momentum,
training)
with tf.variable_scope('shortcut'):
# Pool shortcuts to strided points TODO: max_pool or closest_pool ?
shortcut = ind_max_pool(features, inputs['pools'][layer_ind])
# shortcut = closest_pool(features, neighbors_indices)
# Regular upsample of the features if not the same dimension
if int(shortcut.shape[1]) != 2 * fdim:
w = weight_variable([int(shortcut.shape[1]), 2 * fdim])
shortcut = conv_ops.unary_convolution(shortcut, w)
shortcut = batch_norm(shortcut,
config.use_batch_norm,
config.batch_norm_momentum,
training)
return leaky_relu(x + shortcut)
def inception_deformable_block(layer_ind, inputs, features, radius, fdim, config, training):
"""
Block performing an inception style convolution combining rigid and deformable KPConv
(1conv > rigid) \
> CONCAT > 1conv + shortcut
(1conv > rigid > deform) /
"""
with tf.variable_scope('path1'):
with tf.variable_scope('unary'):
w = weight_variable([int(features.shape[1]), fdim // 2])
x1 = conv_ops.unary_convolution(features, w)
x1 = leaky_relu(batch_norm(x1,
config.use_batch_norm,
config.batch_norm_momentum,
training))
with tf.variable_scope('conv'):
w = weight_variable([config.num_kernel_points, int(x1.shape[1]), fdim // 2])
x1 = KPConv(inputs['points'][layer_ind],
inputs['points'][layer_ind],
inputs['neighbors'][layer_ind],
x1,
w,
radius,
config)
with tf.variable_scope('path2'):
with tf.variable_scope('unary'):
w = weight_variable([int(features.shape[1]), fdim // 2])
x2 = conv_ops.unary_convolution(features, w)
x2 = leaky_relu(batch_norm(x2,
config.use_batch_norm,
config.batch_norm_momentum,
training))
with tf.variable_scope('conv'):
w = weight_variable([config.num_kernel_points, int(x2.shape[1]), fdim // 2])
x2 = KPConv(inputs['points'][layer_ind],
inputs['points'][layer_ind],
inputs['neighbors'][layer_ind],
x2,
w,
radius,
config)
with tf.variable_scope('conv2_deform'):
w = weight_variable([config.num_kernel_points, int(x2.shape[1]), fdim // 2])
x2 = KPConv_deformable_v2(inputs['points'][layer_ind],
inputs['points'][layer_ind],
inputs['neighbors'][layer_ind],
x2,
w,
radius,
config)
with tf.variable_scope('concat'):
x = tf.concat([x1, x2], axis=1)
x = leaky_relu(batch_norm(x,
config.use_batch_norm,
config.batch_norm_momentum,
training))
with tf.variable_scope('unary'):
w = weight_variable([int(x.shape[1]), 2 * fdim])
x = conv_ops.unary_convolution(x, w)
x = batch_norm(x,
config.use_batch_norm,
config.batch_norm_momentum,
training)
with tf.variable_scope('shortcut'):
if int(features.shape[1]) != 2 * fdim:
w = weight_variable([int(features.shape[1]), 2 * fdim])
shortcut = conv_ops.unary_convolution(features, w)
shortcut = batch_norm(shortcut,
config.use_batch_norm,
config.batch_norm_momentum,
training)
else:
shortcut = features
return leaky_relu(x + shortcut)
