flownet.py 文件源码

def flows_to_img(flows):
    """Pyfunc wrapper to transorm flow vectors in color coding"""

    def _flow_transform(flows):
        """ Tensorflow Pyfunc to transorm flow to color coding"""

        flow_imgs = []
        for flow in flows:
            img = computeColor.computeImg(flow)
            # cv2 returns bgr images
            b, g, r = cv2.split(img)
            img = cv2.merge((r, g, b))
            flow_imgs.append(img)
        return [flow_imgs]

    flow_imgs = tf.py_func(_flow_transform, [flows],
                           [tf.uint8], stateful=False, name='flow_transform')

    flow_imgs = tf.squeeze(tf.stack(flow_imgs))
    flow_imgs.set_shape([FLAGS.batchsize] + FLAGS.d_shape_img)
    return flow_imgs