Access 3D tensor (image) using a 2d tensor of indices

user_7034613 0
Viewed 96

With the following 3D tensor representing an image
img.shape=[H,W,F]
And a tensor representing the indices to that img
indices.shape=[N,2]
E.g. if
indices = [[0,1],[5,3],...]] I would like to create a new tensor of shape new.shape=[N,F] where
new[k] == img[indices[k][0],indices[k][1]] Currently to solve this I flatten both tensors:

    idx_flattened = idx_flattened [:,0] * (idx_flattened [:,1].max()+1) + idx_flattened[:,1]
    img = img .reshape(-1,F)
    new = img[idx_flattened ]

But I'm certain there is a better way:)

Here's a full minimal example:

img = torch.arange(8*10*3).reshape(8,10,3)
indices = torch.tensor([[0,0],[3,0],[1,2]])
new = img[indices] <- This does not work
new = [[  0,   1,   2],[ 90,  91,  92],[ 36,  37,  38]]

Ideas?

pytorch python
1 Answers

Slicing would work

img[indices[:,0], indices[:,1]]
tensor([[ 0,  1,  2],
        [90, 91, 92],
        [36, 37, 38]])
user_60756990
Related