With the following 3D tensor representing an imageimg.shape=[H,W,F]
And a tensor representing the indices to that imgindices.shape=[N,2]
Eg ifindices = [[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?
Slicing would work
img[indices[:,0], indices[:,1]]
tensor([[ 0, 1, 2],
[90, 91, 92],
[36, 37, 38]])
The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.