使用TensorFlow在图像上绘制点

Question

给出：

• 以形状为[batch_size, num_points, 2]的张量排列的一组点，
• 以及一批形状为[batch_size, height, width, 3]的图像，

我们如何在图像上绘制点？

tf.image.draw_bounding_boxes可以在图像像素的顶部绘制边界框，但我没有找到绘制关键点的等效项。

不需要使用纯Python和tf.py_func的tensorflow-nic答案将是首选。

Answer 1

我使用零尺寸的框画点：

def draw_points(images, points, tf_summary_writer, step):
    """
    :param images: [batch_size, height, width, 3], tf.float32, from 0 to 1
    :param points: [batch_size, num_points, 2], tf.float32, from 0 to 1
    :param tf_summary_writer: SummaryWriter
    :param step: int
    """
    with tf_summary_writer.as_default():
        x_min = points[:, :, 0]  # shape [bs, np]
        y_min = points[:, :, 1]  # shape [bs, np]
        x_max = points[:, :, 0]  # shape [bs, np]
        y_max = points[:, :, 1]  # shape [bs, np]
        boxes = tf.stack([y_min, x_min, y_max, x_max], axis=-1)  # shape [bs, np, 4]
        image_with_points = tf.image.draw_bounding_boxes(images, boxes, colors=[[1, 0, 0, 1]])  # RGBA red color
        tf.summary.image('points', image_with_points, step)