快捷導(dǎo)航

Pytorch 實現(xiàn)sobel算子的卷積操作詳解

更新時間：2020年01月10日 17:13:21 作者：洪流之源

今天小編就為大家分享一篇Pytorch 實現(xiàn)sobel算子的卷積操作詳解，具有很好的參考價值，希望對大家有所幫助。一起跟隨小編過來看看吧

卷積在pytorch中有兩種實現(xiàn)，一種是torch.nn.Conv2d()，一種是torch.nn.functional.conv2d()，這兩種方式本質(zhì)都是執(zhí)行卷積操作，對輸入的要求也是一樣的，首先需要輸入的是一個torch.autograd.Variable()的類型，大小是（batch，channel， H，W)，其中batch表示輸入的一批數(shù)據(jù)的數(shù)目，channel表示輸入的通道數(shù)。

一般一張彩色的圖片是3，灰度圖片是1，而卷積網(wǎng)絡(luò)過程中的通道數(shù)比較大，會出現(xiàn)幾十到幾百的通道數(shù)。H和W表示輸入圖片的高度和寬度，比如一個batch是32張圖片，每張圖片是3通道，高和寬分別是50和100，那么輸入的大小就是（32，3，50，100）。

如下代碼是卷積執(zhí)行soble邊緣檢測算子的實現(xiàn)：

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算子參數(shù)
  sobel_kernel = np.array([[-1, -1, -1], [-1, 8, -1], [-1, -1, -1]], dtype='float32')
  # 將sobel算子轉(zhuǎn)換為適配卷積操作的卷積核
  sobel_kernel = sobel_kernel.reshape((1, 1, 3, 3))
  # 給卷積操作的卷積核賦值
  conv_op.weight.data = torch.from_numpy(sobel_kernel)
  # 對圖像進行卷積操作
  edge_detect = conv_op(Variable(im))
  # 將輸出轉(zhuǎn)換為圖片格式
  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():
  # 讀入一張圖片，并轉(zhuǎn)換為灰度圖
  im = Image.open('./cat.jpg').convert('L')
  # 將圖片數(shù)據(jù)轉(zhuǎn)換為矩陣
  im = np.array(im, dtype='float32')
  # 將圖片矩陣轉(zhuǎn)換為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數(shù)據(jù)轉(zhuǎn)換為image
  im = Image.fromarray(edge_detect)
  # image數(shù)據(jù)轉(zhuǎn)換為灰度模式
  im = im.convert('L')
  # 保存圖片
  im.save('edge.jpg', quality=95)
 
if __name__ == "__main__":
  main()

原圖片：cat.jpg