def compute_parameter_loss(parameters_pred, parameters_gt,
matching_indices, angle_diff):
axis = batched_gather(parameters_pred['cone_axis'],
matching_indices,
axis=1)
dot_abs = tf.abs(
tf.reduce_sum(axis * parameters_gt['cone_axis'], axis=2))
if angle_diff:
return acos_safe(dot_abs) # BxK
else:
return 1.0 - dot_abs # BxK
def compute_miou_loss(W, I_gt, matching_indices):
# W - BxNxK
# I_gt - BxN
W_reordered = batched_gather(W, indices=matching_indices, axis=2) # BxNxK
depth = tf.shape(W)[2]
# notice in tf.one_hot, -1 will result in a zero row, which is what we want
W_gt = tf.one_hot(I_gt, depth=depth, dtype=tf.float32) # BxNxK
dot = tf.reduce_sum(W_gt * W_reordered, axis=1) # BxK
denominator = tf.reduce_sum(W_gt, axis=1) + tf.reduce_sum(W_reordered,
axis=1) - dot
mIoU = dot / (denominator + DIVISION_EPS) # BxK
return 1.0 - mIoU
def calculate_eval_stats(W, matching_indices, mask_gt, P_in, type_per_point,
T_gt, parameters, residue_losses, parameter_loss,
residue_per_point_array):
'''
Returns a dict containing:
- stats : {
per_instance_type_accuracy: B, average primitive type accuracy for a shape
avg_residue_loss_without_gt: B, average residue loss using the predicted type
parameter_loss_without_gt: B, average parameter loss using the predicted type (over only primitives with matched type)
}
- null_mask: BxK, indicating which predicted primitives are null
- mask_gt_nulled: BxK, indicated which ground truth primitive is not null and is matched with a predicted (non-null) primitive
- instance_per_point: BxN, non-one-hot version of W
- type_per_intance: BxK, type for predicted primitives
- residue_gt_primitive: BxKxN', distance from sampled points on ground truth S_k to the predicted primitive matched with S_k
- residue_to_closest: BxN, distance from each input point to the closest predicted primitive
'''
batch_size = tf.shape(W)[0]
n_points = tf.shape(W)[1]
n_max_instances = W.get_shape()[2] # n_max_instances should not be dynamic
n_registered_primitives = fitter_factory.get_n_registered_primitives()
null_mask = tf.reduce_sum(W, axis=1) < 0.5 # BxK
# null_mask indicates which predicted primitives are null
I = tf.where(
tf.reduce_sum(W, axis=2) > 0.5,
tf.argmax(W, axis=2, output_type=tf.int32),
tf.fill([batch_size, n_points], -1)) # BxN
# I can have -1 entries, indicating unassigned points, just like I_gt, and tf.one_hot(I) == W
per_point_type_one_hot = tf.one_hot(type_per_point,
depth=n_registered_primitives,
dtype=tf.float32) # BxNxT
instance_type_prob = tf.reduce_sum(
tf.expand_dims(W, axis=3) *
tf.expand_dims(per_point_type_one_hot, axis=2),
axis=1) # BxKxT
instance_type = tf.argmax(instance_type_prob, axis=2,
output_type=tf.int32) # BxK
null_mask_gt = batched_gather(
null_mask, matching_indices,
axis=1) # BxK, indicating which gt primitive is not matched
mask_gt_nulled = tf.logical_and(mask_gt, tf.logical_not(
null_mask_gt)) # only count these gt primitives towards some metrics
residue_loss_without_gt = aggregate_loss_from_stacked(
residue_losses, batched_gather(instance_type, matching_indices,
axis=1)) # BxK
avg_residue_loss_without_gt = reduce_mean_masked_instance(
residue_loss_without_gt, mask_gt_nulled) # B
# for parameter loss w/o gt, only count when the predicted type matches the ground truth type
instance_matched_mask = tf.equal(T_gt,
batched_gather(instance_type,
matching_indices,
axis=1)) # BxK, boolean
per_instance_type_accuracy = reduce_mean_masked_instance(
tf.to_float(instance_matched_mask), mask_gt_nulled) # B
parameter_loss_without_gt = reduce_mean_masked_instance(
parameter_loss, tf.logical_and(instance_matched_mask, mask_gt_nulled))
result = {
'stats': {
'per_instance_type_accuracy': per_instance_type_accuracy, # B
'avg_residue_loss_without_gt': avg_residue_loss_without_gt, # B
'parameter_loss_without_gt': parameter_loss_without_gt, # B
},
}
residue_matrix = []
identity_matching_indices = tf.tile(
tf.expand_dims(tf.range(n_max_instances), axis=0),
[batch_size, 1]) # BxK
for fitter_cls in fitter_factory.get_all_fitter_classes():
residue_per_point = fitter_cls.compute_residue_loss(
parameters,
tf.tile(tf.expand_dims(P_in, axis=1), [1, n_max_instances, 1, 1]),
matching_indices=identity_matching_indices) # BxKxN
residue_matrix.append(residue_per_point)
residue_matrix = tf.stack(residue_matrix,
axis=3) # BxKxNxT, this matrix might be large!
# residue_to_primitive[b,n,k] = residue_matrix[b, k, n, instance_type[b,k]]
indices_0 = tf.tile(
tf.expand_dims(tf.expand_dims(tf.range(batch_size), axis=1), axis=2),
[1, n_points, n_max_instances]) # BxNxK
indices_1 = tf.tile(
tf.expand_dims(tf.expand_dims(tf.range(n_max_instances), axis=0),
axis=1), [batch_size, n_points, 1]) # BxNxK
indices_2 = tf.tile(
tf.expand_dims(tf.expand_dims(tf.range(n_points), axis=0), axis=2),
[batch_size, 1, n_max_instances]) # BxNxK
indices_3 = tf.tile(tf.expand_dims(instance_type, axis=1),
[1, n_points, 1]) # BxNxK
residue_to_primitive = tf.gather_nd(
residue_matrix,
indices=tf.stack([indices_0, indices_1, indices_2, indices_3],
axis=3)) # BxNxK
# set null primitive residues to a large number
null_mask_W_like = tf.tile(tf.expand_dims(null_mask, axis=1),
[1, n_points, 1]) # BxNxK
residue_to_primitive = tf.where(null_mask_W_like,
1e8 * tf.ones_like(residue_to_primitive),
residue_to_primitive) # BxNxK
residue_to_closest = tf.reduce_min(residue_to_primitive, axis=2) # BxN
residue_gt_primitive = aggregate_per_point_loss_from_stacked(
residue_per_point_array,
batched_gather(instance_type, matching_indices,
axis=1)) # BxKxN', squared distance
# Save information for downstream analysis
result['null_mask'] = null_mask # BxK
result['mask_gt_nulled'] = mask_gt_nulled # BxK
result['instance_per_point'] = I # BxN
result['type_per_instance'] = instance_type # BxK
result['residue_gt_primitive'] = tf.sqrt(residue_gt_primitive) # BxKxN'
result['residue_to_closest'] = tf.sqrt(residue_to_closest) # BxN
return result
def get_direct_regression_model(scope, P, n_max_instances, gt_dict,
is_training, bn_decay):
'''
Inputs:
- P: BxNx3 tensor, the input point cloud
- K := n_max_instances
- gt_dict: ground truth dictionary, needed since we are also computing the loss
Outputs: (pred_dict, total_loss), where pred_dict contains
- W: BxNxK, segmentation instances, binary
- normal_per_point: BxNx3, normal per point, except in DPPN we don't predict normal, so all normals are constant
- type_per_point: BxNxT, type per points, binary
- parameters - a dict, each entry is a BxKx... tensor
'''
n_registered_primitives = fitter_factory.get_n_registered_primitives()
batch_size = tf.shape(P)[0]
n_points = tf.shape(P)[1]
param_pair_list = get_param_dims_pair_list(n_max_instances)
flattened_param_dims = [pr[1][0] * pr[1][1] for pr in param_pair_list]
reg_result = build_pointnet2_cls('direct_reg_net',
point_cloud=P,
out_dims=flattened_param_dims,
is_training=is_training,
bn_decay=bn_decay)
parameters = {}
for idx, cls_result in enumerate(reg_result):
param_name, param_dim = param_pair_list[idx]
if param_dim[1] == 1:
parameters[param_name] = cls_result # BxK
else:
parameters[param_name] = tf.reshape(cls_result, [-1, *param_dim])
# normalize quantities
for fitter_cls in fitter_factory.get_all_fitter_classes():
fitter_cls.normalize_parameters(parameters)
residue_losses = []
for fitter_cls in fitter_factory.get_all_fitter_classes():
residue_per_point = fitter_cls.compute_residue_loss_pairwise(
parameters, gt_dict['points_per_instance']) # BxKxKxN'
residue_avg = tf.reduce_mean(residue_per_point, axis=3) # BxKxK
# residue_avg[b, k1, k2] is roughly the distance between gt instance k1 and predicted instance k2
residue_losses.append(residue_avg)
residue_matrix = tf.stack(residue_losses, axis=3) # BxKxKxT
residue_matrix_flattened = tf.reshape(
residue_matrix, shape=[batch_size, n_max_instances, -1]) # BxKxKT
n_instance_gt = tf.reduce_max(gt_dict['instance_per_point'], axis=1) + 1
mask_gt = tf.sequence_mask(n_instance_gt, maxlen=n_max_instances)
matching_indices = tf.stop_gradient(
tf.py_func(hungarian_matching,
[residue_matrix_flattened, n_instance_gt],
Tout=tf.int32)) # BxK
matching_matrix = tf.reshape(
tf.one_hot(matching_indices,
depth=n_max_instances * n_registered_primitives), [
batch_size, n_max_instances, n_max_instances,
n_registered_primitives
]) # BxKxKxT
# only 1 element in matching_matrix[b, k, ..., ...] is nonzero
direct_loss = tf.reduce_sum(matching_matrix * residue_matrix,
axis=[2, 3]) # BxK
direct_loss = tf.reduce_sum(direct_loss, axis=1) / tf.to_float(
n_instance_gt) # B
# reorder parameters
matching_instance_id = tf.argmax(tf.reduce_sum(matching_matrix, axis=3),
axis=2,
output_type=tf.int32) # BxK
matching_instance_type = tf.argmax(tf.reduce_sum(matching_matrix, axis=2),
axis=2,
output_type=tf.int32) # BxK
for param in parameters:
parameters[param] = batched_gather(parameters[param],
matching_instance_id,
axis=1)
# now kth instance has type matching_instance_type[b, k] with parameters[b, k,...]
# next construct W: BxNxK
residue_per_point_matrix = []
identity_matching_indices = tf.tile(
tf.expand_dims(tf.range(n_max_instances), axis=0),
[batch_size, 1]) # BxK
for fitter_cls in fitter_factory.get_all_fitter_classes():
residue_per_point = fitter_cls.compute_residue_loss(
parameters,
tf.tile(tf.expand_dims(P, axis=1), [1, n_max_instances, 1, 1]),
matching_indices=identity_matching_indices) # BxKxN
residue_per_point_matrix.append(residue_per_point)
residue_per_point_matrix = tf.stack(residue_per_point_matrix,
axis=3) # BxKxNxT
# dist(P[b,n], instance k) = residue_per_point_matrix[b, k, n, matching_instance_type[b, k]]
indices_0 = tf.tile(
tf.expand_dims(tf.expand_dims(tf.range(batch_size), axis=1), axis=1),
[1, n_points, n_max_instances]) # BxNxK
indices_1 = tf.tile(
tf.expand_dims(tf.expand_dims(tf.range(n_max_instances), axis=0),
axis=0), [batch_size, n_points, 1]) # BxNxK
indices_2 = tf.tile(
tf.expand_dims(tf.expand_dims(tf.range(n_points), axis=0), axis=2),
[batch_size, 1, n_max_instances]) # BxNxK
indices_3 = tf.tile(tf.expand_dims(matching_instance_type, axis=1),
[1, n_points, 1]) # BxNxK
P_to_instance_dist = tf.gather_nd(
residue_per_point_matrix,
indices=tf.stack([indices_0, indices_1, indices_2, indices_3],
axis=3)) # BxNxK
instance_per_point = tf.argmin(P_to_instance_dist,
axis=2,
output_type=tf.int32) # BxN
W = tf.one_hot(instance_per_point, depth=n_max_instances) # BxNxK
type_per_point = batched_gather(matching_instance_type,
instance_per_point,
axis=1)
# we do not predict normal
normal_per_point = tf.tile(
tf.expand_dims(tf.expand_dims(tf.constant([1, 0, 0], dtype=tf.float32),
axis=0),
axis=0), [batch_size, n_points, 1]) # BxNx3
return {
'W': W,
'normal_per_point': normal_per_point,
'type_per_point': tf.one_hot(type_per_point,
depth=n_registered_primitives),
'parameters': parameters,
}, direct_loss
def adaptor_matching(param_list, matching_indices):
return [tf.expand_dims(batched_gather(param, matching_indices, axis=1), axis=2) for param in param_list]