def align_2d_skeletons(coords_true, coords_pred, joint_validity_mask):
mean_pred, stdev_pred = tfu.mean_stdev_masked(
coords_pred, joint_validity_mask, items_axis=1, dimensions_axis=2)
mean_true, stdev_true = tfu.mean_stdev_masked(
coords_true, joint_validity_mask, items_axis=1, dimensions_axis=2)
coords_pred_result = tf.div_no_nan(coords_pred - mean_pred, stdev_pred) * stdev_true
coords_true_result = coords_true - mean_true
return coords_true_result, coords_pred_result
def align_2d_skeletons(coords_pred, coords_true, joint_validity_mask):
mean_pred, stdev_pred = tfu.mean_stdev_masked(coords_pred,
joint_validity_mask,
items_axis=1,
dimensions_axis=2)
mean_true, stdev_true = tfu.mean_stdev_masked(coords_true,
joint_validity_mask,
items_axis=1,
dimensions_axis=2)
return tf.math.divide_no_nan(coords_pred - mean_pred,
stdev_pred) * stdev_true + mean_true
def reconstruct_ref_weak(normalized_2d, coords3d_rel, validity_mask):
mean3d, stdev3d = tfu.mean_stdev_masked(coords3d_rel[..., :2],
validity_mask,
items_axis=1,
dimensions_axis=2)
mean2d, stdev2d = tfu.mean_stdev_masked(normalized_2d[..., :2],
validity_mask,
items_axis=1,
dimensions_axis=2)
stdev2d = tf.maximum(stdev2d, 1e-5)
stdev3d = tf.maximum(stdev3d, 1e-5)
old_mean = tfu.reduce_mean_masked(coords3d_rel,
validity_mask,
axis=1,
keepdims=True)
new_mean_z = tf.math.divide_no_nan(stdev3d, stdev2d)
new_mean = model.util.to_homogeneous(mean2d) * new_mean_z
return tf.squeeze(new_mean - old_mean, 1)
def compute_losses(self, inps, preds):
losses = AttrDict()
if FLAGS.scale_agnostic_loss:
mean_true, scale_true = tfu.mean_stdev_masked(
inps.coords3d_true,
inps.joint_validity_mask,
items_axis=1,
dimensions_axis=2)
mean_pred, scale_pred = tfu.mean_stdev_masked(
preds.coords3d_rel_pred,
inps.joint_validity_mask,
items_axis=1,
dimensions_axis=2)
coords3d_pred_rootrel = tf.math.divide_no_nan(
preds.coords3d_rel_pred - mean_pred, scale_pred) * scale_true
coords3d_true_rootrel = inps.coords3d_true - mean_true
else:
coords3d_true_rootrel = tfu3d.center_relative_pose(
inps.coords3d_true, inps.joint_validity_mask,
FLAGS.mean_relative)
coords3d_pred_rootrel = tfu3d.center_relative_pose(
preds.coords3d_rel_pred, inps.joint_validity_mask,
FLAGS.mean_relative)
rootrel_absdiff = tf.abs(
(coords3d_true_rootrel - coords3d_pred_rootrel) / 1000)
losses.loss3d = tfu.reduce_mean_masked(rootrel_absdiff,
inps.joint_validity_mask)
if FLAGS.scale_agnostic_loss:
_, scale_true = tfu.mean_stdev_masked(inps.coords3d_true,
inps.joint_validity_mask,
items_axis=1,
dimensions_axis=2,
fixed_ref=tf.zeros_like(
inps.coords3d_true))
_, scale_pred = tfu.mean_stdev_masked(preds.coords3d_pred_abs,
inps.joint_validity_mask,
items_axis=1,
dimensions_axis=2,
fixed_ref=tf.zeros_like(
inps.coords3d_true))
preds.coords3d_pred_abs = tf.math.divide_no_nan(
preds.coords3d_pred_abs, scale_pred) * scale_true
if self.global_step > 5000:
absdiff = tf.abs(
(inps.coords3d_true - preds.coords3d_pred_abs) / 1000)
losses.loss3d_abs = tfu.reduce_mean_masked(
absdiff, inps.joint_validity_mask)
else:
losses.loss3d_abs = tf.constant(0, tf.float32)
scale_2d = 1 / FLAGS.proc_side * FLAGS.box_size_mm / 1000
losses.loss23d = tfu.reduce_mean_masked(
tf.abs((inps.coords2d_true - preds.coords2d_pred) * scale_2d),
inps.joint_validity_mask)
preds.coords32d_pred_2d = models.util.align_2d_skeletons(
preds.coords32d_pred_2d, inps.coords2d_true_2d,
inps.joint_validity_mask_2d)
losses.loss32d = tfu.reduce_mean_masked(
tf.abs(
(inps.coords2d_true_2d - preds.coords32d_pred_2d) * scale_2d),
inps.joint_validity_mask_2d)
losses.loss22d = tfu.reduce_mean_masked(
tf.abs(
(inps.coords2d_true_2d - preds.coords22d_pred_2d) * scale_2d),
inps.joint_validity_mask_2d)
losses3d = [
losses.loss3d, losses.loss23d,
FLAGS.absloss_factor * losses.loss3d_abs
]
losses2d = [losses.loss22d, losses.loss32d]
losses.loss = tf.add_n(
losses3d) + FLAGS.loss2d_factor * tf.add_n(losses2d)
return losses