Есть ли в Pytorch `тензорная 'операция или функция, которая работает как cv2.dlate в OpenCV?

Question

Я построил несколько масок через сеть. Эти маски хранятся в переменной torch.tensor. Я хотел бы выполнить операцию cv2.dilate на каждом канале tensor.

Я знаю, что есть способ преобразовать tensor в numpy.ndarray и затем применить cv2.dilate к каждому каналу, используя цикл for. Но поскольку имеется около 32 каналов, этот метод может замедлить операцию пересылки в сети.

Answer 1

Я думаю, что дилат - это, по сути, операция в факеле.Смотрите код ниже

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.]]]])