【python计算机视觉】3、求图像的幅度图

图像的幅度
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