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 get_edge_feature(point_cloud, k=16, idx=None):
"""Construct edge feature for each point
Args:
point_cloud: (batch_size, num_points, 1, num_dims)
nn_idx: (batch_size, num_points, k, 2)
k: int
Returns:
edge features: (batch_size, num_points, k, num_dims)
"""
if idx is None:
_, idx = knn_point_2(k + 1,
point_cloud,
point_cloud,
unique=True,
sort=True)
idx = idx[:, :, 1:, :]
# [N, P, K, Dim]
point_cloud_neighbors = tf.gather_nd(point_cloud, idx)
point_cloud_central = tf.expand_dims(point_cloud, axis=-2)
point_cloud_central = tf.tile(point_cloud_central, [1, 1, k, 1])
edge_feature = tf.concat(
[point_cloud_central, point_cloud_neighbors - point_cloud_central],
axis=-1)
return edge_feature, idx
def get_repulsion_loss(pred,
nsample=20,
radius=0.07,
knn=False,
use_l1=False,
h=0.001):
if knn:
_, idx = knn_point_2(nsample, pred, pred)
pts_cnt = tf.constant(nsample, shape=(30, 1024))
else:
idx, pts_cnt = query_ball_point(radius, nsample, pred, pred)
tf.summary.histogram('smooth/unque_index', pts_cnt)
grouped_pred = group_point(pred, idx) # (batch_size, npoint, nsample, 3)
grouped_pred -= tf.expand_dims(pred, 2)
# get the uniform loss
if use_l1:
dists = tf.reduce_sum(tf.abs(grouped_pred), axis=-1)
else:
dists = tf.reduce_sum(grouped_pred**2, axis=-1)
val, idx = tf.nn.top_k(-dists, 5)
val = val[:, :, 1:] # remove the first one
if use_l1:
h = np.sqrt(h) * 2
print(("h is ", h))
val = tf.maximum(0.0, h + val) # dd/np.sqrt(n)
repulsion_loss = tf.reduce_mean(val)
return repulsion_loss
def get_perulsion_loss(pred, nsample=15, radius=0.07, knn=False, numpoint=512, use_l1=False):
# pred: (batch_size, npoint,3)
if knn:
with tf.device('/gpu:1'):
_, idx = knn_point_2(nsample, pred, pred)
pts_cnt = tf.constant(nsample, shape=(30, numpoint))
else:
idx, pts_cnt = query_ball_point(radius, nsample, pred, pred)
tf.summary.histogram('smooth/unque_index', pts_cnt)
grouped_pred = group_point(pred, idx) # (batch_size, npoint, nsample, 3)
grouped_pred -= tf.expand_dims(pred, 2)
##get the uniform loss
dists = tf.reduce_sum(grouped_pred ** 2, axis=-1)
if use_l1:
dists = tf.sqrt(dists+1e-12)
val, idx = tf.nn.top_k(-dists, 5)
val = val[:, :, 1:] # remove the first one
if use_l1:
h = np.sqrt(0.001)*2
else:
h = 0.01
print ("h is ",h)
val = tf.maximum(0.0, h + val) # dd/np.sqrt(n)
perulsion_loss = tf.reduce_mean(val)
return perulsion_loss