import matplotlib.pyplot as plt
import matplotlib.image as mpimg
import torch
import torchvision.transforms as transforms
import os
os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
### 1 载入图片并显示
myimg = mpimg.imread('img.jpg')
plt.imshow(myimg)
plt.axis('off')
plt.show()
print(myimg.shape)
# 输出 (497, 687, 3)

### 2 图片数据转化为张量
plt2tensor = transforms.ToTensor() # 实例化ToTensor()
rgb_image = plt2tensor(myimg) # 进行图片转换
print(rgb_image[0][0]) # 输出图片的部分数据
print(rgb_image.shape) # 输出图片的形状
# torch.Size([3, 497, 687])

### 3 定义Sobel卷积核
sobelfilter = torch.tensor([
    [-1.0,0.0,1.0],
    [-2.0,0.0,2.0],
    [-1.0,0.0,1.0]]*3).reshape([1,3,3,3]) # 定义Sobel卷积核
print(sobelfilter)

### 4 卷积操作
op = torch.nn.functional.conv2d(rgb_image.unsqueeze(0),sobelfilter,stride=3,padding=1) # 3个输入通道，生成1哥特征图
ret = (op - op.min()).div(op.max() - op.min()) # 卷积结果进行处理
ret = ret.clamp(0.0,1.0).mul(255).int() # 将卷积结果转化图片
print(ret)
### 灰度图
plt.imshow(ret.squeeze(),cmap='Greys_r') # 显示图片
plt.axis('off') # 不显示坐标轴
plt.show()