PyTorch: Vectorizing patch selection from a batch of images

user_3642151 0
Viewed 346

Suppose I have a batch of images as a tensor, for example:

images = torch.zeros(64, 3, 1024, 1024)

Now, I want to select a patch from each of those images. All the patches are of the same size, but have different starting positions for each image in the batch.

size_x = 100
size_y = 100
start_x = torch.zeros(64)
start_y = torch.zeros(64)

I can achieve the desired result like that:

result = []
for i in range(arr.shape[0]):
    result.append(arr[i, :, start_x[i]:start_x[i]+size_x, start_y[i]:start_y[i]+size_y])
result = torch.stack(result, dim=0)

The question is -- is it possible to do the same thing faster, without a loop? Perhaps there is some form of advanced indexing, or a PyTorch function that can do this?

pytorch torchvision vectorization python
1 Answers

You can use torch.take to get rid of a for loop. But first, an array of indices should be created with this function

def convert_inds(img_a,img_b,patch_a,patch_b,start_x,start_y):
    
    all_patches = np.zeros((len(start_x),3,patch_a,patch_b))
    
    patch_src = np.zeros((patch_a,patch_b))
    inds_src = np.arange(patch_b)
    patch_src[:] = inds_src
    for ind,info in enumerate(zip(start_x,start_y)):
        
        x,y = info
        if x + patch_a + 1 > img_a: return False
        if y + patch_b + 1 > img_b: return False
        start_ind = img_b * x + y
        end_ind = img_b * (x + patch_a -1) + y
        col_src = np.linspace(start_ind,end_ind,patch_b)[:,None]
        all_patches[ind,:] = patch_src + col_src
        
    return all_patches.astype(np.int)

As you can see, this function essentially creates the indices for each patch you want to slice. With this function, the problem can be easily solved by

size_x = 100
size_y = 100
start_x = torch.zeros(64)
start_y = torch.zeros(64)

images = torch.zeros(64, 3, 1024, 1024)
selected_inds = convert_inds(1024,1024,100,100,start_x,start_y)
selected_inds = torch.tensor(selected_inds)
res = torch.take(images,selected_inds)

UPDATE

OP's observation is correct, the approach above is not faster than a naive approach. In order to avoid building indices every time, here is another solution based on unfold

First, build a tensor of all the possible patches

# create all possible patches
all_patches = images.unfold(2,size_x,1).unfold(3,size_y,1)

Then, slice the desired patches from all_patches

img_ind = torch.arange(images.shape[0])
selected_patches = all_patches[img_ind,:,start_x,start_y,:,:]
user_83668050
Related