Problem Description:

I have a loop that iterates over object IDs and assigns them to a 2D tensor (pixels_buffer). The conditions are:

• Calculate the pixel position from the object index.
• If the pixel position is unoccupied (value is -1), assign the object ID to that position, otherwise continue.

Loop Details:

• pixels_buffer: 2D tensor of shape (H, W) initialized with -1.
• object_indexers: 1D tensor of linear indices for the pixel positions (for example:[389456, 345970, …, 1661]), each index corresponds to object index in object_indices, has duplicated values.
• object_indices: 1D tensor of object IDs (for example [0, 1, 2, …, N]).

• clarification: object_indexers will contain values of the pixels indices, in their flattened form. for example: (3, 4), in (10, 10) image = 3*10 + 4 = 34.

for i in object_indices:
    idx = object_indexers[i]
    pixel_y = idx // W
    pixel_x = idx % W
    if pixels_buffer[pixel_y, pixel_x] == -1:
        pixels_buffer[pixel_y, pixel_x] = i

I’d like to vectorize this loop, in a way which preserve results and improve time complexity.
Any suggestions on how to achieve this would be greatly appreciated!

I tried using torch.unique, but it didn’t help much because the reverse_indices correspond to the unique values rather than the first occurrences.