图像的幅度
python代码实现
import torch
import torch.nn as nn
import torch.nn.functional as F
#输入四维张量,输出四维张量
class Gradient(nn.Module):
def __init__(self, channels):
super(Gradient, self).__init__()
kernel_gauss = torch.FloatTensor([[0.0924, 0.1192, 0.0924], [0.1192, 0.1538, 0.1192], [0.0924, 0.1192, 0.0924]]).unsqueeze(0).unsqueeze(0)
kernel_gauss = np.repeat(kernel_gauss,channels,axis=0)
kernel_x = torch.FloatTensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]]).unsqueeze(0).unsqueeze(0)
kernel_x = np.repeat(kernel_x,channels,axis=0)
kernel_y = torch.FloatTensor([[1, 2, 1], [0, 0, 0], [-1, -2, -1]]).unsqueeze(0).unsqueeze(0)
kernel_y = np.repeat(kernel_y,channels,axis=0)
self.channels =channels
self.weight_gauss = nn.Parameter(data=kernel_gauss, requires_grad=False)
self.weight_x = nn.Parameter(data=kernel_x, requires_grad=False)
self.weight_y = nn.Parameter(data=kernel_y, requires_grad=False)
def forward(self, image):
b,c,w,h = image.size()
gauss = F.conv2d(image, self.weight_gauss,padding=1,groups=self.channels)
image = F.conv2d(image, self.weight_gauss, padding=1,groups=self.channels)
image_x = F.conv2d(image, self.weight_x, padding=1,groups=self.channels)
image_y = F.conv2d(image, self.weight_y, padding=1,groups=self.channels)
image_x = torch.pow(image_x, 2)
image_y = torch.pow(image_y, 2)
edge = image_x + image_y
edge = torch.sqrt(edge)
return edge