python - 使用 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)