def while_body(i, est_pose, est_inputs):
with tf.variable_scope('encoder_0', reuse=tf.AUTO_REUSE):
features = mlp_conv(est_inputs, [128, 256], bn=self.bn)
features_global = tf.reduce_max(features, axis=1, keep_dims=True, name='maxpool_0')
features = tf.concat([features, tf.tile(features_global, [1, tf.shape(inputs)[1], 1])], axis=2)
with tf.variable_scope('encoder_1', reuse=tf.AUTO_REUSE):
features = mlp_conv(features, [512, 1024], bn=self.bn)
features = tf.reduce_max(features, axis=1, name='maxpool_1')
with tf.variable_scope('fc', reuse=tf.AUTO_REUSE):
if self.rot_representation == 'quat':
est_pose_rep_i = mlp(features, [1024, 1024, 512, 512, 256, 7], bn=self.bn)
elif self.rot_representation == '6dof':
est_pose_rep_i = mlp(features, [1024, 1024, 512, 512, 256, 9], bn=self.bn)
with tf.variable_scope('est', reuse=tf.AUTO_REUSE):
if self.rot_representation == 'quat':
t = tf.expand_dims(est_pose_rep_i[:, 4:], axis=2)
q = est_pose_rep_i[:, 0:4]
R = quat2rotm_tf(q)
elif self.rot_representation == '6dof':
t = tf.expand_dims(est_pose_rep_i[:, 6:], axis=2)
mat6d = est_pose_rep_i[:, 0:6]
R = mat6d2rotm_tf(mat6d)
est_pose_T_i = tf.concat([
tf.concat([R, t], axis=2),
tf.concat([tf.zeros([B, 1, 3]), tf.ones([B, 1, 1])], axis=2)],
axis=1)
est_inputs = transform_tf(est_inputs, est_pose_T_i)
est_pose = tf.linalg.matmul(est_pose_T_i, est_pose)
return [tf.add(i, 1), est_pose, est_inputs]
def create_encoder(self, inputs):
with tf.variable_scope('encoder_0', reuse=tf.AUTO_REUSE):
features = mlp_conv(inputs, [128, 256])
features_global = tf.reduce_max(features, axis=1, keep_dims=True, name='maxpool_0')
features = tf.concat([features, tf.tile(features_global, [1, tf.shape(inputs)[1], 1])], axis=2)
with tf.variable_scope('encoder_1', reuse=tf.AUTO_REUSE):
features = mlp_conv(features, [512, 1024])
features = tf.reduce_max(features, axis=1, name='maxpool_1')
return features
def encoder(inputs, embed_size=1024):
with tf.variable_scope('encoder_0', reuse=tf.AUTO_REUSE):
features = mlp_conv(inputs, [128, 256], bn=None, bn_params=None)
features_global = tf.reduce_max(features, axis=1, keepdims=True, name='maxpool_0')
features = tf.concat([features, tf.tile(features_global, [1, tf.shape(inputs)[1], 1])], axis=2)
with tf.variable_scope('encoder_1', reuse=tf.AUTO_REUSE):
features = mlp_conv(features, [512, embed_size], bn=None, bn_params=None)
features = tf.reduce_max(features, axis=1, name='maxpool_1')
return features_global,features
def create_decoder(self, features):
with tf.variable_scope('decoder', reuse=tf.AUTO_REUSE):
x = tf.linspace(-self.grid_scale, self.grid_scale, self.grid_size)
y = tf.linspace(-self.grid_scale, self.grid_scale, self.grid_size)
grid = tf.meshgrid(x, y)
grid = tf.reshape(tf.stack(grid, axis=2), [-1, 2])
grid = tf.tile(tf.expand_dims(grid, 0), [features.shape[0], 1, 1])
features = tf.tile(tf.expand_dims(features, 1), [1, self.num_output_points, 1])
with tf.variable_scope('folding_1'):
fold1 = mlp_conv(tf.concat([features, grid], axis=2), [512, 512, 3])
with tf.variable_scope('folding_2'):
fold2 = mlp_conv(tf.concat([features, fold1], axis=2), [512, 512, 3])
return fold1, fold2
def create_decoder(code,
inputs,
step_ratio,
num_extract=512,
mean_feature=None):
with tf.variable_scope('decoder', reuse=tf.AUTO_REUSE):
level0 = tf_util.mlp(code, [1024, 1024, 512 * 3],
bn=None,
bn_params=None,
name='coarse') # ,name='coarse'
level0 = tf.tanh(level0)
level0 = tf.reshape(level0, [-1, 512, 3])
coarse = level0
input_fps = symmetric_sample(inputs, int(num_extract / 2))
level0 = tf.concat([input_fps, level0], 1)
if num_extract > 512:
level0 = gather_point(level0, farthest_point_sample(1024, level0))
for i in range(int(math.log2(step_ratio))):
num_fine = 2**(i + 1) * 1024
grid = tf_util.gen_grid_up(2**(i + 1))
grid = tf.expand_dims(grid, 0)
grid_feat = tf.tile(grid, [level0.shape[0], 1024, 1])
point_feat = tf.tile(tf.expand_dims(level0, 2), [1, 1, 2, 1])
point_feat = tf.reshape(point_feat, [-1, num_fine, 3])
global_feat = tf.tile(tf.expand_dims(code, 1), [1, num_fine, 1])
mean_feature_use = tf.contrib.layers.fully_connected(
mean_feature, 128, activation_fn=tf.nn.relu, scope='mean_fc')
mean_feature_use = tf.expand_dims(mean_feature_use, 1)
mean_feature_use = tf.tile(mean_feature_use, [1, num_fine, 1])
feat = tf.concat(
[grid_feat, point_feat, global_feat, mean_feature_use], axis=2)
feat1 = tf_util.mlp_conv(feat, [128, 64],
bn=None,
bn_params=None,
name='up_branch')
feat1 = tf.nn.relu(feat1)
feat2 = contract_expand_operation(feat1, 2)
feat = feat1 + feat2
fine = tf_util.mlp_conv(
feat, [512, 512, 3], bn=None, bn_params=None,
name='fine') + point_feat
level0 = fine
return coarse, fine
def create_decoder(self, features):
with tf.variable_scope('decoder', reuse=tf.AUTO_REUSE):
coarse = mlp(features, [1024, 1024, self.num_coarse * 3])
coarse = tf.reshape(coarse, [-1, self.num_coarse, 3])
with tf.variable_scope('folding', reuse=tf.AUTO_REUSE):
grid = tf.meshgrid(tf.linspace(-0.05, 0.05, self.grid_size),
tf.linspace(-0.05, 0.05, self.grid_size))
grid = tf.expand_dims(tf.reshape(tf.stack(grid, axis=2), [-1, 2]),
0)
grid_feat = tf.tile(grid, [features.shape[0], self.num_coarse, 1])
point_feat = tf.tile(tf.expand_dims(coarse, 2),
[1, 1, self.grid_size**2, 1])
point_feat = tf.reshape(point_feat, [-1, self.num_fine, 3])
global_feat = tf.tile(tf.expand_dims(features, 1),
[1, self.num_fine, 1])
feat = tf.concat([grid_feat, point_feat, global_feat], axis=2)
center = tf.tile(tf.expand_dims(coarse, 2),
[1, 1, self.grid_size**2, 1])
center = tf.reshape(center, [-1, self.num_fine, 3])
fine = mlp_conv(feat, [512, 512, 3]) + center
return coarse, fine
def decoder(inputs, features, step_ratio=16, num_fine=16 * 1024):
num_coarse=1024
assert num_fine == num_coarse * step_ratio
with tf.variable_scope('decoder', reuse=tf.AUTO_REUSE):
coarse = mlp(features, [1024, 1024, num_coarse * 3], bn=None, bn_params=None)
coarse = tf.reshape(coarse, [-1, num_coarse, 3])
p1_idx = farthest_point_sample(512, coarse)
coarse_1 = gather_point(coarse, p1_idx)
input_fps = symmetric_sample(inputs, int(512 / 2))
coarse = tf.concat([input_fps, coarse_1], 1)
with tf.variable_scope('folding', reuse=tf.AUTO_REUSE):
if not step_ratio ** .5 % 1 == 0:
grid = gen_1d_grid(step_ratio)
else:
grid = gen_grid(np.round(np.sqrt(step_ratio)).astype(np.int32))
grid = tf.expand_dims(grid, 0)
grid_feat = tf.tile(grid, [features.shape[0], num_coarse, 1])
point_feat = tf.tile(tf.expand_dims(coarse, 2), [1, 1, step_ratio, 1])
point_feat = tf.reshape(point_feat, [-1, num_fine, 3])
global_feat = tf.tile(tf.expand_dims(features, 1), [1, num_fine, 1])
feat = tf.concat([grid_feat, point_feat, global_feat], axis=2)
fine = mlp_conv(feat, [512, 512, 3], bn=None, bn_params=None) + point_feat
return coarse, fine
def create_encoder(inputs, embed_size=1024):
inputs_new = inputs
features = tf_util.mlp_conv(inputs_new, [128, 256],
bn=None,
bn_params=None,
name='encoder_0')
features_global = tf.reduce_max(features,
axis=1,
keepdims=True,
name='maxpool_0')
features = tf.concat(
[features,
tf.tile(features_global, [1, tf.shape(inputs_new)[1], 1])],
axis=2)
features = tf_util.mlp_conv(features, [512, embed_size],
bn=None,
bn_params=None,
name='encoder_1')
features = tf.reduce_max(features, axis=1, name='maxpool_1')
return features
def patch_dection(point_cloud, divide_ratio=1):
with tf.variable_scope('patch', reuse=tf.AUTO_REUSE):
l0_xyz = point_cloud
l0_points = None
num_point = point_cloud.get_shape()[1].value
l1_xyz, l1_points = tf_util.pointnet_sa_module_msg(l0_xyz, l0_points, int(num_point/8), [0.1, 0.2, 0.4], [16, 32, 128],
[[32//divide_ratio, 32//divide_ratio, 64//divide_ratio], \
[64//divide_ratio, 64//divide_ratio, 128//divide_ratio], \
[64//divide_ratio, 96//divide_ratio, 128//divide_ratio]],
scope='layer1', use_nchw=False)
patch_values = tf_util.mlp_conv(l1_points, [1],
bn=None,
bn_params=None,
name='patch')
return patch_values
def completion(self, inputs, is_training):
num_point = inputs.get_shape()[1].value
l0_xyz = inputs[:,:,0:3]
l0_points = None
is_training = is_training
bradius = 1.0
use_bn = False
use_ibn = False
bn_decay = 0.95
up_ratio = 8
self.grid_size = 2
self.num_coarse = int(num_point * up_ratio / 4)
with tf.variable_scope('encoder_0', reuse=tf.AUTO_REUSE):
l1_xyz, l1_points, l1_indices = pointnet_sa_module(l0_xyz, l0_points, npoint=num_point,
radius=bradius * 0.05, bn=use_bn, ibn=use_ibn,
nsample=32, mlp=[32, 32, 64], 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=num_point / 2,
radius=bradius * 0.1, bn=use_bn, ibn=use_ibn,
nsample=32, mlp=[64, 64, 128], 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=num_point / 4,
radius=bradius * 0.2, bn=use_bn, ibn=use_ibn,
nsample=32, mlp=[128, 128, 256], mlp2=None,
group_all=False,
is_training=is_training, bn_decay=bn_decay,
scope='layer3')
l4_xyz, l4_points, l4_indices = pointnet_sa_module(l3_xyz, l3_points, npoint=num_point / 8,
radius=bradius * 0.3, bn=use_bn, ibn=use_ibn,
nsample=32, mlp=[256, 256, 512], mlp2=None,
group_all=False,
is_training=is_training, bn_decay=bn_decay,
scope='layer4')
l5_xyz, l5_points, l5_indices = pointnet_sa_module(l4_xyz, l4_points, npoint=num_point / 16,
radius=bradius * 0.4, bn=use_bn, ibn=use_ibn,
nsample=32, mlp=[512, 512, 1024], mlp2=None,
group_all=False,
is_training=is_training, bn_decay=bn_decay,
scope='layer5')
gl_xyz, gl_points, gl_indices = pointnet_sa_module(l5_xyz, l5_points, npoint=1,
radius=bradius * 0.3, bn=use_bn, ibn=use_ibn,
nsample=32, mlp=[512, 512, 1024], mlp2=None,
group_all=True,
is_training=is_training, bn_decay=bn_decay,
scope='layer6')
gl_feature = tf.reduce_max(gl_points, axis=1)
print('gl_feature', gl_feature)
# Feature Propagation layers
up_gl_points = pointnet_fp_module(l0_xyz, gl_xyz, None, gl_points, [64], is_training, bn_decay,
scope='fa_layer0', bn=use_bn, ibn=use_ibn)
up_l5_points = pointnet_fp_module(l0_xyz, l5_xyz, None, l5_points, [64], is_training, bn_decay,
scope='fa_layer1', bn=use_bn, ibn=use_ibn)
up_l4_points = pointnet_fp_module(l0_xyz, l4_xyz, None, l4_points, [64], is_training, bn_decay,
scope='fa_layer2', bn=use_bn, ibn=use_ibn)
up_l3_points = pointnet_fp_module(l0_xyz, l3_xyz, None, l3_points, [64], is_training, bn_decay,
scope='fa_layer3', bn=use_bn, ibn=use_ibn)
up_l2_points = pointnet_fp_module(l0_xyz, l2_xyz, None, l2_points, [64], is_training, bn_decay,
scope='fa_layer4', bn=use_bn, ibn=use_ibn)
###concat feature
with tf.variable_scope('up_layer', reuse=tf.AUTO_REUSE):
new_points_list = []
for i in range(up_ratio):
if i>3:
transform = input_transform_net(l0_xyz, is_training, bn_decay, K=3)
xyz_transformed = tf.matmul(l0_xyz, transform)
concat_feat = tf.concat([up_gl_points, up_gl_points-up_l5_points, up_gl_points-up_l4_points, up_gl_points-up_l3_points, up_gl_points-up_l2_points, up_gl_points-l1_points, xyz_transformed],
axis=-1)
print('concat_feat1', concat_feat)
else:
concat_feat = tf.concat([up_gl_points, up_l5_points, up_l4_points, up_l3_points, up_l2_points, l1_points, l0_xyz],
axis=-1)
print('concat_feat2', concat_feat)
#concat_feat = tf.concat([up_l4_points, up_l3_points, up_l2_points, l1_points, l0_xyz], axis=-1)
concat_feat = tf.expand_dims(concat_feat, axis=2)
concat_feat = tf_util2.conv2d(concat_feat, 256, [1, 1],
padding='VALID', stride=[1, 1],
bn=False, is_training=is_training,
scope='fc_layer0_%d' % (i), bn_decay=bn_decay)
new_points = tf_util2.conv2d(concat_feat, 128, [1, 1],
padding='VALID', stride=[1, 1],
bn=use_bn, is_training=is_training,
scope='conv_%d' % (i),
bn_decay=bn_decay)
new_points_list.append(new_points)
net = tf.concat(new_points_list, axis=1)
coord_feature = tf_util2.conv2d(net, 64, [1, 1],
padding='VALID', stride=[1, 1],
bn=False, is_training=is_training,
scope='fc_layer1', bn_decay=bn_decay)
coord = tf_util2.conv2d(coord_feature, 3, [1, 1],
padding='VALID', stride=[1, 1],
bn=False, is_training=is_training,
scope='fc_layer2', bn_decay=bn_decay,
activation_fn=None, weight_decay=0.0) # B*(2N)*1*3
coarse_highres = tf.squeeze(coord, [2]) # B*(2N)*3
coord_feature = tf.squeeze(coord_feature, [2])
fps_idx = farthest_point_sample(int(self.num_fine)/2, coarse_highres)
coord_feature = gather_point(coord_feature, fps_idx)
coarse_fps = gather_point(coarse_highres, fps_idx)
coord_feature = tf.expand_dims(coord_feature, 2)
print('coord_feature', coord, coord_feature)
score = tf_util2.conv2d(coord_feature, 16, [1, 1],
padding='VALID', stride=[1, 1],
bn=False, is_training=is_training,
scope='fc_layer3', bn_decay=bn_decay)
score = tf_util2.conv2d(score, 8, [1, 1],
padding='VALID', stride=[1, 1],
bn=False, is_training=is_training,
scope='fc_layer4', bn_decay=bn_decay)
score = tf_util2.conv2d(score, 1, [1, 1],
padding='VALID', stride=[1, 1],
bn=False, is_training=is_training,
scope='fc_layer5', bn_decay=bn_decay)
score = tf.nn.softplus(score)
score = tf.squeeze(score, [2,3])
_, idx = tf.math.top_k(score, self.num_coarse)
coarse = gather_point(coarse_fps, idx)
coord_feature = tf.squeeze(coord_feature, [2])
coord_feature = gather_point(coord_feature, idx)
print('coarse', coord_feature, coarse)
with tf.variable_scope('folding', reuse=tf.AUTO_REUSE):
grid = tf.meshgrid(tf.linspace(-0.05, 0.05, self.grid_size), tf.linspace(-0.05, 0.05, self.grid_size))
print('grid:', grid)
grid = tf.expand_dims(tf.reshape(tf.stack(grid, axis=2), [-1, 2]), 0)
print('grid:', grid)
grid_feat = tf.tile(grid, [coarse.shape[0], self.num_coarse, 1])
print('grid_feat', grid_feat)
point_feat = tf.tile(tf.expand_dims(tf.concat([coarse, coord_feature], axis=-1), 2), [1, 1, self.grid_size ** 2, 1])
point_feat = tf.reshape(point_feat, [coarse.shape[0], self.num_fine, -1])
print('point_feat', point_feat)
global_feat = tf.tile(tf.expand_dims(gl_feature, 1), [1, self.num_fine, 1])
#print('global_feat', global_feat)
feat = tf.concat([grid_feat, point_feat, global_feat], axis=2)
print('feat:', feat)
center = tf.tile(tf.expand_dims(coarse, 2), [1, 1, self.grid_size ** 2, 1])
center = tf.reshape(center, [-1, self.num_fine, 3])
print('center', center)
fine = mlp_conv(feat, [512, 512, 3]) + center
print('fine:', fine)
return coarse_highres, coarse, fine
