pre_process_poses.py 文件源码

def sim_occlusions(poses, dm_shape, batch_size, max_length, n_dims, body_splits, _int_type=tf.int32, _float_type=tf.float32):
    def occluded_poses():
        body_splits_tf = tf.constant(body_splits, dtype=_int_type)
        occ_idcs = tf.random_uniform([batch_size, 1], minval=0, maxval=len(body_splits), dtype=_int_type)
        occ_idcs = tf.gather_nd(body_splits_tf, occ_idcs)
        noise_mask = tf.tile(
            tf.reshape(
                tf.cast(tf.reduce_sum(tf.one_hot(occ_idcs, dm_shape[0]), axis=1), dtype=tf.bool),
                [batch_size, 1, dm_shape[0], 1]),
            [1, max_length, 1, n_dims])
        noisy_poses = poses * tf.random_uniform([batch_size, max_length, 1, n_dims], minval=0.8, maxval=1.2, dtype=_float_type)
        return tf.where(noise_mask, noisy_poses, poses)

    occlude_rate = 0.5
    return tf.cond(tf.cast(tf.round(tf.random_uniform([], minval=-0.5, maxval=0.5) + occlude_rate), tf.bool),
                   occluded_poses, lambda: poses)