Mapping tensor in pytorch

user_1448257 0
Viewed 392

I have the following two tensors:

  • img is a RGB image of shape (224,224,3)
  • uvs is a tensor with same spacial size e.g. (224, 224, 2) that maps to coordinates (x,y). In other words it provides (x,y) coordinates for every pixel of the input image.

I want to create now a new output image tensor that contains on index (x,y) the value of the input image. So the output should be an image as well with the pixels rearranged according to the mapping tensor.

Small toy example:

img = [[c1,c2], [c3, c4]] where c is a RGB color [r, g, b]
uvs = [[[0,0], [1,1]],[[0,1], [1,0]]]
out = [[c1, c3], [c4, c2]]

How would one achieve such a thing in pytorch in a fast vectorized manner?

pytorch numpy python
2 Answers

Try with:

out = img[idx[...,0], idx[...,1]]
user_42384080

I was able to solve it (with the help of Quang Hoang answer)

out[idx[...,0], idx[...,1]] = img
user_14482570
Related