在tensorflow中画线

Question

我想创建一个人体姿态骨架估计网络，为此，我有一个两部分的网络，第一部分生成16个热图作为输出（每个热图针对不同的关节，因此可以提取一个关键点），使用这16个关键点，我希望创建一个人体骨架，并将其馈送到我的网络的第二部分。我的问题是，我如何在关键点之间画线来创建骨架？我找不到使用tensorflow或Keras在张量对象上做这件事的方法。

Answer 1

我知道我有点晚了，但这里有一些代码，我认为可以满足您的需要（在TFV2.3中）。希望它能为将来的人节省时间！

它只使用tensorflow ops，所以您可以在数据加载器等中使用它。这里真正的麻烦是，tensorflow不允许急切赋值，所以您不能仅仅通过索引更新张量。通过创建两个稀疏张量，一个用于mask（在何处应用该行），另一个用于new_values（在该行应用什么值），可以解决这一问题。用于简单设计行的代码可能不适用于您的情况（基于https://stackoverflow.com/a/47381058)，而是从Numpy移植。

import tensorflow as tf


def trapez(y, y0, w):
    return tf.clip_by_value(tf.minimum(y + 1 + w/2 - y0, -y + 1 + w/2 + y0), 0, 1)


def apply_output(img, yy, xx, val):
    stack = tf.stack([yy, xx], axis=1)
    stack = tf.cast(stack, tf.int64)
    values = tf.ones(stack.shape[0], tf.float32)
    mask = tf.sparse.SparseTensor(indices=stack, values=values, dense_shape=img.shape)
    mask = tf.sparse.reorder(mask)
    mask = tf.sparse.to_dense(mask)
    mask = tf.cast(mask, tf.float32)

    new_values = tf.sparse.SparseTensor(indices=stack, values=val, dense_shape=img.shape)
    new_values = tf.sparse.reorder(new_values)
    new_values = tf.sparse.to_dense(new_values)

    img = img * (1 - mask) + new_values * mask
    img = tf.cast(tf.expand_dims(img * 255, axis=-1), tf.uint8)
    return img


def weighted_line(img, r0, c0, r1, c1, w):
    output = img
    x = tf.range(c0, c1 + 1, dtype=tf.float32)
    slope = (r1-r0) / (c1-c0)
    w *= tf.sqrt(1 + tf.abs(slope)) / 2
    y = x * slope + (c1*r0-c0*r1) / (c1-c0)

    thickness = tf.math.ceil(w/2)
    yy = (tf.reshape(tf.math.floor(y), [-1, 1]) + tf.reshape(tf.range(-thickness-1, thickness+2), [1, -1]))
    xx = tf.repeat(x, yy.shape[1])
    values = tf.reshape(trapez(yy, tf.reshape(y, [-1, 1]), w), [-1])
    yy = tf.reshape(yy, [-1])

    limits_y = tf.math.logical_and(yy >= 0, yy < img.shape[0])
    limits_x = tf.math.logical_and(xx >= 0, xx < img.shape[1])
    limits = tf.math.logical_and(limits_y, limits_x)
    limits = tf.math.logical_and(limits, values > 0)
    yy = tf.cast(yy[limits], tf.float32)
    xx = tf.cast(xx[limits], tf.float32)

    return yy, xx, values[limits], apply_output(output, yy, xx, values[limits])

只是为了检查是否正常，您可以使用以下命令调用它，并使用opencv显示它

if __name__ == "__main__":
    IMG = tf.zeros((500, 500), tf.float32)
    yy, xx, vals, FINAL_IMG = weighted_line(IMG, 10, 20, 100, 200, 5)
    jpeg_string = tf.io.encode_jpeg(FINAL_IMG)
    tf.io.write_file("output.jpg", jpeg_string)
    import cv2
    img = cv2.imread("output.jpg")
    cv2.imshow("Output", img)
    cv2.waitKey(0)