Batched 4D Tensor Tensorflow Indizierung

Question

Bei

• batch_images: 4D Tensor Form (B, H, W, C)
• x: 3D-Tensor Form (B, H, W)
• y: 3D-Tensor Form (B, H, W)

Ziel

Wie kann ich Index in batch_images mit der x und y koordiniert ein 4D-Tensor Form B, H, W, C zu erhalten. Das heißt, ich möchte für jede Charge und für jedes Paar (x, y) einen Tensor der Form C erhalten.

In numpy, würde dies zum Beispiel mit input_img[np.arange(B)[:,None,None], y, x] erreicht werden, aber ich kann nicht scheinen es in Tensorflow arbeiten.

Mein Versuch so weit

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) 

die (B, H, W) einen falschen Tensor Form zurückkehrt.

Answer 1

Es sollte möglich sein, dies zu tun, indem Sie den Tensor abflachen, wie Sie es getan haben, aber die Batch-Dimension muss in der Indexberechnung berücksichtigt werden. Dazu müssen Sie einen zusätzlichen Dummy-Batch-Index-Tensor mit der gleichen Form wie x und y erstellen, der immer den Index des aktuellen Stapels enthält. Dies ist im Grunde die np.arange(B) von Ihrem numpy Beispiel, das in Ihrem TensorFlow-Code fehlt.

Sie können die Dinge auch ein wenig vereinfachen, indem Sie tf.gather_nd verwenden, die die Indexberechnungen für Sie durchführt.

Hier ist ein Beispiel:

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.