draw a line in tensorflow

Question

I want to create a human pose skeleton estimation network and for this, I have a two-part network, first part generates 16 heatmaps as output(each heatmap for different joint and hence a key point can be extracted), using these 16 key points I wish to create a human skeleton and feed it to second half of my network. My problem is, how do I draw lines between the key points to create the skeleton? I couldn't find a way to do it on a tensor object using tensorflow or keras.

Answer 1

I know i'm a bit late but here is some code that I think does what you're after (in TFv2.3). Hopefully it will save someone time in the future!

It uses solely tensorflow ops, so you can use it in data loaders etc. The real pain here is that Tensorflow doesn't allow Eager Assignment, so you can't just update tensors by index. This works around that by creating two sparse tensors, one for the mask (where to apply the line) and another for the new_values (what value to apply at the line). The code for simply designing the line might not be applicable in your case (based on https://stackoverflow.com/a/47381058) but ported away from 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])

Just for a sanity check, you can call it with the following, and display it using 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)