首页 > 解决方案 > TensorFlow 地标热图

问题描述

我正在尝试使用 tensorflow 绘制具有里程碑意义的热图。
我目前的方法是tf.scatter_nd这样使用的:

    def draw_lmarks(x):
        def draw_lmarks_inner(x2):
            return tf.scatter_nd(x2[0], x2[1], shape=(IMGSIZE, IMGSIZE))
        ret = tf.map_fn(draw_lmarks_inner, x, dtype="float32")
        return tf.reshape(tf.reduce_max(ret, axis=0), [IMGSIZE, IMGSIZE, 1])
    return tf.map_fn(draw_lmarks, [locations, vals], dtype="float32")

但这很慢,因为我必须为每个批次时间地标创建一个 IMAGESIZE*IMAGESIZE 图像。
所以我四处寻找tf.tensor_scatter_nd_update,发现我可以使用如下:

    img = tf.zeros((IMGSIZE,IMGSIZE), dtype="float32")
    def draw_lmarks(x):
        return tf.tensor_scatter_nd_update(img, x[0], x[1])
    imgs = tf.map_fn(draw_lmarks, [locations, vals], dtype="float32")

这允许我只生成运行速度相当快的 batch_size 图像。
...但是,这不会在某一时刻使用最高值,而是简单地覆盖。
tf.scatter_max一个听起来像我需要的功能,但这似乎期望不同形状的输入。有没有办法使用第二种方法,而不是覆盖值在某一点取最大值?

形状:
location= (-1, 68, 16, 16, 2)
vals= (-1, 68, 16, 16)

可视化:
这是第二个(更快)函数返回的内容: 虽然我需要类似的东西
在此处输入图像描述

在此处输入图像描述

标签: tensorflow

解决方案


我认为首先设置地标的种子然后将结果与热图模板进行卷积会更好。就像是

import tensorflow as tf

num_loc = 10
im_dim = 32
locations = tf.random.uniform((num_loc, 2), maxval=im_dim, dtype=tf.int32)
centers = tf.scatter_nd(locations, [1]*num_loc, (im_dim, im_dim))
heatmap  = tf.nn.conv2d(centers[None, :, :, None], heatmap_template[:, :, None, None], (1, 1, 1, 1), 'SAME')[0, :, :, 0]

推荐阅读