批量 4D 张量 Tensorflow 索引

Question

Given

• batch_images : 4D tensor of shape (B, H, W, C)
• x : 3D tensor of shape (B, H, W)
• y : 3D tensor of shape (B, H, W)

Goal

How can I index into batch_images using the x and y coordinates to obtain a 4D tensor of shape B, H, W, C . That is, I want to obtain for each batch, and for each pair (x, y) a tensor of shape C .

In numpy, this would be achieved using input_img[np.arange(B)[:,None,None], y, x] for example but I can't seem to make it work in tensorflow.

My attempt so far

def get_pixel_value(img, x, y):
    """
    Utility function to get pixel value for 
    coordinate vectors x and y from a  4D tensor image.
    """
    H = tf.shape(img)[1]
    W = tf.shape(img)[2]
    C = tf.shape(img)[3]

    # flatten image
    img_flat = tf.reshape(img, [-1, C])

    # flatten idx
    idx_flat = (x*W) + y

    return tf.gather(img_flat, idx_flat)

which is returning an incorrect tensor of shape (B, H, W) .

Answer 1

It should be possible to do it by flattening the tensor as you've done, but the batch dimension has to be taken into account in the index calculation. In order to do this, you'll have to make an additional dummy batch index tensor with the same shape as x and y that always contains the index of the current batch. This is basically the np.arange(B) from your numpy example, which is missing from your TensorFlow code.

You can also simplify things a bit by using tf.gather_nd , which does the index calculations for you.

Here's an example:

import numpy as np
import tensorflow as tf

# Example tensors
M = np.random.uniform(size=(3, 4, 5, 6))
x = np.random.randint(0, 5, size=(3, 4, 5))
y = np.random.randint(0, 4, size=(3, 4, 5))

def get_pixel_value(img, x, y):
    """
    Utility function that composes a new image, with pixels taken
    from the coordinates given in x and y.
    The shapes of x and y have to match.
    The batch order is preserved.
    """

    # We assume that x and y have the same shape.
    shape = tf.shape(x)
    batch_size = shape[0]
    height = shape[1]
    width = shape[2]

    # Create a tensor that indexes into the same batch.
    # This is needed for gather_nd to work.
    batch_idx = tf.range(0, batch_size)
    batch_idx = tf.reshape(batch_idx, (batch_size, 1, 1))
    b = tf.tile(batch_idx, (1, height, width))

    indices = tf.pack([b, y, x], 3)
    return tf.gather_nd(img, indices)

s = tf.Session()
print(s.run(get_pixel_value(M, x, y)).shape)
# Should print (3, 4, 5, 6).
# We've composed a new image of the same size from randomly picked x and y
# coordinates of each original image.