Is there a `tensor` operation or function in Pytorch that works like cv2.dilate in OpenCV?

Question

I built several masks through a network. These masks are stored in a torch.tensor variable. I would like to do a cv2.dilate like operation on every channel of the tensor.

I know there is a way that convert the tensor to numpy.ndarray and then apply cv2.dilate to every channel using a for loop. But since there are about 32 channels, this method might slow down the forward operation in the network.

Answer 1

I think dilate is essentially conv2d operation in torch. See the code below

import cv2
import numpy as np
import torch

im = np.array([ [0, 0, 0, 0, 0],
                [0, 1, 0, 0, 0],
                [0, 1, 1, 0, 0],
                [0, 0, 0, 1, 0],
                [0, 0, 0, 0, 0] ], dtype=np.float32)
kernel = np.array([ [1, 1, 1],
                    [1, 1, 1],
                    [1, 1, 1] ], dtype=np.float32)
print(cv2.dilate(im, kernel))
# [[1. 1. 1. 0. 0.]
#  [1. 1. 1. 1. 0.]
#  [1. 1. 1. 1. 1.]
#  [1. 1. 1. 1. 1.]
#  [0. 0. 1. 1. 1.]]
im_tensor = torch.Tensor(np.expand_dims(np.expand_dims(im, 0), 0)) # size:(1, 1, 5, 5)
kernel_tensor = torch.Tensor(np.expand_dims(np.expand_dims(kernel, 0), 0)) # size: (1, 1, 3, 3)
torch_result = torch.clamp(torch.nn.functional.conv2d(im_tensor, kernel_tensor, padding=(1, 1)), 0, 1)
print(torch_result)
# tensor([[[[1., 1., 1., 0., 0.],
#           [1., 1., 1., 1., 0.],
#           [1., 1., 1., 1., 1.],
#           [1., 1., 1., 1., 1.],
#           [0., 0., 1., 1., 1.]]]])