Pyotrch实现单幅图像边缘检测

随笔记~

内容来自吴恩达老师的Coursera课程 Convolutional Neural Networks。

参考代码：

import torch
import torch.nn.functional as F
from PIL import Image
from torchvision import transforms

# 转换列表
tfs = [transforms.ToTensor(), transforms.ToPILImage()]

# 读取图片并进行类型转换
src = Image.open(r'C:\Users\ykx\Desktop\mess\1\cat.jpg').convert('L')
src_tensor = tfs[0](src).unsqueeze(0)  # PILimg --> tensor 并补充维度


# 自定义卷积核
# kernel = torch.tensor([[1, 0, -1], [1, 0, -1], [1, 0, -1]]).float()  # 垂直
kernel = torch.tensor([[1, 1, 1], [0, 0, 0], [-1, -1, -1]]).float()  # 水平
kernel = (kernel.unsqueeze(0)).unsqueeze(0)


# 边缘检测
dst_tensor = F.conv2d(src_tensor, kernel)
dst = tfs[1](dst_tensor.squeeze(0))  # tensor --> PILimg
dst.show()

原图 👇

垂直边缘检测 👇

水平边缘检测 👇