def sample_and_group(npoint, radius, nsample, xyz, points, knn=False, use_xyz=True):
'''
Input:
npoint: int32
radius: float32
nsample: int32
xyz: (batch_size, ndataset, 3) TF tensor
points: (batch_size, ndataset, channel) TF tensor, if None will just use xyz as points
knn: bool, if True use kNN instead of radius search
use_xyz: bool, if True concat XYZ with local point features, otherwise just use point features
Output:
new_xyz: (batch_size, npoint, 3) TF tensor
new_points: (batch_size, npoint, nsample, 3+channel) TF tensor
idx: (batch_size, npoint, nsample) TF tensor, indices of local points as in ndataset points
grouped_xyz: (batch_size, npoint, nsample, 3) TF tensor, normalized point XYZs
(subtracted by seed point XYZ) in local regions
'''
new_xyz = gather_point(xyz, farthest_point_sample(npoint, xyz)) # (batch_size, npoint, 3)
if knn:
_,idx = knn_point(nsample, xyz, new_xyz)
else:
idx, pts_cnt = query_ball_point(radius, nsample, xyz, new_xyz)
grouped_xyz = group_point(xyz, idx) # (batch_size, npoint, nsample, 3)
grouped_xyz -= tf.tile(tf.expand_dims(new_xyz, 2), [1,1,nsample,1]) # translation normalization
if points is not None:
grouped_points = group_point(points, idx) # (batch_size, npoint, nsample, channel)
if use_xyz:
new_points = tf.concat([grouped_xyz, grouped_points], axis=-1) # (batch_size, npoint, nample, 3+channel)
else:
new_points = grouped_points
else:
new_points = grouped_xyz
return new_xyz, new_points, idx, grouped_xyz
def pointnet_sa_module_msg(xyz, points, npoint, radius_list, nsample_list, mlp_list, is_training, bn_decay, scope, bn=True, use_xyz=True, use_nchw=False):
''' PointNet Set Abstraction (SA) module with Multi-Scale Grouping (MSG)
Input:
xyz: (batch_size, ndataset, 3) TF tensor
points: (batch_size, ndataset, channel) TF tensor
npoint: int32 -- #points sampled in farthest point sampling
radius: list of float32 -- search radius in local region
nsample: list of int32 -- how many points in each local region
mlp: list of list of int32 -- output size for MLP on each point
use_xyz: bool, if True concat XYZ with local point features, otherwise just use point features
use_nchw: bool, if True, use NCHW data format for conv2d, which is usually faster than NHWC format
Return:
new_xyz: (batch_size, npoint, 3) TF tensor
new_points: (batch_size, npoint, \sum_k{mlp[k][-1]}) TF tensor
'''
data_format = 'NCHW' if use_nchw else 'NHWC'
with tf.variable_scope(scope) as sc:
new_xyz = gather_point(xyz, farthest_point_sample(npoint, xyz))
new_points_list = []
for i in range(len(radius_list)):
radius = radius_list[i]
nsample = nsample_list[i]
idx, pts_cnt = query_ball_point(radius, nsample, xyz, new_xyz)
grouped_xyz = group_point(xyz, idx)
grouped_xyz -= tf.tile(tf.expand_dims(new_xyz, 2), [1,1,nsample,1])
if points is not None:
grouped_points = group_point(points, idx)
if use_xyz:
grouped_points = tf.concat([grouped_points, grouped_xyz], axis=-1)
else:
grouped_points = grouped_xyz
if use_nchw: grouped_points = tf.transpose(grouped_points, [0,3,1,2])
for j,num_out_channel in enumerate(mlp_list[i]):
grouped_points = tf_util.conv2d(grouped_points, num_out_channel, [1,1],
padding='VALID', stride=[1,1], bn=bn, is_training=is_training,
scope='conv%d_%d'%(i,j), bn_decay=bn_decay)
if use_nchw: grouped_points = tf.transpose(grouped_points, [0,2,3,1])
new_points = tf.reduce_max(grouped_points, axis=[2])
new_points_list.append(new_points)
new_points_concat = tf.concat(new_points_list, axis=-1)
return new_xyz, new_points_concat
def pointnet_sa_module_msg(xyz, points, npoint, radius_list, nsample_list, mlp_list, is_training, bn_decay, scope, bn=True, use_xyz=True, use_nchw=False):
''' PointNet Set Abstraction (SA) module with Multi-Scale Grouping (MSG)
Input:
xyz: (batch_size, ndataset, 3) TF tensor
points: (batch_size, ndataset, channel) TF tensor
npoint: int32 -- #points sampled in farthest point sampling
radius: list of float32 -- search radius in local region
nsample: list of int32 -- how many points in each local region
mlp: list of list of int32 -- output size for MLP on each point
use_xyz: bool, if True concat XYZ with local point features, otherwise just use point features
use_nchw: bool, if True, use NCHW data format for conv2d, which is usually faster than NHWC format
Return:
new_xyz: (batch_size, npoint, 3) TF tensor
new_points: (batch_size, npoint, \sum_k{mlp[k][-1]}) TF tensor
'''
data_format = 'NCHW' if use_nchw else 'NHWC'
with tf.variable_scope(scope) as sc:
new_xyz = gather_point(xyz, farthest_point_sample(npoint, xyz))
new_points_list = []
for i in range(len(radius_list)):
radius = radius_list[i]
nsample = nsample_list[i]
idx, pts_cnt = query_ball_point(radius, nsample, xyz, new_xyz)
grouped_xyz = group_point(xyz, idx)
grouped_xyz -= tf.tile(tf.expand_dims(new_xyz, 2), [1,1,nsample,1])
if points is not None:
grouped_points = group_point(points, idx)
if use_xyz:
grouped_points = tf.concat([grouped_points, grouped_xyz], axis=-1)
else:
grouped_points = grouped_xyz
if use_nchw: grouped_points = tf.transpose(grouped_points, [0,3,1,2])
for j,num_out_channel in enumerate(mlp_list[i]):
grouped_points = tf_util.conv2d(grouped_points, num_out_channel, [1,1],
padding='VALID', stride=[1,1], bn=bn, is_training=is_training,
scope='conv%d_%d'%(i,j), bn_decay=bn_decay)
if use_nchw: grouped_points = tf.transpose(grouped_points, [0,2,3,1])
new_points = tf.reduce_max(grouped_points, axis=[2])
new_points_list.append(new_points)
new_points_concat = tf.concat(new_points_list, axis=-1)
return new_xyz, new_points_concat
def sample_and_group(npoint, radius, nsample, xyz, points, global_idx, knn=False, use_xyz=True):
fps_idx, g_idx = farthest_point_sample(npoint, xyz, global_idx)
new_xyz = gather_point(xyz, fps_idx) # (batch_size, npoint, 3)
if knn:
_,idx = knn_point(nsample, xyz, new_xyz)
else:
idx, pts_cnt = query_ball_point(radius, nsample, xyz, new_xyz)
grouped_xyz = group_point(xyz, idx) # (batch_size, npoint, nsample, 3)
center_xyz = tf.tile(tf.expand_dims(new_xyz, 2), [1, 1, nsample, 1])
diff_xyz = grouped_xyz - center_xyz # translation normalization
euclid_dist = tf.norm(diff_xyz, axis=-1, keep_dims=True)
grouped_xyz = tf.concat([grouped_xyz, center_xyz, diff_xyz, euclid_dist], axis=-1)
if points is not None:
grouped_points = group_point(points, idx) # (batch_size, npoint, nsample, channel)
if use_xyz:
new_points = tf.concat([diff_xyz, grouped_points], axis=-1) # (batch_size, npoint, nample, 3+channel)
else:
new_points = grouped_points
else:
new_points = diff_xyz
return new_xyz, new_points, fps_idx, g_idx, diff_xyz
def sample_points(npoint, xyz):
'''
Input:
npoint: int32
radius: float32
nsample: int32
xyz: (batch_size, ndataset, 3) TF tensor
points: (batch_size, ndataset, channel) TF tensor, if None will just use xyz as points
tnet_spec: dict (keys: mlp, mlp2, is_training, bn_decay), if None do not apply tnet
knn: bool, if True use kNN instead of radius search
use_xyz: bool, if True concat XYZ with local point features, otherwise just use point features
Output:
new_xyz: (batch_size, npoint, 3) TF tensor
new_points: (batch_size, npoint, nsample, 3+channel) TF tensor
idx: (batch_size, npoint, nsample) TF tensor, indices of local points as in ndataset points
grouped_xyz: (batch_size, npoint, nsample, 3) TF tensor, normalized point XYZs
(subtracted by seed point XYZ) in local regions
'''
new_xyz = gather_point(xyz, farthest_point_sample(npoint, xyz)) # (batch_size, npoint, 3)
return new_xyz
def get_uniform_loss2(
pcd,
percentages=[0.002, 0.004, 0.006, 0.008, 0.010, 0.012, 0.015],
radius=1.0):
B, N, C = pcd.get_shape().as_list()
npoint = int(N * 0.05)
loss = []
for p in percentages:
nsample = int(N * p)
r = math.sqrt(p * radius)
#print(npoint,nsample)
new_xyz = gather_point(pcd, farthest_point_sample(
npoint, pcd)) # (batch_size, npoint, 3)
idx, pts_cnt = query_ball_point(
r, nsample, pcd, new_xyz) #(batch_size, npoint, nsample)
uniform_val = tf.py_func(py_uniform_loss, [pcd, idx, pts_cnt, r],
tf.float32)
loss.append(uniform_val * math.sqrt(p * 100))
return tf.add_n(loss) / len(percentages)
def sample_and_group(npoint, radius, nsample, xyz, points, tnet_spec=None, knn=False, use_xyz=True):
'''
Input:
npoint: int32
radius: float32
nsample: int32
xyz: (batch_size, ndataset, 3) TF tensor
points: (batch_size, ndataset, channel) TF tensor, if None will just use xyz as points
tnet_spec: dict (keys: mlp, mlp2, is_training, bn_decay), if None do not apply tnet
knn: bool, if True use kNN instead of radius search
use_xyz: bool, if True concat XYZ with local point features, otherwise just use point features
Output:
new_xyz: (batch_size, npoint, 3) TF tensor
new_points: (batch_size, npoint, nsample, 3+channel) TF tensor
idx: (batch_size, npoint, nsample) TF tensor, indices of local points as in ndataset points
grouped_xyz: (batch_size, npoint, nsample, 3) TF tensor, normalized point XYZs
(subtracted by seed point XYZ) in local regions
'''
#npoint个采样点，每个采样点周围采样nsample个点 #sampling layer
new_xyz = gather_point(xyz, farthest_point_sample(npoint, xyz)) # (batch_size, npoint, 3)
# grouping layer
if knn:
_,idx = knn_point(nsample, xyz, new_xyz)
else: #idx应该是用数字表示是第几个shape的第几个点周围的neightbor，也就是一个sample内部的序号
idx, pts_cnt = query_ball_point(radius, nsample, xyz, new_xyz)
grouped_xyz = group_point(xyz, idx) # (batch_size, npoint, nsample, 3) #把sample的内部序号和点的序号对应的坐标组合起来
#每个点和中心点的相对距离
grouped_xyz -= tf.tile(tf.expand_dims(new_xyz, 2), [1,1,nsample,1]) # translation normalization 对每个npoint插入一维，大小是nsample，值和中间点一样。
if tnet_spec is not None:
grouped_xyz = tnet(grouped_xyz, tnet_spec)
if points is not None:
grouped_points = group_point(points, idx) # (batch_size, npoint, nsample, channel)#points是一些featrue，xyz才是坐标
if use_xyz:
new_points = tf.concat([grouped_xyz, grouped_points], axis=-1) # (batch_size, npoint, nample, 3+channel)
else:
new_points = grouped_points
else:
new_points = grouped_xyz
return new_xyz, new_points, idx, grouped_xyz #中心点，分组后的点（可能包含channel），idx和分组后的xyz
def sample_and_group(npoint, radius, nsample, xyz, points, knn=False, use_xyz=True, return_new=False):
'''
Input:
npoint: int32
radius: float32
nsample: int32
xyz: (batch_size, ndataset, 3) TF tensor
points: (batch_size, ndataset, channel) TF tensor, if None will just use xyz as points
knn: bool, if True use kNN instead of radius search
use_xyz: bool, if True concat XYZ with local point features, otherwise just use point features
Output:
new_xyz: (batch_size, npoint, 3) TF tensor
new_points: (batch_size, npoint, nsample, 3+channel) TF tensor
idx: (batch_size, npoint, nsample) TF tensor, indices of local points as in ndataset points
grouped_xyz: (batch_size, npoint, nsample, 3) TF tensor, normalized point XYZs
(subtracted by seed point XYZ) in local regions
'''
new_idx = farthest_point_sample(npoint, xyz)
new_xyz = gather_point(xyz, new_idx) # (batch_size, npoint, 3)
if knn:
_,idx = knn_point(nsample, xyz, new_xyz)
else:
idx, pts_cnt = query_ball_point(radius, nsample, xyz, new_xyz)
grouped_xyz = group_point(xyz, idx) # (batch_size, npoint, nsample, 3)
grouped_xyz -= tf.tile(tf.expand_dims(new_xyz, 2), [1,1,nsample,1]) # translation normalization
if points is not None:
grouped_points = group_point(points, idx) # (batch_size, npoint, nsample, channel)
if use_xyz:
new_points = tf.concat([grouped_xyz, grouped_points], axis=-1) # (batch_size, npoint, nample, 3+channel)
else:
new_points = grouped_points
else:
new_points = grouped_xyz
if return_new:
return new_xyz, new_points, idx, grouped_xyz, new_idx
else:
return new_xyz, new_points, idx, grouped_xyz
def test(self):
self.opts.batch_size = 1
final_ratio = self.opts.final_ratio
step_ratio = 4
self.opts.up_ratio = step_ratio
self.build_model_test(final_ratio=self.opts.final_ratio, step_ratio=step_ratio)
saver = tf.train.Saver()
restore_epoch, checkpoint_path = model_utils.pre_load_checkpoint(self.opts.log_dir)
print(checkpoint_path)
self.saver.restore(self.sess, checkpoint_path)
#self.restore_model(self.opts.log_dir, epoch=self.opts.restore_epoch, verbose=True)
samples = glob(self.opts.test_data)
point = pc_util.load(samples[0])
self.opts.num_point = point.shape[0]
out_point_num = int(self.opts.num_point * final_ratio)
for point_path in samples:
logging.info(point_path)
start = time()
pc = pc_util.load(point_path)[:, :3]
pc, centroid, furthest_distance = pc_util.normalize_point_cloud(pc)
input_list, pred_list, coarse_list = self.pc_prediction(pc)
end = time()
print("total time: ", end - start)
pred_pc = np.concatenate(pred_list, axis=0)
pred_pc = (pred_pc * furthest_distance) + centroid
pred_pc = np.reshape(pred_pc, [-1, 3])
idx = farthest_point_sample(out_point_num, pred_pc[np.newaxis, ...]).eval(session=self.sess)[0]
pred_pc = pred_pc[idx, 0:3]
# path = os.path.join(self.opts.out_folder, point_path.split('/')[-1][:-4] + '.ply')
# np.savetxt(path[:-4] + '.xyz',pred_pc,fmt='%.6f')
in_folder = os.path.dirname(self.opts.test_data)
path = os.path.join(self.opts.out_folder, point_path.split('/')[-1][:-4] + '_X%d.xyz' % final_ratio)
np.savetxt(path, pred_pc, fmt='%.6f')
def pointcnn_xconv_module(xyz, points, npoint, c_fts_out, c_x, k_neighbors,
d_rate, is_training, depth_multiplier, with_global,
scope):
'''
input:
xyz: TF tensor, input point clouds coords, B x N x 3
points: TF tensor, input point clouds features, B x N x fts_channel
npoint: int32, number of representative samplings
c_fts_out: int32, output channels number
c_x: int32, channels number of lifted features for x-transformation matrix
k_neighbors: int32, neighbor size
d_rate: int32, dilation rate
is_training: TF tensor, flag indicating training status
output:
new_xyz: TF tensor, output point clouds coords, B x npoint x 3
new_points: TF tensor, output point clouds features, B x npoint x c_fts_out
'''
with tf.variable_scope(scope):
new_xyz = gather_point(xyz, farthest_point_sample(npoint, xyz))
new_points = xconv(xyz,
points,
new_xyz,
scope,
xyz.shape[0],
k_neighbors,
d_rate,
npoint,
c_fts_out,
c_x,
is_training,
True,
depth_multiplier,
sorting_method=None,
with_global=with_global)
return new_xyz, new_points
def get_Geometric_Loss(predictedPts, targetpoints, FLAGS):
gen_points = FLAGS.generate_num
targetpoints = gather_point(targetpoints,
farthest_point_sample(
gen_points,
targetpoints)) #将targetpoints按照输出的点的个数采样
# calculate shape loss
square_dist = pairwise_l2_norm2_batch(targetpoints, predictedPts)
dist = tf.sqrt(square_dist) # 开方
minRow = tf.reduce_min(dist, axis=2) ## 在降维后的第二维，即y那维找最小
minCol = tf.reduce_min(dist, axis=1) ## 在x那一维找最小值
shapeLoss = tf.reduce_mean(minRow) + tf.reduce_mean(
minCol) ## 在[batchsize,x]取平均
# calculate density loss
square_dist2 = pairwise_l2_norm2_batch(targetpoints, targetpoints)
dist2 = tf.sqrt(square_dist2)
knndis = tf.nn.top_k(tf.negative(dist), k=FLAGS.nnk) #返回每行最大的8个数
knndis2 = tf.nn.top_k(tf.negative(dist2), k=FLAGS.nnk)
densityLoss = tf.reduce_mean(tf.abs(knndis.values - knndis2.values))
data_loss = shapeLoss + densityLoss * FLAGS.densityWeight
return data_loss, shapeLoss, densityLoss
def point2sequence_module(xyz,
points,
npoint,
nsample_list,
mlp_list,
hidden_size,
output_size,
is_training,
bn_decay,
scope,
bn=True,
use_xyz=True,
use_nchw=False,
batch_size=16):
''' Point2sequence module
assume mlp[k][-1] are all the same for rnn input
Input:
xyz: (batch_size, ndataset, 3) TF tensor
points: (batch_size, ndataset, channel) TF tensor
npoint: int32 -- #points sampled in farthest point sampling
nsample: list of int32 -- how many points in each local region
mlp: list of list of int32 -- output size for MLP on each point
hidden_size: int32 -- hidden size of the RNN hidden state
use_xyz: bool, if True concat XYZ with local point features, otherwise just use point features
use_nchw: bool, if True, use NCHW data format for conv2d, which is usually faster than NHWC format
Return:
new_xyz: (batch_size, npoint, 3) TF tensor
feature: (batch_size, npoint, output_size}) TF tensor
'''
with tf.variable_scope(scope) as sc:
new_xyz = gather_point(xyz, farthest_point_sample(npoint, xyz))
new_points_list = []
for i in range(len(nsample_list)):
nsample = nsample_list[i]
vla, idx = knn_point(nsample, xyz, new_xyz)
grouped_xyz = group_point(xyz, idx)
# convert to the relative coordinate system
grouped_xyz -= tf.tile(tf.expand_dims(new_xyz, 2),
[1, 1, nsample, 1])
if points is not None:
grouped_points = group_point(points, idx)
if use_xyz:
grouped_points = tf.concat([grouped_points, grouped_xyz],
axis=-1)
else:
grouped_points = grouped_xyz
if use_nchw:
grouped_points = tf.transpose(grouped_points, [0, 3, 1, 2])
# MLP layers
for j, num_out_channel in enumerate(mlp_list[i]):
grouped_points = tf_util.conv2d(grouped_points,
num_out_channel, [1, 1],
padding='VALID',
stride=[1, 1],
bn=bn,
is_training=is_training,
scope='conv%d_%d' % (i, j),
bn_decay=bn_decay)
if use_nchw:
grouped_points = tf.transpose(grouped_points, [0, 2, 3, 1])
# max pool
new_points = tf.reduce_max(grouped_points, axis=[2])
new_points_list.append(new_points)
# multi-scale area feature sequence
feature = tf.reshape(tf.concat(new_points_list, axis=-1),
(-1, npoint, len(nsample_list), mlp_list[-1][-1]))
# RNN-based sequence model with attention
feature = tf_util.seq2seq_with_attention(feature,
hidden_size,
scope='seq_attention',
bn=bn,
is_training=is_training,
bn_decay=bn_decay)
if output_size > 0:
feature = tf.expand_dims(feature, 1)
feature = tf_util.conv2d(feature,
output_size, [1, 1],
padding='VALID',
stride=[1, 1],
bn=bn,
is_training=is_training,
scope='conv',
bn_decay=bn_decay)
feature = tf.reshape(feature, (-1, npoint, output_size))
return new_xyz, feature
def fps(self, num_gt_points, completion): # farthest point sampling
fps_indices = farthest_point_sample(num_gt_points, completion)
fps_completion = gather_point(completion,
fps_indices) # same point number with gt
return fps_completion, fps_indices
def Invariance_Transformation_Net(point_cloud, features, is_training,
invarians_trans_param):
xconv_params = invarians_trans_param.xconv_params
fc_params = invarians_trans_param.fc_params
with_X_transformation = invarians_trans_param.with_X_transformation
sorting_method = invarians_trans_param.sorting_method
#N=point_cloud.get_shape()[1].value
N = point_cloud.get_shape()[0].value
if invarians_trans_param.sampling == 'fps':
import tf_sampling
layer_pts = [point_cloud]
if features is None:
layer_fts = [features]
else:
features = tf.reshape(features,
(N, -1, invarians_trans_param.data_dim - 3),
name='features_reshape')
C_fts = xconv_params[0]['C'] // 2
features_hd = pf.dense(features, C_fts, 'features_hd', is_training)
layer_fts = [features_hd]
for layer_idx, layer_param in enumerate(xconv_params):
tag = 'xconv_' + str(
layer_idx +
1) + '_' #####xconv_1_ #####xconv_2_ #####xconv_3_ #####xconv_4_
K = layer_param['K']
D = layer_param['D']
P = layer_param['P']
C = layer_param['C']
links = layer_param['links'] ## type(layer_param) is dict
#get k-nearest points
pts = layer_pts[-1]
fts = layer_fts[-1]
if P == -1 or (layer_idx > 0
and P == xconv_params[layer_idx - 1]['P']):
qrs = layer_pts[-1]
else:
if invarians_trans_param.sampling == 'fps':
fps_indices = tf_sampling.farthest_point_sample(P, pts)
batch_indices = tf.tile(tf.reshape(tf.range(N), (-1, 1, 1)),
(1, P, 1))
indices = tf.concat(
[batch_indices,
tf.expand_dims(fps_indices, -1)], axis=-1)
qrs = tf.gather_nd(pts, indices, name=tag + 'qrs') ### (N,P,3)
print(tf.shape(qrs))
elif invarians_trans_param.sampling == 'ids':
indices = pf.inverse_density_sampling(pts, K, P)
qrs = tf.gather_nd(pts, indices)
elif invarians_trans_param == 'random':
qrs = tf.slice(pts, (0, 0, 0), (-1, P, -1),
name=tag + 'qrs') ### (N,P,3)
else:
print('unknown sampling method')
exit()
layer_pts.append(qrs)
if layer_idx == 0:
C_pts_fts = C // 2 if fts is None else C // 4
depth_multipiler = 4
else:
C_prev = xconv_params[layer_idx - 1]['C']
C_pts_fts = C_prev // 4
depth_multipiler = math.ceil(C / C_prev)
with_global = (invarians_trans_param.with_global
and layer_idx == len(xconv_params) - 1)
fts_xconv = xconv(pts, fts, qrs, tag, N, K, D, P, C, C_pts_fts,
is_training, with_X_transformation, depth_multipiler,
sorting_method, with_global)
fts_list = []
for link in links:
fts_from_link = layer_fts[link]
if fts_from_link is not None:
fts_slice = tf.slice(fts_from_link, (0, 0, 0), (-1, P, -1),
name=tag + 'fts_slice' + str(-link))
fts_list.append(fts_slice)
if fts_list:
fts_list.append(fts_xconv)
layer_fts.append(
tf.concat(fts_list, axis=-1, name=tag + 'fts_list_concat'))
else:
layer_fts.append(fts_xconv)
if hasattr(invarians_trans_param, 'xdconv_params'):
for layer_idx, layer_param in enumerate(
invarians_trans_param.xdconv_params):
tag = 'xdconv_' + str(layer_idx + 1) + '_'
K = layer_param['K']
D = layer_param['D']
pts_layer_idx = layer_param['pts_layer_idx']
qrs_layer_idx = layer_param['qrs_layer_idx']
pts = layer_pts[pts_layer_idx + 1]
fts = layer_fts[pts_layer_idx +
1] if layer_idx == 0 else layer_fts[-1]
qrs = layer_pts[qrs_layer_idx + 1]
fts_qrs = layer_fts[qrs_layer_idx + 1]
P = xconv_params[qrs_layer_idx]['P']
C = xconv_params[qrs_layer_idx]['C']
C_prev = xconv_params[pts_layer_idx]['C']
C_pts_fts = C_prev // 4
depth_multipiler = 1
fts_xdconv = xconv(pts, fts, qrs, tag, N, K, D, P, C, C_pts_fts,
is_training, with_X_transformation,
depth_multipiler, sorting_method)
fts_concat = tf.concat([fts_xdconv, fts_qrs],
axis=-1,
name=tag + 'fts_concat')
fts_fuse = pf.dense(fts_concat, C, tag + 'fts_fuse', is_training)
layer_pts.append(qrs)
layer_fts.append(fts_fuse)
fc_layers = layer_fts[-1]
for layer_idx, layer_param in enumerate(fc_params):
C = layer_param['C']
dropout_rate = layer_param['dropout_rate']
fc = pf.dense(fc_layers[-1], C, 'fc{:d}'.format(layer_idx),
is_training)
fc_drop = tf.layers.dropout(fc,
dropout_rate,
training=is_training,
name='fc{:d}_drop'.format(layer_idx))
fc_layers.append(fc_drop)
return fc_layers
def get_model(point_cloud, is_training, global_config, bn_decay=None):
"""
Contact-GraspNet model consisting of a PointNet++ backbone and multiple output heads
Arguments:
point_cloud {tf.placeholder} -- batch of point clouds
is_training {bool} -- train or eval mode
global_config {dict} -- config
Keyword Arguments:
bn_decay {tf.variable} -- batch norm decay (default: {None})
Returns:
[dict] -- endpoints of the network
"""
model_config = global_config['MODEL']
data_config = global_config['DATA']
radius_list_0 = model_config['pointnet_sa_modules_msg'][0]['radius_list']
radius_list_1 = model_config['pointnet_sa_modules_msg'][1]['radius_list']
radius_list_2 = model_config['pointnet_sa_modules_msg'][2]['radius_list']
nsample_list_0 = model_config['pointnet_sa_modules_msg'][0]['nsample_list']
nsample_list_1 = model_config['pointnet_sa_modules_msg'][1]['nsample_list']
nsample_list_2 = model_config['pointnet_sa_modules_msg'][2]['nsample_list']
mlp_list_0 = model_config['pointnet_sa_modules_msg'][0]['mlp_list']
mlp_list_1 = model_config['pointnet_sa_modules_msg'][1]['mlp_list']
mlp_list_2 = model_config['pointnet_sa_modules_msg'][2]['mlp_list']
npoint_0 = model_config['pointnet_sa_modules_msg'][0]['npoint']
npoint_1 = model_config['pointnet_sa_modules_msg'][1]['npoint']
npoint_2 = model_config['pointnet_sa_modules_msg'][2]['npoint']
fp_mlp_0 = model_config['pointnet_fp_modules'][0]['mlp']
fp_mlp_1 = model_config['pointnet_fp_modules'][1]['mlp']
fp_mlp_2 = model_config['pointnet_fp_modules'][2]['mlp']
input_normals = data_config['input_normals']
offset_bins = data_config['labels']['offset_bins']
joint_heads = model_config['joint_heads']
# expensive, rather use random only
if 'raw_num_points' in data_config and data_config[
'raw_num_points'] != data_config['ndataset_points']:
point_cloud = gather_point(
point_cloud,
farthest_point_sample(data_config['ndataset_points'], point_cloud))
end_points = {}
l0_xyz = tf.slice(point_cloud, [0, 0, 0], [-1, -1, 3])
l0_points = tf.slice(point_cloud, [0, 0, 3],
[-1, -1, 3]) if input_normals else None
# Set abstraction layers
l1_xyz, l1_points = pointnet_sa_module_msg(l0_xyz,
l0_points,
npoint_0,
radius_list_0,
nsample_list_0,
mlp_list_0,
is_training,
bn_decay,
scope='layer1')
l2_xyz, l2_points = pointnet_sa_module_msg(l1_xyz,
l1_points,
npoint_1,
radius_list_1,
nsample_list_1,
mlp_list_1,
is_training,
bn_decay,
scope='layer2')
if 'asymmetric_model' in model_config and model_config['asymmetric_model']:
l3_xyz, l3_points = pointnet_sa_module_msg(l2_xyz,
l2_points,
npoint_2,
radius_list_2,
nsample_list_2,
mlp_list_2,
is_training,
bn_decay,
scope='layer3')
l4_xyz, l4_points, _ = pointnet_sa_module(
l3_xyz,
l3_points,
npoint=None,
radius=None,
nsample=None,
mlp=model_config['pointnet_sa_module']['mlp'],
mlp2=None,
group_all=model_config['pointnet_sa_module']['group_all'],
is_training=is_training,
bn_decay=bn_decay,
scope='layer4')
# Feature Propagation layers
l3_points = pointnet_fp_module(l3_xyz,
l4_xyz,
l3_points,
l4_points,
fp_mlp_0,
is_training,
bn_decay,
scope='fa_layer1')
l2_points = pointnet_fp_module(l2_xyz,
l3_xyz,
l2_points,
l3_points,
fp_mlp_1,
is_training,
bn_decay,
scope='fa_layer2')
l1_points = pointnet_fp_module(l1_xyz,
l2_xyz,
l1_points,
l2_points,
fp_mlp_2,
is_training,
bn_decay,
scope='fa_layer3')
l0_points = l1_points
pred_points = l1_xyz
else:
l3_xyz, l3_points, _ = pointnet_sa_module(
l2_xyz,
l2_points,
npoint=None,
radius=None,
nsample=None,
mlp=model_config['pointnet_sa_module']['mlp'],
mlp2=None,
group_all=model_config['pointnet_sa_module']['group_all'],
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,
fp_mlp_0,
is_training,
bn_decay,
scope='fa_layer1')
l1_points = pointnet_fp_module(l1_xyz,
l2_xyz,
l1_points,
l2_points,
fp_mlp_1,
is_training,
bn_decay,
scope='fa_layer2')
l0_points = tf.concat([l0_xyz, l0_points],
axis=-1) if input_normals else l0_xyz
l0_points = pointnet_fp_module(l0_xyz,
l1_xyz,
l0_points,
l1_points,
fp_mlp_2,
is_training,
bn_decay,
scope='fa_layer3')
pred_points = l0_xyz
if joint_heads:
head = tf_util.conv1d(l0_points,
128,
1,
padding='VALID',
bn=True,
is_training=is_training,
scope='fc1',
bn_decay=bn_decay)
head = tf_util.dropout(head,
keep_prob=0.7,
is_training=is_training,
scope='dp1')
head = tf_util.conv1d(head,
4,
1,
padding='VALID',
activation_fn=None,
scope='fc2')
grasp_dir_head = tf.slice(head, [0, 0, 0], [-1, -1, 3])
grasp_dir_head = tf.math.l2_normalize(grasp_dir_head, axis=2)
binary_seg_head = tf.slice(head, [0, 0, 3], [-1, -1, 1])
else:
# Head for grasp direction
grasp_dir_head = tf_util.conv1d(l0_points,
128,
1,
padding='VALID',
bn=True,
is_training=is_training,
scope='fc1',
bn_decay=bn_decay)
grasp_dir_head = tf_util.dropout(grasp_dir_head,
keep_prob=0.7,
is_training=is_training,
scope='dp1')
grasp_dir_head = tf_util.conv1d(grasp_dir_head,
3,
1,
padding='VALID',
activation_fn=None,
scope='fc3')
grasp_dir_head_normed = tf.math.l2_normalize(grasp_dir_head, axis=2)
# Head for grasp approach
approach_dir_head = tf_util.conv1d(l0_points,
128,
1,
padding='VALID',
bn=True,
is_training=is_training,
scope='fc1_app',
bn_decay=bn_decay)
approach_dir_head = tf_util.dropout(approach_dir_head,
keep_prob=0.7,
is_training=is_training,
scope='dp1_app')
approach_dir_head = tf_util.conv1d(approach_dir_head,
3,
1,
padding='VALID',
activation_fn=None,
scope='fc3_app')
approach_dir_head_orthog = tf.math.l2_normalize(
approach_dir_head - tf.reduce_sum(
tf.multiply(grasp_dir_head_normed, approach_dir_head),
axis=2,
keepdims=True) * grasp_dir_head_normed,
axis=2)
# Head for grasp width
if model_config['dir_vec_length_offset']:
grasp_offset_head = tf.norm(grasp_dir_head, axis=2, keepdims=True)
elif model_config['bin_offsets']:
grasp_offset_head = tf_util.conv1d(l0_points,
128,
1,
padding='VALID',
bn=True,
is_training=is_training,
scope='fc1_off',
bn_decay=bn_decay)
grasp_offset_head = tf_util.conv1d(grasp_offset_head,
len(offset_bins) - 1,
1,
padding='VALID',
activation_fn=None,
scope='fc2_off')
else:
grasp_offset_head = tf_util.conv1d(l0_points,
128,
1,
padding='VALID',
bn=True,
is_training=is_training,
scope='fc1_off',
bn_decay=bn_decay)
grasp_offset_head = tf_util.dropout(grasp_offset_head,
keep_prob=0.7,
is_training=is_training,
scope='dp1_off')
grasp_offset_head = tf_util.conv1d(grasp_offset_head,
1,
1,
padding='VALID',
activation_fn=None,
scope='fc2_off')
# Head for contact points
binary_seg_head = tf_util.conv1d(l0_points,
128,
1,
padding='VALID',
bn=True,
is_training=is_training,
scope='fc1_seg',
bn_decay=bn_decay)
binary_seg_head = tf_util.dropout(binary_seg_head,
keep_prob=0.5,
is_training=is_training,
scope='dp1_seg')
binary_seg_head = tf_util.conv1d(binary_seg_head,
1,
1,
padding='VALID',
activation_fn=None,
scope='fc2_seg')
end_points['grasp_dir_head'] = grasp_dir_head_normed
end_points['binary_seg_head'] = binary_seg_head
end_points['binary_seg_pred'] = tf.math.sigmoid(binary_seg_head)
end_points['grasp_offset_head'] = grasp_offset_head
end_points['grasp_offset_pred'] = tf.math.sigmoid(
grasp_offset_head
) if model_config['bin_offsets'] else grasp_offset_head
end_points['approach_dir_head'] = approach_dir_head_orthog
end_points['pred_points'] = pred_points
return end_points
def get_model(layer_pts, is_training, RIconv_params, RIdconv_params, fc_params, sampling='fps', weight_decay=0.0, bn_decay=None, part_num=50):
if sampling == 'fps':
sys.path.append(os.path.join(BASE_DIR, 'tf_ops/sampling'))
from tf_sampling import farthest_point_sample, gather_point
layer_fts_list = [None]
layer_pts_list = [layer_pts]
for layer_idx, layer_param in enumerate(RIconv_params):
tag = 'xconv_' + str(layer_idx + 1) + '_'
K = layer_param['K']
D = layer_param['D']
P = layer_param['P']
C = layer_param['C']
# qrs = layer_pts if P == -1 else layer_pts[:,:P,:] # (N, P, 3)
if P == -1:
qrs = layer_pts
else:
if sampling == 'fps':
qrs = gather_point(layer_pts, farthest_point_sample(P, layer_pts))
elif sampling == 'random':
qrs = tf.slice(layer_pts, (0, 0, 0), (-1, P, -1), name=tag + 'qrs') # (N, P, 3)
else:
print('Unknown sampling method!')
exit()
layer_fts= RIConv(layer_pts_list[-1], layer_fts_list[-1], qrs, is_training, tag, K, D, P, C, True, bn_decay)
layer_pts = qrs
layer_pts_list.append(qrs)
layer_fts_list.append(layer_fts)
if RIdconv_params is not None:
fts = layer_fts_list[-1]
for layer_idx, layer_param in enumerate(RIdconv_params):
tag = 'xdconv_' + str(layer_idx + 1) + '_'
K = layer_param['K']
D = layer_param['D']
pts_layer_idx = layer_param['pts_layer_idx'] # 2 1 0
qrs_layer_idx = layer_param['qrs_layer_idx'] # 1 0 -1
pts = layer_pts_list[pts_layer_idx + 1]
qrs = layer_pts_list[qrs_layer_idx + 1]
fts_qrs = layer_fts_list[qrs_layer_idx + 1]
C = fts_qrs.get_shape()[-1].value if fts_qrs is not None else C//2
P = qrs.get_shape()[1].value
layer_fts= RIConv(pts, fts, qrs, is_training, tag, K, D, P, C, True, bn_decay)
if fts_qrs is not None: # this is for last layer
fts_concat = tf.concat([layer_fts, fts_qrs], axis=-1, name=tag + 'fts_concat')
fts = pf.dense(fts_concat, C, tag + 'fts_fuse', is_training)
else:
fts = layer_fts
for layer_idx, layer_param in enumerate(fc_params):
C = layer_param['C']
dropout_rate = layer_param['dropout_rate']
layer_fts = pf.dense(layer_fts, C, 'fc{:d}'.format(layer_idx), is_training)
layer_fts = tf.layers.dropout(layer_fts, dropout_rate, training=is_training, name='fc{:d}_drop'.format(layer_idx))
logits_seg = pf.dense(layer_fts, part_num, 'logits', is_training, with_bn=False, activation=None)
return logits_seg
def cache_file(self, src_data_path, cache_path):
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
with tf.device('/gpu:0'):
pc_tf = tf.placeholder(tf.float32)
ind_tf = tf_sampling.farthest_point_sample(self.npoint, pc_tf)
ind_ins_tf = tf_sampling.farthest_point_sample(
self.npoint_ins, pc_tf)
sess = tf.Session(config=config)
chunk_file = []
if isinstance(src_data_path, list):
for src in src_data_path:
chunk_file.extend(
glob.glob(os.path.join(src, self.data_type + '*.h5')))
else:
chunk_file.extend(
glob.glob(os.path.join(src_data_path,
self.data_type + '*.h5')))
# chunk_file = glob.glob(os.path.join(src_data_path, self.data_type+'*-00.h5'))
print chunk_file
chunk_file.sort()
data_index = 0
for i, cur_chunk in enumerate(chunk_file):
hf = h5py.File(cur_chunk, 'r')
pc_chunk = hf['pts'].value #[B, N, 3]
seg_chunk = hf[
'gt_label'].value # [B, N], (0~curnseg), 0 is background points
group_chunk = np.concatenate((np.expand_dims(
hf['gt_other_mask'].value, 1), hf['gt_mask'].value),
1) # [B, 201, N], binary mask
group_chunk = np.argmax(
group_chunk, 1) # [B, N], (0~curngroup), 0 is background
nfile = pc_chunk.shape[0]
hf.close()
#### update ngroup and seg
if np.max(group_chunk) + 1 > self.ngroup:
self.ngroup = np.max(group_chunk) + 1
if np.max(seg_chunk) + 1 > self.nseg:
self.nseg = np.max(seg_chunk) + 1
for index in range(nfile):
print(i, np.float32(index) / nfile)
curpc = pc_chunk[index, :, :].astype(np.float32) # [npoint, 3]
curgroup = group_chunk[index, :].astype(
np.int32) # group zero is background, [npoint]
curseg = seg_chunk[index, :].astype(
np.int32) # 0 is background, [npoint]
curngroup = np.max(
curgroup) + 1 # group zero is background, [ngroup+1]
#### sample instance for each group
pc_ins = np.zeros((curngroup, self.npoint_ins, 3),
dtype=np.float32)
for j in range(1, curngroup):
if np.sum(curgroup == j) < 5:
continue
curins = curpc[curgroup == j, :]
if self.npoint_ins < curins.shape[0]:
choice = sess.run(
ind_ins_tf,
feed_dict={pc_tf: np.expand_dims(curins, 0)})[0]
pc_ins[j, :, :] = curins[choice, :]
elif self.npoint_ins == curins.shape[0]:
pc_ins[j, :, :] = copy.deepcopy(curins)
else:
choice = np.random.choice(
curins.shape[0], self.npoint_ins - curins.shape[0])
pc_ins[j, :, :] = np.concatenate(
(curins, curins[choice, :]), 0)
#### remove group less than 5 points
valid_group_indicator = np.ones(curngroup - 1)
target_group_idx = np.zeros(curngroup)
count = 0
for j in range(1, curngroup):
if np.sum(curgroup == j) < 5:
valid_group_indicator[j - 1] = 0
else:
count += 1
target_group_idx[j] = count
curgroup = self.changem(curgroup,
np.arange(curngroup).astype(np.int32),
target_group_idx).astype('int32')
valid_group_indicator = np.concatenate(
([1], valid_group_indicator))
pc_ins = pc_ins[valid_group_indicator == 1, :, :]
curngroup = count + 1 # group zero is background, [ngroup+1]
pc = curpc
group_label = curgroup
seg_label = curseg
seg_label_per_group = np.zeros(curngroup).astype(np.int32)
for j in range(1, curngroup):
seg_label_per_group[j] = stats.mode(
seg_label[group_label == j])[0][0]
self.data_list[data_index] = (pc, group_label, seg_label,
seg_label_per_group, pc_ins,
curngroup)
data_index += 1
np.savez_compressed(cache_path,
data_list=self.data_list,
ngroup=self.ngroup,
nseg=self.nseg)
def edge_preserve_sampling(feature_input,
point_input,
num_samples,
adj_input,
dist_threshold,
knn,
atrous,
radius_min,
radius_max,
PFS_flag=False,
atrous_flag=True):
'''
Input:
feature_input: (batch_size, num_points, num_features)
point_input: (batch_size, num_points, 3)
num_samples: int32
adj_input: (batch_size, num_points, num_points)
dist_threshold: bool
knn: int32
atrous: int32
radius_min: float32
radius_max: float32
PFS_flag: bool
Returns:
net: (batch_size, num_samples, 1, 2 * num_features)
p_idx: (batch_size, num_samples)
pn_idx: (batch_size, num_samples, knn)
point_output: (batch_size, num_samples, 3)
'''
feature_shape = feature_input.get_shape()
batch_size = feature_shape[0]
num_points = int(feature_shape[1])
if PFS_flag == False:
p_idx = farthest_point_sample(num_samples, point_input)
else:
p_idx = tf_util.gather_principal_feature(feature_input, num_samples)
point_output = gather_point(point_input, p_idx)
padj_matrix = tf.squeeze(
group_point(adj_input, tf.expand_dims(p_idx, axis=-1)))
if batch_size == 1:
padj_matrix = tf.expand_dims(padj_matrix, axis=0)
if num_points == 1:
padj_matrix = tf.expand_dims(padj_matrix, axis=1)
if dist_threshold:
pdist_matrix = padj_matrix
else:
pdist_matrix = None
if atrous_flag:
pk = int(min(knn, num_points / atrous))
pn_idx = tf_util.get_atrous_knn(padj_matrix, pk, atrous, pdist_matrix,
radius_min, radius_max)
else:
pk = int(min(knn, num_points))
_, pn_idx = knn_point(pk, point_input, point_output)
neighbor_feature = group_point(feature_input, pn_idx)
neighbor_feature = tf.reduce_max(neighbor_feature, axis=-2, keepdims=True)
center_feature = group_point(feature_input, tf.expand_dims(p_idx, 2))
net = tf.concat([center_feature, neighbor_feature], axis=-1)
return net, p_idx, pn_idx, point_output
def get_model(point_cloud, is_training, bn_decay=None):
""" ConvNet baseline, input is BxNx9 gray image """
batch_size = point_cloud.get_shape()[0].value
num_point = point_cloud.get_shape()[1].value
end_points = {}
point_cloud_0 = tf.slice(point_cloud, [0, 0, 0], [-1, -1, 3])
l0_net = tf.slice(point_cloud, [0, 0, 3], [-1, -1, 6])
# l0_net = tf.expand_dims(l0_net, 2)
k = 32
npoints = 1024
new_point_cloud_1 = gather_point(
point_cloud_0, farthest_point_sample(npoints, point_cloud_0))
net = get_sampled_edgeconv_groupconv(point_cloud_0,
new_point_cloud_1,
k, [32, 32, 64],
is_training=is_training,
bn_decay=bn_decay,
scope='layer1',
bn=True,
is_dist=True)
l1_net = tf.squeeze(net, [2])
k = 20
npoints = 256
sampled_net, new_point_cloud_2 = get_sampled_feature(
new_point_cloud_1, net, npoints)
net = get_sampled_edgeconv_groupconv(net,
sampled_net,
k, [64, 64, 128],
is_training=is_training,
bn_decay=bn_decay,
scope='layer2',
bn=True,
sampled_pc=new_point_cloud_2,
pc=new_point_cloud_1,
is_dist=True)
net = get_edgeconv(
net,
k, [128],
is_training=is_training,
bn_decay=bn_decay,
scope='layer3',
bn=True,
associated=[sampled_net,
tf.expand_dims(new_point_cloud_2, axis=-2)],
is_dist=True)
l2_net = tf.squeeze(net, [2])
npoints = 64
sampled_net, new_point_cloud_3 = get_sampled_feature(
new_point_cloud_2, net, npoints)
net = get_sampled_edgeconv_groupconv(net,
sampled_net,
k, [128, 128, 256],
is_training=is_training,
bn_decay=bn_decay,
scope='layer4',
bn=True,
sampled_pc=new_point_cloud_3,
pc=new_point_cloud_2,
is_dist=True)
l3_net = tf.squeeze(net, [2])
k = 16
npoints = 16
sampled_net, new_point_cloud_4 = get_sampled_feature(
new_point_cloud_3, net, npoints)
net = get_sampled_edgeconv_groupconv(net,
sampled_net,
k, [256, 256, 512],
is_training=is_training,
bn_decay=bn_decay,
scope='layer5',
bn=True,
sampled_pc=new_point_cloud_4,
pc=new_point_cloud_3,
is_dist=True)
net = get_edgeconv(
net,
k, [512],
is_training=is_training,
bn_decay=bn_decay,
scope='layer6',
bn=True,
associated=[sampled_net,
tf.expand_dims(new_point_cloud_4, axis=-2)],
is_dist=True)
l4_net = tf.squeeze(net, [2])
l3_net = pointnet_fp_module(new_point_cloud_3,
new_point_cloud_4,
l3_net,
l4_net, [256, 256],
is_training,
bn_decay,
scope='fa_layer1',
is_dist=True)
l2_net = pointnet_fp_module(new_point_cloud_2,
new_point_cloud_3,
l2_net,
l3_net, [256, 256],
is_training,
bn_decay,
scope='fa_layer2',
is_dist=True)
l1_net = pointnet_fp_module(new_point_cloud_1,
new_point_cloud_2,
l1_net,
l2_net, [256, 128],
is_training,
bn_decay,
scope='fa_layer3',
is_dist=True)
l0_net = pointnet_fp_module(point_cloud_0,
new_point_cloud_1,
l0_net,
l1_net, [128, 128, 128],
is_training,
bn_decay,
scope='fa_layer4',
is_dist=True)
net = tf.expand_dims(l0_net, axis=-2)
# FC layers
net = tf_util.conv2d(net,
128, [1, 1],
padding='VALID',
stride=[1, 1],
bn=True,
is_training=is_training,
scope='fc1',
bn_decay=bn_decay,
is_dist=True)
end_points['feats'] = net
net = tf_util.dropout(net,
keep_prob=0.5,
is_training=is_training,
scope='dp1')
net = tf_util.conv2d(net,
13, [1, 1],
padding='VALID',
activation_fn=None,
scope='fc2',
is_dist=True)
net = tf.squeeze(net, [2])
return net
def cache_file(self, src_mesh_path, src_label_path, cache_path):
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
with tf.device('/gpu:0'):
pc_tf = tf.placeholder(tf.float32)
ind_tf = tf_sampling.farthest_point_sample(self.npoint, pc_tf)
ind_ins_tf = tf_sampling.farthest_point_sample(self.npoint_ins, pc_tf)
sess = tf.Session(config=config)
nfile = len(self.file_list)
#### valid class ids defined in ScanNet
VALID_CLASS_IDS = np.array([3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 16, 24, 28, 33, 34, 36, 39])
target_sem_idx = np.arange(40)
count = 0
for i in range(40):
if i in VALID_CLASS_IDS:
count += 1
target_sem_idx[i] = count
else:
target_sem_idx[i] = 0
for index in range(nfile):
print(np.float32(index)/nfile)
curpc_color = io_util.read_color_ply(os.path.join(src_mesh_path, self.file_list[index]+'.ply'))
curpc = curpc_color[:,:3].astype(np.float32)
curcolor = curpc_color[:,3:].astype(np.float32)/255.0
curgroup = io_util.read_label_txt(os.path.join(src_label_path, 'group_'+self.file_list[index]+'.txt')).astype(np.int32)
curseg = io_util.read_label_txt(os.path.join(src_label_path, 'sem_'+self.file_list[index]+'.txt')).astype(np.int32)
curseg[curseg>=40] = 0
curseg[curseg<0] = 0
curseg = self.changem(curseg, np.arange(40), target_sem_idx)
curngroup = np.max(curgroup)+1
#### collect instance information
valid_group_indicator = np.zeros(curngroup)
target_idx = np.zeros(1+curngroup)
count = 0
for i in range(curngroup):
if np.sum(curgroup==i)==0:
target_idx[i+1] = 0
elif np.round(np.mean(curseg[curgroup==i])).astype('int32')!=0:
valid_group_indicator[i] = 1
count += 1
target_idx[i+1] = count
else:
target_idx[i+1] = 0
curgroup = self.changem(curgroup, np.arange(-1, curngroup), target_idx).astype('int32')
curgroup[curgroup<0] = 0
curngroup = 1+np.sum(valid_group_indicator).astype('int32') # group zero is background
#### resample each scene to a fix number of points
if curngroup>self.ngroup:
self.ngroup = curngroup
if self.npoint<curpc.shape[0]:
choice = sess.run(ind_tf, feed_dict={pc_tf: np.expand_dims(curpc,0)})[0]
pc = curpc[choice,:]
color = curcolor[choice,:]
group_label = curgroup[choice]
seg_label = curseg[choice]
elif self.npoint==curpc.shape[1]:
pc = copy.deepcopy(curpc)
color = copy.deepcopy(curcolor)
group_label = copy.deepcopy(curgroup)
seg_label = copy.deepcopy(curseg)
else:
choice = np.random.choice(curpc.shape[0], self.npoint - curpc.shape[0])
pc = np.concatenate((curpc,curpc[choice,:]), 0)
color = np.concatenate((curcolor,curcolor[choice,:]), 0)
group_label = np.concatenate((curgroup, curgroup[choice]), 0)
seg_label = np.concatenate((curseg, curseg[choice]), 0)
#### resample each instance to a fix number of points
pc_ins = np.zeros((curngroup, self.npoint_ins, 3), dtype=np.float32)
for j in range(1,curngroup):
curins = curpc[curgroup==j,:]
if self.npoint_ins<curins.shape[0]:
choice = sess.run(ind_ins_tf, feed_dict={pc_tf: np.expand_dims(curins,0)})[0]
pc_ins[j,:,:] = curins[choice,:]
elif self.npoint_ins==curins.shape[0]:
pc_ins[j,:,:] = copy.deepcopy(curins)
else:
choice = np.random.choice(curins.shape[0], self.npoint_ins - curins.shape[0])
pc_ins[j,:,:] = np.concatenate((curins,curins[choice,:]), 0)
#### data tuple for each scene
#### group_label indicates instance label
#### seg_label indicates semantic label
self.data_list[index] = (pc, color, group_label, seg_label, pc_ins, curngroup)
np.savez_compressed(cache_path, data_list=self.data_list, ngroup=self.ngroup)
def meteor_direct_module(xyz,
time,
points,
npoint,
radius,
nsample,
mlp,
mlp2,
group_all,
is_training,
bn_decay,
scope,
bn=True,
pooling='max',
knn=False,
use_xyz=True,
use_nchw=False):
'''
Input:
xyz: (batch_size, ndataset, 3) TF tensor
time: (batch_size, ndataset, 1) TF tensor
points: (batch_size, ndataset, channel) TF tensor
npoint: int32 -- #points sampled in farthest point sampling
radius: list of float32 -- search radiuses in local region
nsample: int32 -- how many points in each local region
mlp: list of int32 -- output size for MLP on each point
mlp2: list of int32 -- output size for MLP on each region
group_all: bool -- group all points into one PC if set true, OVERRIDE
npoint, radius and nsample settings
use_xyz: bool, if True concat XYZ with local point features, otherwise just use point features
use_nchw: bool, if True, use NCHW data format for conv2d, which is usually faster than NHWC format
Return:
new_xyz: (batch_size, npoint, 3) TF tensor
new_points: (batch_size, npoint, mlp[-1] or mlp2[-1]) TF tensor
idx: (batch_size, npoint, nsample) int32 -- indices for local regions
'''
data_format = 'NCHW' if use_nchw else 'NHWC'
sample_idx = None
batch_size = xyz.get_shape()[0].value
with tf.variable_scope(scope) as sc:
##### sample and group with variable radius
sample_idx = farthest_point_sample(npoint, xyz)
new_xyz = gather_point(xyz, sample_idx) # (batch_size, npoint, 3)
new_time = gather_point(time, sample_idx) # (batch_size, npoint, 1)
time_ = tf.reshape(time,
[batch_size, 1, -1]) # (batch_size, 1, ndataset)
new_time_ = tf.abs(new_time - time_) # (batch_size, npoint, ndataset)
radius_ = tf.gather(radius, tf.cast(
new_time_, tf.int32)) # (batch_size, npoint, ndataset)
idx, pts_cnt = query_ball_point_var_rad(radius_, nsample, xyz, new_xyz)
grouped_xyz = group_point(xyz, idx) # (batch_size, npoint, nsample, 3)
grouped_xyz -= tf.tile(tf.expand_dims(new_xyz, 2),
[1, 1, nsample, 1]) # translation normalization
if points is not None:
grouped_points = group_point(
points, idx) # (batch_size, npoint, nsample, channel)
grouped_time = group_point(
time, idx) # (batch_size, npoint, nsample, channel)
if use_xyz:
new_points = tf.concat(
[grouped_xyz, grouped_points, grouped_time],
axis=-1) # (batch_size, npoint, nample, 3+channel)
else:
new_points = grouped_points
else:
new_points = grouped_xyz
# Point Feature Embedding
for i, num_out_channel in enumerate(mlp):
new_points = tf_util.conv2d(new_points,
num_out_channel, [1, 1],
padding='VALID',
stride=[1, 1],
bn=bn,
is_training=is_training,
scope='conv%d' % (i),
bn_decay=bn_decay,
data_format=data_format)
new_points = tf.reduce_max(new_points, axis=[2], name='maxpool')
return new_xyz, new_time, new_points, idx
def sample_and_group(npoint,
radius,
nsample,
xyz,
points,
knn=False,
use_xyz=True):
'''
Input:
npoint: int32
radius: float32
nsample: int32
xyz: (batch_size, ndataset, 3) TF tensor
points: (batch_size, ndataset, channel) TF tensor, if None will just use xyz as points
knn: bool, if True use kNN instead of radius search
use_xyz: bool, if True concat XYZ with local point features, otherwise just use point features
Output:
new_xyz: (batch_size, npoint, 3) TF tensor
new_points: (batch_size, npoint, nsample, 3+channel) TF tensor
idx: (batch_size, npoint, nsample) TF tensor, indices of local points as in ndataset points
grouped_xyz: (batch_size, npoint, nsample, 3) TF tensor, normalized point XYZs
(subtracted by seed point XYZ) in local regions
'''
new_xyz = gather_point(xyz, farthest_point_sample(
npoint, xyz)) # (batch_size, npoint, 3) B*N*3
if knn:
_, idx = knn_point(nsample, xyz, new_xyz)
else:
idx, pts_cnt = query_ball_point(radius, nsample, xyz, new_xyz)
# idx (batch_size, npoint, nsample) pts_cnt (batch_size, npoint)
grouped_xyz = group_point(xyz, idx) # (batch_size, npoint, nsample, 3)
G_p = get_neigh_geo_feat(new_xyz, grouped_xyz, nsample) # B * N * 32 *12
# batch= xyz.get_shape()[0].value
# one=tf.ones([batch,npoint],dtype=tf.float32)
# V=tf.constant(4/3*math.pi*radius*radius*radius)
#
# D=tf.divide(tf.cast(pts_cnt,tf.float32),tf.cast(V,tf.float32)) # (batch_size, npoint) maybe need to be normalized
# D=tf.expand_dims(D, -1)
# D = tf.tile(tf.expand_dims(D, -1), [1, 1, nsample, 1]) # 16,1024,32,1
#
# G_p=tf.concat([G_p,D], axis=-1) # 16,10244,32,13
# now, grouped_xyz is local coordinate
grouped_xyz -= tf.tile(tf.expand_dims(new_xyz, 2),
[1, 1, nsample, 1]) # translation normalization
if points is not None:
grouped_points = group_point(
points, idx) # (batch_size, npoint, nsample, channel)
if use_xyz:
new_points = tf.concat(
[grouped_xyz, grouped_points],
axis=-1) # (batch_size, npoint, nample, 3+channel)
else:
new_points = grouped_points
else:
new_points = grouped_xyz
return new_xyz, new_points, idx, grouped_xyz, G_p
def pointnet_sa_module_msg(xyz, points, npoint, radius_list, nsample_list, mlp_list, \
is_training, bn_decay, scope, bn=True, use_xyz=True, use_nchw=False):
'''
new pointnet set abstraction (sa) module with multi-scale grouping (msg)
'''
data_format = 'NCHW' if use_nchw else 'NHWC'
with tf.variable_scope(scope) as sc:
input_points = xyz
point_cloud_shape = points.get_shape()
batch_size = point_cloud_shape[0].value
# sampled_idx = tf.random_uniform(shape=(batch_size,npoint),maxval=npoint-1,dtype=tf.int32)
sampled_idx = farthest_point_sample(npoint, xyz)
new_xyz = gather_point(xyz, sampled_idx)
sampled_idx = tf.expand_dims(sampled_idx, -1)
new_points_list = []
for i in range(len(radius_list)):
input_points = xyz
if points is not None:
if use_xyz:
input_points = tf.concat([input_points, points], axis=-1)
else:
input_points = points
else:
input_points = xyz
# fit for mlp
input_points = tf.expand_dims(input_points, -2)
print("[MSG-MLP]", input_points.shape, input_points.dtype)
if use_nchw:
input_points = tf.transpose(input_points, [0, 3, 1, 2])
for j, num_out_channel in enumerate(mlp_list[i][0:1]):
input_points = tf_util.conv2d(input_points,
num_out_channel, [1, 1],
padding='VALID',
stride=[1, 1],
bn=bn,
is_training=is_training,
scope='conv%d_%d' % (i, j),
bn_decay=bn_decay)
if use_nchw:
input_points = tf.transpose(input_points, [0, 2, 3, 1])
radius = radius_list[i]
nsample = nsample_list[i]
idx, _ = query_ball_point(radius, nsample, xyz, new_xyz)
# recover for grouping
input_points = tf.squeeze(input_points, -2)
sampled_points = new_group_point(input_points, sampled_idx)
new_points = new_group_point(input_points, idx)
new_points -= sampled_points
if use_nchw: new_points = tf.transpose(new_points, [0, 3, 1, 2])
for j, num_out_channel in enumerate(mlp_list[i][1:]):
new_points = tf_util.conv2d(new_points,
num_out_channel, [1, 1],
padding='VALID',
stride=[1, 1],
bn=bn,
is_training=is_training,
scope='conv%d_%d' % (i, j + 1),
bn_decay=bn_decay)
if use_nchw: new_points = tf.transpose(new_points, [0, 2, 3, 1])
# sampled_points = tf.squeeze(sampled_points, -2)
# new_points -= sampled_points
new_points = tf.reduce_max(new_points, axis=[2])
# print(tf.shape(input_points), tf.shape(new_points))
# sampled_points = gather_point(input_points, sampled_idx)
new_points_list.append(new_points)
new_points_concat = tf.concat(new_points_list, axis=-1)
print("[MSG-MLP] output:", new_points_concat.shape)
return new_xyz, new_points_concat
def extract_patches(batch_xyz,
k,
patch_num=1,
batch_features=None,
gt_xyz=None,
gt_k=None,
is_training=None):
"""
:param batch_xyz [B, P, 3]
"""
batch_size, num_point, _ = batch_xyz.shape.as_list()
with tf.name_scope("extract_input"):
if is_training:
use_random = False
if patch_num > 1:
batch_seed_point = gather_point(
batch_xyz, farthest_point_sample(patch_num, batch_xyz))
else:
# B, 1, 3
idx = tf.random_uniform([batch_size, patch_num],
minval=0,
maxval=num_point,
dtype=tf.int32)
# idx = tf.constant(250, shape=[batch_size, 1], dtype=tf.int32)
batch_seed_point = gather_point(batch_xyz, idx)
#patch_num = 1
else:
assert (batch_size == 1)
# remove residual, (B P 1) and (B, P, 1, 2)
closest_d, _ = knn_point_2(2, batch_xyz, batch_xyz, unique=False)
closest_d = closest_d[:, :, 1:]
# (B, P)
mask = tf.squeeze(
closest_d < 5 *
(tf.reduce_mean(closest_d, axis=1, keepdims=True)),
axis=-1)
# filter (B, P', 3)
batch_xyz = tf.expand_dims(tf.boolean_mask(batch_xyz, mask),
axis=0)
# batch_xyz = tf.Print(batch_xyz, [tf.shape(batch_xyz)])
# B, M, 3
# batch_seed_point = batch_xyz[:, -1:, :]
# patch_num = 1
patch_num = int(num_point / k * 5)
# idx = tf.random_uniform([batch_size, patch_num], minval=0, maxval=num_point, dtype=tf.int32)
idx = tf.squeeze(farthest_point_sample(patch_num, batch_xyz),
axis=0)
# idx = tf.random_uniform([patch_num], minval=0, maxval=tf.shape(batch_xyz)[1], dtype=tf.int32)
# B, P, 3 -> B, k, 3 (idx B, k, 1)
# idx = tf.Print(idx, [idx], message="idx")
batch_seed_point = tf.gather(batch_xyz, idx, axis=1)
k = tf.minimum(k, tf.shape(batch_xyz)[1])
# batch_seed_point = gather_point(batch_xyz, idx)
# B, M, k, 2
_, new_patch_idx = knn_point_2(k,
batch_xyz,
batch_seed_point,
unique=False)
# B, M, k, 3
batch_xyz = tf.gather_nd(batch_xyz, new_patch_idx)
# MB, k, 3
batch_xyz = tf.concat(tf.unstack(batch_xyz, axis=1), axis=0)
if batch_features is not None:
with tf.name_scope("extract_feature"):
batch_features = tf.gather_nd(batch_features, new_patch_idx)
batch_features = tf.concat(tf.unstack(batch_features, axis=1),
axis=0)
if is_training and (gt_xyz is not None and gt_k is not None):
with tf.name_scope("extract_gt"):
_, new_patch_idx = knn_point_2(gt_k,
gt_xyz,
batch_seed_point,
unique=False)
gt_xyz = tf.gather_nd(gt_xyz, new_patch_idx)
gt_xyz = tf.concat(tf.unstack(gt_xyz, axis=1), axis=0)
else:
gt_xyz = None
return batch_xyz, batch_features, gt_xyz
def __init__(self, points, features, is_training, setting):
xconv_params = setting.xconv_params
fc_params = setting.fc_params
with_X_transformation = setting.with_X_transformation
sorting_method = setting.sorting_method
network_name = setting.network_name
N = tf.shape(points)[0]
#if setting.sampling == 'fps':
# from sampling import tf_sampling
with tf.variable_scope(network_name):
self.layer_pts = [points]
if features is None:
self.layer_fts = [features]
else:
features = tf.reshape(features, (N, -1, setting.data_dim - 3),
name='features_reshape')
C_fts = xconv_params[0]['C'] // 2
features_hd = pf.dense(features, C_fts, 'features_hd',
is_training)
self.layer_fts = [features_hd]
for layer_idx, layer_param in enumerate(xconv_params):
tag = 'xconv_' + str(layer_idx + 1) + '_'
K = layer_param['K']
D = layer_param['D']
P = layer_param['P']
C = layer_param['C']
links = layer_param['links']
if setting.sampling != 'random' and links:
print(
'Error: flexible links are supported only when random sampling is used!'
)
exit()
# get k-nearest points
pts = self.layer_pts[-1]
fts = self.layer_fts[-1]
if P == -1 or (layer_idx > 0
and P == xconv_params[layer_idx - 1]['P']):
qrs = self.layer_pts[-1]
else:
if setting.sampling == 'fps':
fps_indices = tf_sampling.farthest_point_sample(P, pts)
batch_indices = tf.tile(
tf.reshape(tf.range(N), (-1, 1, 1)), (1, P, 1))
indices = tf.concat(
[batch_indices,
tf.expand_dims(fps_indices, -1)],
axis=-1)
qrs = tf.gather_nd(pts, indices,
name=tag + 'qrs') # (N, P, 3)
elif setting.sampling == 'ids':
indices = pf.inverse_density_sampling(pts, K, P)
qrs = tf.gather_nd(pts, indices)
elif setting.sampling == 'random':
qrs = tf.slice(pts, (0, 0, 0), (-1, P, -1),
name=tag + 'qrs') # (N, P, 3)
else:
print('Unknown sampling method!')
exit()
self.layer_pts.append(qrs)
if layer_idx == 0:
C_pts_fts = C // 2 if fts is None else C // 4
depth_multiplier = 4
else:
C_prev = xconv_params[layer_idx - 1]['C']
C_pts_fts = C_prev // 4
depth_multiplier = math.ceil(C / C_prev)
with_global = (setting.with_global
and layer_idx == len(xconv_params) - 1)
fts_xconv = xconv(pts, fts, qrs, tag, N, K, D, P, C, C_pts_fts,
is_training, with_X_transformation,
depth_multiplier, sorting_method,
with_global)
fts_list = []
for link in links:
fts_from_link = self.layer_fts[link]
if fts_from_link is not None:
fts_slice = tf.slice(
fts_from_link, (0, 0, 0), (-1, P, -1),
name=tag + 'fts_slice_' + str(-link))
fts_list.append(fts_slice)
if fts_list:
fts_list.append(fts_xconv)
self.layer_fts.append(
tf.concat(fts_list,
axis=-1,
name=tag + 'fts_list_concat'))
else:
self.layer_fts.append(fts_xconv)
if hasattr(setting, 'xdconv_params'):
for layer_idx, layer_param in enumerate(setting.xdconv_params):
tag = 'xdconv_' + str(layer_idx + 1) + '_'
K = layer_param['K']
D = layer_param['D']
pts_layer_idx = layer_param['pts_layer_idx']
qrs_layer_idx = layer_param['qrs_layer_idx']
pts = self.layer_pts[pts_layer_idx + 1]
fts = self.layer_fts[
pts_layer_idx +
1] if layer_idx == 0 else self.layer_fts[-1]
qrs = self.layer_pts[qrs_layer_idx + 1]
fts_qrs = self.layer_fts[qrs_layer_idx + 1]
P = xconv_params[qrs_layer_idx]['P']
C = xconv_params[qrs_layer_idx]['C']
C_prev = xconv_params[pts_layer_idx]['C']
C_pts_fts = C_prev // 4
depth_multiplier = 1
fts_xdconv = xconv(pts, fts, qrs, tag, N, K, D, P, C,
C_pts_fts, is_training,
with_X_transformation, depth_multiplier,
sorting_method)
fts_concat = tf.concat([fts_xdconv, fts_qrs],
axis=-1,
name=tag + 'fts_concat')
fts_fuse = pf.dense(fts_concat, C, tag + 'fts_fuse',
is_training)
self.layer_pts.append(qrs)
self.layer_fts.append(fts_fuse)
self.fc_layers = [self.layer_fts[-1]]
for layer_idx, layer_param in enumerate(fc_params):
C = layer_param['C']
dropout_rate = layer_param['dropout_rate']
fc = pf.dense(self.fc_layers[-1], C,
'fc{:d}'.format(layer_idx), is_training)
fc_drop = tf.layers.dropout(
fc,
dropout_rate,
training=is_training,
name='fc{:d}_drop'.format(layer_idx))
self.fc_layers.append(fc_drop)
all_points = data_utils.px2cam_w_sample(intrinsics, depth_raw)
# hand segmentation from the input 3D scene
# depth-based segmentation, assume hand is the cloest object to the camera
hand_points = all_points[all_points[:, 2] < 800]
# OBB-based hand points preprocessing
hand_points_rotate, coeff = data_utils.pca_obb(hand_points)
hand_points_normalized_sampled, rand_idx, volume_length, offset = data_utils.normalize_w_sample(
hand_points_rotate, hand_points.shape[0])
# FPS: farthest point sampling
# 1st leveln
hand_point = tf.convert_to_tensor(
hand_points_normalized_sampled[np.newaxis, :],
dtype=tf.float32)
sampled_idx_l1 = tf_sampling.farthest_point_sample(
sample_num_level1, hand_point).eval().squeeze()
other_idx = np.setdiff1d(np.arange(SAMPLE_NUM), sampled_idx_l1)
new_idx = np.concatenate((sampled_idx_l1, other_idx))
hand_points_normalized_sampled = hand_points_normalized_sampled[
new_idx, :]
# 2nd level
hand_point = tf.convert_to_tensor(
hand_points_normalized_sampled[:sample_num_level1, :3][
np.newaxis, :],
dtype=tf.float32)
sampled_idx_l2 = tf_sampling.farthest_point_sample(
sample_num_level2, hand_point).eval().squeeze()
other_idx = np.setdiff1d(np.arange(sample_num_level1),
sampled_idx_l2)
new_idx = np.concatenate((sampled_idx_l2, other_idx))
def corrs_flow_pred_net(xyz1,
xyz2,
net1,
net2,
scopename,
reuse,
is_training,
bn_decay,
nsmp=256,
nfea=64):
#########################################
# input
# xyz1, xyz2: (B x N x 3)
# net1, net2: (B x N x nfea)
# output
# pred_flow: (B x N x 3)
# pred_vismask: (B x N)
# fpsidx1, fpsidx2: (B x nsmp)
# matching_score_sub: (B x nsmp x nsmp)
#########################################
num_point = xyz1.get_shape()[1].value
with tf.variable_scope(scopename) as myscope:
if reuse:
myscope.reuse_variables()
# sub-sample to predict vismask and flow
fpsidx1 = farthest_point_sample(nsmp, xyz1) # Bxnsmp
idx = tf.where(tf.greater_equal(fpsidx1, 0))
fpsidx1 = tf.concat((tf.expand_dims(tf.cast(
idx[:, 0], tf.int32), -1), tf.reshape(fpsidx1, [-1, 1])), 1)
xyz1_sub = tf.reshape(tf.gather_nd(xyz1, fpsidx1), [-1, nsmp, 3])
net1_sub = tf.reshape(tf.gather_nd(net1, fpsidx1), [-1, nsmp, nfea])
fpsidx2 = farthest_point_sample(nsmp, xyz2) # Bxnsmp
idx = tf.where(tf.greater_equal(fpsidx2, 0))
fpsidx2 = tf.concat((tf.expand_dims(tf.cast(
idx[:, 0], tf.int32), -1), tf.reshape(fpsidx2, [-1, 1])), 1)
xyz2_sub = tf.reshape(tf.gather_nd(xyz2, fpsidx2), [-1, nsmp, 3])
net2_sub = tf.reshape(tf.gather_nd(net2, fpsidx2), [-1, nsmp, nfea])
net_combined_sub = tf.concat(
(tf.tile(tf.expand_dims(net1_sub, 2), [1, 1, nsmp, 1]),
tf.tile(tf.expand_dims(net2_sub, 1), [1, nsmp, 1, 1])), -1)
mlp_maskpred = [128, 128, 128]
for i, num_out_channel in enumerate(mlp_maskpred):
net_combined_sub = tf_util.conv2d(net_combined_sub,
num_out_channel, [1, 1],
padding='VALID',
stride=[1, 1],
bn=True,
is_training=is_training,
scope='conv%d_maskpred' % (i),
bn_decay=bn_decay)
pred_vismask_sub = tf.reduce_max(net_combined_sub, 2, keep_dims=True)
mlp2_maskpred = [128, 64, 32]
for i, num_out_channel in enumerate(mlp2_maskpred):
pred_vismask_sub = tf_util.conv2d(pred_vismask_sub,
num_out_channel, [1, 1],
padding='VALID',
stride=[1, 1],
bn=True,
is_training=is_training,
scope='conv_post_%d_maskpred' %
(i),
bn_decay=bn_decay)
pred_vismask_sub = tf_util.conv2d(pred_vismask_sub,
1, [1, 1],
padding='VALID',
stride=[1, 1],
scope='conv_out_maskpred',
activation_fn=None)
pred_vismask_sub = tf.squeeze(pred_vismask_sub, [2])
pred_vismask = pointnet_fp_module(xyz1,
xyz1_sub,
None,
pred_vismask_sub, [],
tf.constant(True),
None,
scope='interp_layer')
pred_vismask = tf.squeeze(pred_vismask, 2) # B x nsmp
pred_vismask_sub = tf.stop_gradient(
tf.sigmoid(pred_vismask_sub)) # B x nsmp x 1
mlp0 = [8]
for i, num_out_channel in enumerate(mlp0):
net_combined_sub = tf_util.conv2d(net_combined_sub,
num_out_channel, [1, 1],
padding='VALID',
stride=[1, 1],
bn=True,
is_training=is_training,
scope='conv_prev_%d' % (i),
bn_decay=bn_decay)
net_combined_sub = tf_util.conv2d(net_combined_sub,
1, [1, 1],
padding='VALID',
stride=[1, 1],
scope='conv_prev_3',
activation_fn=None)
U = tf.nn.softmax(net_combined_sub, 2) # B x nsmp x nsmp x 1
matching_score_sub = tf.squeeze(net_combined_sub, -1)
#### mask prob
U = tf.concat(
(tf.multiply(U, tf.expand_dims(pred_vismask_sub, 2)),
tf.expand_dims(xyz2_sub, 1) - tf.expand_dims(xyz1_sub, 2)), -1)
mlp = [32, 64, 128, 256]
for i, num_out_channel in enumerate(mlp):
U = tf_util.conv2d(U,
num_out_channel, [1, 1],
padding='VALID',
stride=[1, 1],
bn=True,
is_training=is_training,
scope='conv%d' % (i),
bn_decay=bn_decay)
U = tf.reduce_max(U, 2)
l1_xyz = xyz1_sub
#### mask energy
l1_points = tf.concat((U, pred_vismask_sub), -1)
l2_xyz, l2_points, l2_indices = pointnet_sa_module(
l1_xyz,
l1_points,
npoint=128,
radius=0.4,
nsample=32,
mlp=[128, 128, 256],
mlp2=None,
group_all=False,
is_training=is_training,
bn_decay=bn_decay,
scope='corrs_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,
use_xyz=False,
is_training=is_training,
bn_decay=bn_decay,
scope='corrs_layer3')
# Feature Propagation layers
l2_points = pointnet_fp_module(l2_xyz,
l3_xyz,
l2_points,
l3_points, [256, 256],
is_training,
bn_decay,
scope='corrs_fa_layer1')
l1_points = pointnet_fp_module(l1_xyz,
l2_xyz,
l1_points,
l2_points, [256, 128],
is_training,
bn_decay,
scope='corrs_fa_layer2')
l0_points = pointnet_fp_module(xyz1,
l1_xyz,
None,
l1_points, [128, 128, 64],
is_training,
bn_decay,
scope='corrs_fa_layer3')
# FC layers
net = tf_util.conv1d(l0_points,
64,
1,
padding='VALID',
bn=True,
is_training=is_training,
scope='corrs_fc1',
bn_decay=bn_decay)
net = tf_util.conv1d(net,
3,
1,
padding='VALID',
activation_fn=None,
scope='corrs_fc2')
pred_flow = tf.reshape(net, [-1, num_point, 3])
return pred_flow, pred_vismask, fpsidx1, fpsidx2, matching_score_sub
def setup_graph(general_opts, train_opts, hyperparameters):
tf.reset_default_graph()
now = datetime.now()
target_cls_choosen = general_opts['target_cls']
for i in valid_seq_id[target_cls_choosen]:
filename = str(i).zfill(4) + "_pcnn.tfrecord"
valid_file_lists.append(filename)
valid_file_list = [
os.path.join(data_dir, file) for file in valid_file_lists
]
BATCH_SIZE = hyperparameters['batch_size']
NUM_POINT = general_opts['num_point']
GPU_INDEX = general_opts['gpu']
MODEL = importlib.import_module(
general_opts['model']) # import network module
minimum_points_in_segment = NUM_POINT
tf.set_random_seed(123456789)
# double check
BN_INIT_DECAY = 0.5
BN_DECAY_DECAY_RATE = 0.5
BN_DECAY_DECAY_STEP = float(40)
BN_DECAY_CLIP = 0.99
# threshold distance: for class index i, remove points further away from segment mean than threshold_distance_per_class[i]
threshold_distance_per_class = 0.2 * np.ones(
(NUM_CLASS, ), dtype=np.float32)
with tf.Graph().as_default():
with tf.device('/cpu:0'):
with tf.name_scope('prepare_data'):
start_time_seg = datetime.now()
val_datasets = [
create_tfrecord_dataset(f, NUM_POINT,
minimum_points_in_segment,
threshold_distance_per_class,
target_cls_choosen)
for f in valid_file_list
]
val_dataset = tf.data.experimental.sample_from_datasets(
val_datasets)
val_dataset = val_dataset.batch(
BATCH_SIZE, drop_remainder=False).prefetch(1)
val_iterator = val_dataset.make_initializable_iterator()
iter_handle = tf.placeholder(tf.string,
shape=[],
name='iterator_handle')
iterator = tf.data.Iterator.from_string_handle(
iter_handle, val_dataset.output_types,
val_dataset.output_shapes)
next_element = iterator.get_next()
next_element = reshape_element(next_element,
batch_size=BATCH_SIZE,
num_point=NUM_POINT)
time_elapsed_seg = datetime.now() - start_time_seg
print 'seg time elapsed (hh:mm:ss) {}'.format(time_elapsed_seg)
with tf.device('/gpu:' + str(GPU_INDEX)):
is_training_pl_encoder = tf.placeholder(tf.bool, shape=())
is_training_pl = tf.placeholder(tf.bool, shape=())
print(is_training_pl)
batch = tf.Variable(0.)
bn_momentum = tf.train.exponential_decay(BN_INIT_DECAY,
batch * BATCH_SIZE,
BN_DECAY_DECAY_STEP,
BN_DECAY_DECAY_RATE,
staircase=True)
bn_decay = tf.minimum(BN_DECAY_CLIP, 1 - bn_momentum)
tf.summary.scalar('bn_decay', bn_decay)
next_element_xyz_inlier = next_element[
'xyz_inlier'][:, 0:NUM_POINT, :]
next_element_xyz_inlier = tf.reshape(next_element_xyz_inlier,
[BATCH_SIZE, NUM_POINT, 3])
cls_gt_onehot = tf.one_hot(indices=next_element['class_id'],
depth=len(target_cls))
cls_gt_onehot_expand = tf.expand_dims(cls_gt_onehot, axis=1)
cls_gt_onehot_tile = tf.tile(cls_gt_onehot_expand,
[1, NUM_POINT, 1])
visiblePoints = tf.reshape(next_element['visiblePoints_org'],
[BATCH_SIZE, 2048 + 1, 3])
visiblePoints_final = visiblePoints[:, 0:NUM_POINT, :]
xyz_graph_input = next_element_xyz_inlier
with tf.name_scope('6d_pose'):
element_mean = tf.reduce_mean(xyz_graph_input, axis=1)
next_element_xyz_inlier_normalized = xyz_graph_input - tf.expand_dims(
element_mean, 1)
# dgcnn
xyz_recon_res, rot_pred, trans_pred_res, endpoint = MODEL.get_model_dgcnn_mean_6d(
tf.concat([
next_element_xyz_inlier_normalized, cls_gt_onehot_tile
],
axis=2),
is_training_pl_encoder,
is_training_pl,
10,
bn_decay=bn_decay)
xyz_recon = xyz_recon_res + tf.tile(
tf.expand_dims(element_mean, 1),
[1, xyz_recon_res.shape[1], 1])
trans_pred = trans_pred_res + element_mean
# for all decoder
xyz_recon_FPS = gather_point(
xyz_recon, farthest_point_sample(NUM_POINT, xyz_recon))
xyz_loss, _ = chamfer_loss.get_loss(xyz_recon_FPS,
visiblePoints_final)
tf.summary.scalar('chamfer_loss', xyz_loss)
with tf.name_scope('translation'):
trans_loss, trans_loss_perSample = trans_distance.get_translation_error(
trans_pred, next_element['translation'])
mean_dist_loss, mean_dist_loss_perSample = trans_distance.get_translation_error(
element_mean, next_element['translation'])
tf.summary.scalar('trans_loss', trans_loss)
tf.summary.scalar('mean_dist_loss', mean_dist_loss)
tf.summary.scalar('trans_loss_min',
tf.reduce_min(trans_loss_perSample))
tf.summary.scalar('trans_loss_max',
tf.reduce_max(trans_loss_perSample))
xyz_remove_trans = next_element['xyz'] - tf.expand_dims(trans_pred,
axis=1)
with tf.name_scope('rotation'):
current_batch_axag = tf.reshape(next_element['axisangle'],
(1, 3))
rot_pred = tf.cast(rot_pred, tf.float64)
axag_loss, axag_loss_perSample = angular_distance_taylor.get_rotation_error(
rot_pred, tf.cast(current_batch_axag, tf.float64))
axag_loss = tf.cast(axag_loss, tf.float32)
tf.summary.scalar('axag_loss', axag_loss)
tf.summary.scalar('axag_loss_min',
tf.reduce_min(axag_loss_perSample))
tf.summary.scalar('axag_loss_max',
tf.reduce_max(axag_loss_perSample))
# Create a session
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
config.log_device_placement = False
sess = tf.Session(config=config)
# Add summary writers
merged = tf.summary.merge_all()
# Init variables
init = tf.global_variables_initializer()
sess.run(init, {is_training_pl_encoder: False, is_training_pl: False})
# Add ops to save and restore all the variables.
saver = tf.train.Saver(max_to_keep=None)
# Restore variables from disk.
trained_model = general_opts['trained_model']
saver.restore(sess, trained_model)
print "Model restored."
ops = {
'is_training_pl': is_training_pl,
'is_training_pl_encoder': is_training_pl_encoder,
'trans_loss': trans_loss,
'axag_loss': axag_loss,
'merged': merged,
'step': batch,
'class_id': next_element['class_id'],
'seq_id': next_element['seq_id'],
'frame_id': next_element['frame_id'],
'obj_batch': tf.squeeze(next_element['obj_batch'], axis=1),
'trans_pred': trans_pred,
'rot_pred': rot_pred,
'xyz_recon': xyz_recon_FPS,
'xyz_inlier_full': next_element['xyz_inlier_full'],
'xyz_graph_input': xyz_graph_input,
'handle': iter_handle
}
validation_handle = sess.run(val_iterator.string_handle())
sess.run(val_iterator.initializer)
model, class_id = eval_graph(sess,
ops,
validation_handle,
batch_size=BATCH_SIZE)
def resample(pointcloud, npoint):
new_pointcloud = gather_point(pointcloud,
farthest_point_sample(npoint, pointcloud))
return new_pointcloud
def grouping_pred_net(xyz,
flow,
trans,
scopename,
reuse,
is_training,
bn_decay=None,
nsmp=128):
########################################
# input
# xyz: (B x N x 3)
# flow: (B x N x 3)
# trans: (B x N x nfea)
# output
# pred_grouping_sub: (B x nsmp x nsmp) - logits
# fpsidx: (B x nsmp)
########################################
num_point = xyz.get_shape()[1].value
nfea = trans.get_shape()[2].value
with tf.variable_scope(scopename) as myscope:
if reuse:
myscope.reuse_variables()
# Grouping
fpsidx = farthest_point_sample(nsmp, xyz) # Bxnsmp
idx = tf.where(tf.greater_equal(fpsidx, 0))
fpsidx = tf.concat((tf.expand_dims(tf.cast(
idx[:, 0], tf.int32), -1), tf.reshape(fpsidx, [-1, 1])), 1)
xyz_sub = tf.reshape(tf.gather_nd(xyz, fpsidx), [-1, nsmp, 3])
flow_sub = tf.reshape(tf.gather_nd(flow, fpsidx), [-1, nsmp, 3])
pred_sub = tf.reshape(tf.gather_nd(trans, fpsidx), [-1, nsmp, nfea])
Rs = tf.reshape(pred_sub[:, :, :9], [-1, nsmp, 3, 3])
ts = tf.reshape(pred_sub[:, :, 9:], [-1, nsmp, 1, 3])
ppdist = tf.expand_dims(xyz_sub, 1) - tf.expand_dims(
xyz_sub, 2) # B x nsmp x nsmp x 3
ppdist = tf.matmul(ppdist, Rs) + ts + tf.expand_dims(
flow_sub, 2) # B x nsmp x nsmp x 3
U = tf.concat(
(tf.tile(tf.expand_dims(xyz_sub, 1),
(1, nsmp, 1, 1)), ppdist - tf.expand_dims(flow_sub, 1)),
-1) # B x nsmp x nsmp x 6
mlp = [16, 64, 512]
for i, num_out_channel in enumerate(mlp):
U = tf_util.conv2d(U,
num_out_channel, [1, 1],
padding='VALID',
stride=[1, 1],
bn=True,
is_training=is_training,
scope='conv_stage1_%d' % (i),
bn_decay=bn_decay)
U_glb = tf.reduce_max(U, 2, keep_dims=True)
mlp2 = [256, 256, 256]
for i, num_out_channel in enumerate(mlp2):
U_glb = tf_util.conv2d(U_glb,
num_out_channel, [1, 1],
padding='VALID',
stride=[1, 1],
bn=True,
is_training=is_training,
scope='conv_stage2_%d' % (i),
bn_decay=bn_decay)
U_combined = tf.concat((tf.tile(U_glb, (1, 1, nsmp, 1)), U), -1)
mlp3 = [256, 64, 16]
for i, num_out_channel in enumerate(mlp3):
U_combined = tf_util.conv2d(U_combined,
num_out_channel, [1, 1],
padding='VALID',
stride=[1, 1],
bn=True,
is_training=is_training,
scope='conv_stage3_%d' % (i),
bn_decay=bn_decay)
U_combined = tf_util.conv2d(U_combined,
1, [1, 1],
padding='VALID',
stride=[1, 1],
scope='conv_stage3_3',
activation_fn=None) # B x nsmp x nsmp x 1
pred_grouping_sub = tf.squeeze(U_combined, -1) # B x nsmp x nsmp
return pred_grouping_sub, fpsidx
def pointnet_sa_module_msg3(gt,
pred,
npoint,
radius_list,
nsample_list,
mlp_list,
scope,
use_xyz=True,
use_nchw=False,
knn=True):
''' PointNet Set Abstraction (SA) module with Multi-Scale Grouping (MSG)
Input:
xyz: (batch_size, ndataset, 3) TF tensor
points: (batch_size, ndataset, channel) TF tensor
npoint: int32 -- #points sampled in farthest point sampling
radius: list of float32 -- search radius in local region
nsample: list of int32 -- how many points in each local region
mlp: list of list of int32 -- output size for MLP on each point
use_xyz: bool, if True concat XYZ with local point features, otherwise just use point features
use_nchw: bool, if True, use NCHW data format for conv2d, which is usually faster than NHWC format
Return:
new_xyz: (batch_size, npoint, 3) TF tensor
new_points: (batch_size, npoint, \sum_k{mlp[k][-1]}) TF tensor
'''
with tf.variable_scope(scope) as sc:
p1_idx = farthest_point_sample(npoint, gt)
source_gt = gather_point(gt, p1_idx)
new_points_list = []
for i in range(len(radius_list)):
radius = radius_list[i]
nsample = nsample_list[i]
if knn:
_, idx_gt = knn_point(nsample, gt, source_gt)
else:
idx_gt, _ = query_ball_point(radius, nsample, gt, source_gt)
grouped_gt = group_point(gt, idx_gt) #b*n*k*3
grouped_gt -= tf.tile(tf.expand_dims(source_gt, 2),
[1, 1, nsample, 1])
if knn:
_, idx_pred = knn_point(nsample, pred, source_gt)
else:
idx_pred, _ = query_ball_point(radius, nsample, pred,
source_gt)
grouped_pred = group_point(pred, idx_pred) # b*n*k*3
grouped_pred -= tf.tile(tf.expand_dims(source_gt, 2),
[1, 1, nsample, 1])
grouped_points = tf.concat([grouped_gt, grouped_pred], axis=2)
for j, num_out_channel in enumerate(mlp_list[i]):
grouped_points = conv2d(grouped_points,
num_out_channel, [1, 1],
weight_decay=0,
padding='VALID',
stride=[1, 1],
scope='conv%d_%d' % (i, j),
activation_fn=tf.nn.leaky_relu)
#new_points = tf.reduce_max(grouped_points, axis=[2]) #b*n*c
new_points = tf.layers.max_pooling2d(grouped_points, [1, nsample],
strides=[1, nsample],
padding='VALID',
name='maxpool_%d' % i)
new_points_list.append(new_points)
new_points_concat = tf.concat(new_points_list, axis=-1)
return source_gt, new_points_concat
def sample_and_group(npoint,
radius,
nsample,
xyz,
points,
bn,
is_training,
bn_decay,
mlp,
knn=False,
use_xyz=True,
xyz_feature=None,
end=False,
use_edge_feature=False):
'''
Input:
npoint: int32
radius: float32
nsample: int32
xyz: (batch_size, ndataset, 3) TF tensor
points: (batch_size, ndataset, channel) TF tensor, if None will just use xyz as points
knn: bool, if True use kNN instead of radius search
use_xyz: bool, if True concat XYZ with local point features, otherwise just use point features
Output:
new_xyz: (batch_size, npoint, 3) TF tensor
new_points: (batch_size, npoint, nsample, 3+channel) TF tensor
idx: (batch_size, npoint, nsample) TF tensor, indices of local points as in ndataset points
grouped_xyz: (batch_size, npoint, nsample, 3) TF tensor, normalized point XYZs
(subtracted by seed point XYZ) in local regions
'''
new_xyz = gather_point(xyz, farthest_point_sample(
npoint, xyz)) # (batch_size, npoint, 3)
if knn:
_, idx = knn_point(nsample, xyz, new_xyz)
else:
idx, pts_cnt = query_ball_point(radius, nsample, xyz, new_xyz)
grouped_xyz = group_point(xyz, idx) # (batch_size, npoint, nsample, 3)
grouped_xyz -= tf.tile(tf.expand_dims(new_xyz, 2),
[1, 1, nsample, 1]) # translation normalization
if points is not None:
grouped_points = group_point(
points, idx) # (batch_size, npoint, nsample, channel)
if use_xyz:
new_points = tf.concat(
[grouped_xyz, grouped_points],
axis=-1) # (batch_size, npoint, nample, 3+channel)
else:
new_points = grouped_points
else:
new_points = grouped_xyz
if use_edge_feature == False:
return new_xyz, new_points, idx, grouped_xyz
# [batch_size, npoint, 1, F]
if xyz_feature == None:
xyz_feature = xyz
xyz_feature = group_point(xyz_feature, idx)
edge_feature = grouped_xyz
for i, num_out_channel in enumerate(mlp):
edge_feature = tf_util.conv2d(edge_feature,
num_out_channel, [1, 1],
padding='VALID',
stride=[1, 1],
bn=bn,
is_training=is_training,
scope='xyz_feature_%d' % (i),
bn_decay=bn_decay)
output_feature = tf.concat([xyz_feature, edge_feature], axis=-1)
if end == False:
xyz_feature = tf_util.conv2d(output_feature,
mlp[-1], [1, 1],
padding='VALID',
stride=[1, 1],
bn=bn,
is_training=is_training,
scope='xyz_feature2',
bn_decay=bn_decay)
# we can try sum and mean
xyz_feature = tf.reduce_max(xyz_feature,
axis=[2],
keep_dims=True,
name='maxpool')
xyz_feature = tf.squeeze(xyz_feature, [2])
return new_xyz, new_points, idx, output_feature, xyz_feature, grouped_xyz
