Masking image on GPU

vision
1

Is there a straightforward way given a mask and an image to blur the part of the image the mask corresponds to on torch.cuda (meaning doing everything on the GPU). The GaussianBlur is implemented, but how would I do the masking (would it be the same as I would do it in numpy or is there some optimized method for it)?

Thanks!

2

Yes. You can move your tensors into your device in your transformation before applying the GaussianBlur. Here is an example:

import torch
import torchvision
import torchvision.transforms as T
# ...
device = torch.device('cuda') 
transform = T.Compose([
    # ...
    T.ToTensor(),
    T.Lambda(lambda x: x.to(device))
    T.GaussianBlur((7, 13))
    # ...
])
imgs = torch.stack([copy.deepcopy(transform(img)) for i in range(10)])

for i in range(imgs.shape[0]):
  print(imgs[i].device)
# >>> cuda:0
# ...
# >>> cuda:0

You might consider looking into this post for troubleshooting when having your data loader to have multiple workers.

3

I was talking rather about the masking, but I figured it out. Thanks anyway!

1 Like