Pytorch 实现sobel算子的卷积操作

本是古典何须时尚 2024-04-18 07:10 9阅读 0赞

卷积在pytorch中有两种实现，一种是torch.nn.Conv2d()，一种是torch.nn.functional.conv2d()，这两种方式本质都是执行卷积操作，对输入的要求也是一样的，首先需要输入的是一个torch.autograd.Variable()的类型，大小是（batch，channel， H，W)，其中batch表示输入的一批数据的数目，channel表示输入的通道数。一般一张彩色的图片是3，灰度图片是1，而卷积网络过程中的通道数比较大，会出现几十到几百的通道数。H和W表示输入图片的高度和宽度，比如一个batch是32张图片，每张图片是3通道，高和宽分别是50和100，那么输入的大小就是（32，3，50，100）。如下代码是卷积执行soble边缘检测算子的实现：

import torch
    import numpy as np
    from torch import nn
    from PIL import Image
    from torch.autograd import Variable
    import torch.nn.functional as F
    
    
    def nn_conv2d(im):
        # 用nn.Conv2d定义卷积操作
        conv_op = nn.Conv2d(1, 1, 3, bias=False)
        # 定义sobel算子参数
        sobel_kernel = np.array([[-1, -1, -1], [-1, 8, -1], [-1, -1, -1]], dtype='float32')
        # 将sobel算子转换为适配卷积操作的卷积核
        sobel_kernel = sobel_kernel.reshape((1, 1, 3, 3))
        # 给卷积操作的卷积核赋值
        conv_op.weight.data = torch.from_numpy(sobel_kernel)
        # 对图像进行卷积操作
        edge_detect = conv_op(Variable(im))
        # 将输出转换为图片格式
        edge_detect = edge_detect.squeeze().detach().numpy()
        return edge_detect
    
    def functional_conv2d(im):
        sobel_kernel = np.array([[-1, -1, -1], [-1, 8, -1], [-1, -1, -1]], dtype='float32')  #
        sobel_kernel = sobel_kernel.reshape((1, 1, 3, 3))
        weight = Variable(torch.from_numpy(sobel_kernel))
        edge_detect = F.conv2d(Variable(im), weight)
        edge_detect = edge_detect.squeeze().detach().numpy()
        return edge_detect
    
    def main():
        # 读入一张图片，并转换为灰度图
        im = Image.open('./cat.jpg').convert('L')
        # 将图片数据转换为矩阵
        im = np.array(im, dtype='float32')
        # 将图片矩阵转换为pytorch tensor,并适配卷积输入的要求
        im = torch.from_numpy(im.reshape((1, 1, im.shape[0], im.shape[1])))
        # 边缘检测操作
        # edge_detect = nn_conv2d(im)
        edge_detect = functional_conv2d(im)
        # 将array数据转换为image
        im = Image.fromarray(edge_detect)
        # image数据转换为灰度模式
        im = im.convert('L')
        # 保存图片
        im.save('edge.jpg', quality=95)
    
    if __name__ == "__main__":
        main()

原图片：cat.jpg

![watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3dlaWNhbzE5OTA_size_16_color_FFFFFF_t_70][]

结果图片：edge.jpg

![watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3dlaWNhbzE5OTA_size_16_color_FFFFFF_t_70 1][]

[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3dlaWNhbzE5OTA_size_16_color_FFFFFF_t_70]: https://image.dandelioncloud.cn/pgy_files/images/2024/04/17/558a63c7fa5745e69837b7904a178ec8.jpg
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3dlaWNhbzE5OTA_size_16_color_FFFFFF_t_70 1]: https://image.dandelioncloud.cn/pgy_files/images/2024/04/17/9365a936c95545968d3d2d2e0caa7d19.jpg