def get_labels_seg(pcl_gt, pcl_pred, metric):
'''
Point wise correspondences between two point sets
args:
pcl_gt: (batch_size, n_pts, 3), gt pcl
pcl_pred: (batch_size, n_pts, 3), predicted pcl
metric: str, 'chamfer' or 'emd'
metric to be considered for returning corresponding points
returns:
pts_match_fwd: gt to pred point-wise correspondence
each point in gt is mapped to nearest point in pred
pts_match_bwd: pred to gt point-wise correspondence
each point in pred is mapped to nearest point in gt
pts_match: one-to-one mapping between pred and gt, acc. to emd
'''
if metric == 'chamfer':
_, pts_match_fwd, _, pts_match_bwd = tf_nndistance.nn_distance(\
pcl_gt, pcl_pred)
return pts_match_fwd, pts_match_bwd
elif metric == 'emd':
pts_match, _ = auction_match(pcl_gt, pcl_pred)
return pts_match
else:
print 'Undefined metric'
return None
def get_metrics(gt_pcl, pred_pcl, args):
'''
Obtain chamfer and emd distances between GT and predicted pcl
args:
gt_pcl: float, (BS,N_PTS,3); GT point cloud
pred_pcl: float, (BS,N_PTS,3); predicted point cloud
returns:
dists_forward: float, (BS); forward chamfer distance
dists_backward: float, (BS); backward chamfer distance
chamfer_distance: float, (BS); chamfer distance
emd: float, (BS); earth mover's distance
'''
dists_forward, _, dists_backward, _ = tf_nndistance.nn_distance(
gt_pcl, pred_pcl)
dists_forward = tf.reduce_mean(
dists_forward, axis=1) # (BATCH_SIZE,args.N_PTS) --> (BATCH_SIZE)
dists_backward = tf.reduce_mean(dists_backward, axis=1)
chamfer_distance = dists_backward + dists_forward
X, _ = tf.meshgrid(range(args.batch_size),
range(args.OUTPUT_PCL_SIZE),
indexing='ij')
ind, _ = auction_match(pred_pcl,
gt_pcl) # Ind corresponds to points in pcl_gt
ind = tf.stack((X, ind), -1)
emd = tf.reduce_mean(
tf.reduce_sum((tf.gather_nd(gt_pcl, ind) - pred_pcl)**2, axis=-1),
axis=1
) # (BATCH_SIZE,args.N_PTS,3) --> (BATCH_SIZE,args.N_PTS) --> (BATCH_SIZE)
return dists_forward, dists_backward, chamfer_distance, emd
def get_loss(reconstructed_points, original_points, type='chamfer', w=0, mu=0, sigma=0,fv=0, add_fv_loss=False):
n_points = reconstructed_points.get_shape()[1].value
d = tf.constant(0)
matched_out = tf.constant(0)
if type == 'chamfer':
#Chamfer Distance
s1_s2 = tf.reduce_sum(tf.reduce_min(pairwise_diff(reconstructed_points, original_points), axis=2), axis=1)
s2_s1 = tf.reduce_sum(tf.reduce_min(pairwise_diff(original_points, reconstructed_points), axis=2), axis=1)
loss = (s1_s2 + s2_s1)/ n_points
elif type == 'emd':
matchl_out, matchr_out = tf_auctionmatch.auction_match(original_points, reconstructed_points)
matched_out = tf_sampling.gather_point(reconstructed_points, matchl_out)
d = tf.sqrt(tf.reduce_sum(tf.square(original_points - matched_out), axis=2))
loss = tf.reduce_sum(d,axis=1)/ n_points
elif type=='joint':
#Use both loss functions
s1_s2 = tf.reduce_sum(tf.reduce_min(pairwise_diff(reconstructed_points, original_points), axis=2), axis=1)
s2_s1 = tf.reduce_sum(tf.reduce_min(pairwise_diff(original_points, reconstructed_points), axis=2), axis=1)
loss_chamfer = s1_s2 + s2_s1
matchl_out, matchr_out = tf_auctionmatch.auction_match(original_points, reconstructed_points)
matched_out = tf_sampling.gather_point(reconstructed_points, matchl_out)
d = tf.sqrt(tf.reduce_sum(tf.square(original_points - matched_out), axis=2))
loss_emd = tf.reduce_sum(d,axis=1)
loss = (loss_chamfer + loss_emd) / n_points
loss = tf.reduce_mean(loss)
tf.summary.scalar('emd_loss', loss)
if add_fv_loss:
fv_rec = tf_util.get_fv_tf_no_mvn(reconstructed_points, w, mu, sigma, flatten=False)
fv_loss = tf.nn.l2_loss(fv_rec-fv)
tf.summary.scalar('fv_loss', fv_loss)
loss = loss + 0.001*fv_loss
# tf.summary.scalar('weight_decay_loss', weight_decay_loss)
return loss, d, matched_out
def get_emd_loss(pred, pc, radius):
""" pred: BxNxC,
label: BxN, """
batch_size = pred.get_shape()[0].value
matchl_out, matchr_out = tf_auctionmatch.auction_match(pred, pc)
matched_out = tf_sampling.gather_point(pc, matchl_out)
dist = tf.reshape((pred - matched_out)**2, shape=(batch_size, -1))
dist = tf.reduce_mean(dist, axis=1, keep_dims=True)
dist_norm = dist / radius
emd_loss = tf.reduce_mean(dist_norm)
return emd_loss
def get_emd_dist(gt_pcl, pred_pcl):
'''
Calculate emd between two point clouds
Args:
gt_pcl: (BS,N_pts,3); GT point cloud
pred_pcl: (BS,N_pts,3); predicted point cloud
Returns:
emd: (BS); averaged emd
'''
batch_size, num_points = (gt_pcl.shape)[:2]
X,_ = tf.meshgrid(range(batch_size), range(num_points), indexing='ij')
ind, _ = auction_match(pred_pcl, gt_pcl) # Ind corresponds to points in pcl_gt
ind = tf.stack((X, ind), -1)
emd = tf.reduce_mean(tf.reduce_sum((tf.gather_nd(gt_pcl, ind) - pred_pcl)**2, axis=-1), axis=1) # (BATCH_SIZE,NUM_POINTS,3) --> (BATCH_SIZE,NUM_POINTS) --> (BATCH_SIZE)
return emd
def get_rec_metrics(gt_pcl, pred_pcl, batch_size=10, num_points=1024):
'''
Calculate chamfer and emd metrics
args:
gt_pcl: float, (BS,N_PTS,3); ground truth point cloud
pred_pcl: float, (BS,N_PTS,3); predicted point cloud
'''
dists_forward, _, dists_backward, _ = tf_nndistance.nn_distance(gt_pcl,
pred_pcl)
dists_forward = tf.reduce_mean(tf.sqrt(dists_forward), axis=1) # (B, )
dists_backward = tf.reduce_mean(tf.sqrt(dists_backward), axis=1) # (B, )
chamfer_distance = dists_backward + dists_forward
X,_ = tf.meshgrid(range(batch_size), range(num_points), indexing='ij')
ind, _ = auction_match(pred_pcl, gt_pcl) # Ind corresponds to points in pcl_gt
print X.get_shape()
print ind.get_shape()
ind = tf.stack((X, ind), -1)
print gt_pcl.get_shape()
print ind.get_shape()
print
emd = tf.reduce_mean(tf.sqrt(tf.reduce_sum((tf.gather_nd(gt_pcl, ind) - \
pred_pcl)**2, axis=-1)), axis=1) # (BATCH_SIZE,NUM_POINTS,3) --> (BATCH_SIZE,NUM_POINTS) --> (BATCH_SIZE)
return dists_forward, dists_backward, chamfer_distance, emd