用pytorch 直接卷积一张图片的灰度图并生成轮廓

用pytorch 直接卷积一张图片.

import torch.nn.functional as F
import torch
import torch.nn as nn
from torch.nn.modules.conv import Conv2d 
import numpy as np
import torch
from torch import nn
from torch.autograd import Variable
import torch.nn.functional as F
from PIL import Image
import matplotlib.pyplot as plt 
import cv2

torch.set_default_tensor_type(torch.DoubleTensor)

img = cv2.imread("46455close.jpg",0) # 0 参数是读取的灰度图
print(img.shape)
 
# 将图片矩阵转化为 pytorch tensor，并适配卷积输入的要求 
# print( img.reshape((1, 3, img.shape[0], img.shape[1])) )
img = img.astype('uint8')
w = img.shape[0]
h =  img.shape[1]
imgtensor = torch.from_numpy(img.reshape((1, 1, w, h)))
print(imgtensor.shape)
# imgtensor = imgtensor.unsqueeze(0) # 三通道数据如果输入到conv2d中应该是四维数据。
# conv = nn.Conv2d(in_channels=3, out_channels=1, kernel_size=1)
# res = conv(Variable(imgtensor))
conv1 = nn.Conv2d(in_channels=3, out_channels=1, kernel_size=(3,3)) # 定义卷积
sobel_kernel = np.array([
  # [-1, -1, -1], [-1, 8, -1], [-1, -1, -1],
  # [-1, -1, -1], [-1, 8, -1], [-1, -1, -1],
  [-1, -1, -1], [-1, 8, -1], [-1, -1, -1]
  ], dtype='float32') # 定义轮廓检测算子
sobel_kernel = sobel_kernel.reshape((1,1, 3, 3)) # 适配卷积的输入输出
conv1.weight.data = torch.from_numpy(sobel_kernel).double() # 给卷积的 kernel 赋值
# conv1 = conv1.double()
edge1 = conv1(Variable(imgtensor).double()) # 作用在图片上
edge1 = edge1.data.squeeze().numpy() # 将输出转换为图片的格式 
# plt.imshow(edge1, cmap='gray')
# plt.waitforbuttonpress()
print(edge1.shape)
imgaaa = Image.fromarray(edge1)
imgaaa.show()

# plt.imshow(edge1)

参考文章
https://blog.csdn.net/weixin_40123108/article/details/83510592