[FIXED] PyTorch - read single pixel from each image channel and write it to another image

Issue

I have a tensor [C, H, W], where C is a number of channels (does not have to be 3). I have another array with indices and I would like to read pixels at those indices and write them to another image.

For example:

two images A (source) and B (target)

C = 3
H, W = 64

indicesRead = [[0, 0], [13, 15], [32, 43]]
indicesWrite = [[7, 5], [1, 1], [4, 4]]

and I would like to get from image A pixel for channel 0 at [0, 0], channel 1 at [13, 15] and channel 2 at [32, 43].

Once I have these values, I want to write them to image B to channel 0 to position [7, 5] (basically copy A[0, 0] to B[7, 5]) etc.

Can it be done with torch methods or I have to iterate tensor manually?

Solution

You can use list-based indexing in Pytorch.

# add channel index to indicesRead
indicesRead = [[0,0, 0], [1,13, 15], [2,32, 43]]
indicesWrite = [[7, 5], [1, 1], [4, 4]]


cRead, aRead , bRead  = zip(*indicesRead)
aWrite, bWrite = zip(*indicesWrite)

B[aWrite,bWrite]=  A[cRead,aRead,bRead] 

Note I use a and b to denote height and width dimensions (no correlation to input and output matrices A and B because the x y (column-first) convention becomes a bit confusing with arrays. I use c to indicate channel.

Answered By - DerekG