Tensorflow地标热图

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

tf.scatter\n和
如下：

    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\u scatter\u nd\u 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")

这使得我只能生成运行速度相当快的批量大小的图像。

... 但是，这不会在某一点使用最高值，而是简单地覆盖。

有一个
tf.scatter\u max
函数，听起来像是我需要的，但这似乎需要不同的形状输入。
是否有一种方法可以使用第二种方法，但不是覆盖值，而是在某一点上取最大值

形状：

位置
=（-1,68,16,16,2）

vals
=（-1,68,16,16）

可视化：

这是第二个（更快的）函数返回的结果：



而我需要像

我认为，首先设定地标的种子，然后将结果与热图模板进行卷积，会让你的境况更好。差不多

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]

谢谢，我还没有试过，但实际上这听起来是更好的方法。既然赏金快用完了，我确实发出去了，但在我调查了它的实际效果之后，我会接受的。