选中内容（绿色）时除了会搜索文章名，还会搜索文章内容
点击结果中的文章名进入文章界面后可以按Ctrl+F在页面内搜索
  • 版权：CC BY-SA 4.0
  • 创建：2021-07-26
  • 更新：2021-08-03
Pytorch 使用模型(卷积/conv)实现 sobel(索贝尔) 边缘检测实现源码
版权声明：本文为 neucrack 的原创文章，遵循 CC 4.0 BY-SA 版权协议，转载请附上原文出处链接及本声明。
原文链接(持续更新）：https://neucrack.com/p/377

边缘检测效果

sobel_edge2

final

sobel_edge

sobel_v831

源码在末尾

边缘检测原理

边缘就是值变化剧烈的地方, 如果对值的变化求导, 则边缘部分就是导数局部最大.
但是在图像处理时没有具体的函数让我们求导, 使用卷积运算则可以很好的近似替代

如下图, 假设左上为坐标原点, 横轴为 x, 纵轴为y, 如下图左上角9个像素点, P(x, y)表示坐标(x, y)的点, 要求P(1, 1)处在x轴的变化率, 则只需将P(2, 1) - P(0, 1) 得到值为0, P(1, 0)处为1-3 = -2, 这个差值即变化率, 类比成导数, 我们就能知道横轴在哪些地方变化率更大.

上面这种方法我们可以得到横轴的变化率, 这里使用卷积核

  1. [-1, 0, 1],
  2. [-2, 0, 2],
  3. [-1, 0, 1]

对图像进行卷积运算, 如图中的计算方法, 像素点左右权值取2, 角上的也参与计算,但是权值为1,没有左右边的权值高. 这样我们就得到了横轴的变化率图, 即边缘检测图.

注意, 这里是对横轴计算了, 比较的点左右的值变化, 所以实际看到的图像会出现明显的纵轴边缘, 如下图左边

vertical_horizontal

同理, 上图右边的图使用卷积核

  1. [1,2,1],
  2. [0,0,0],
  3. [-1, -2, -1]

得到的纵轴的边缘图.

注意这里用右边减左边, 如果右边的值比左边的小会是负数, 如果我们希望只检测颜色值变大(变白)则可以直接使用, 如果两个变化方向都要检测, 则可以取绝对值. 比如下图左边是没有取绝对值, 右边取了绝对值

without_with_abs

得到两个方向的图后, 对其进行合并, 对每个像素平方和开方即可

final

这张图左边是使用 GIMP 的 sobel 边缘检测(垂直+水平)的效果, 略微有点不同:

sobel_edge2

不同的原因是使用水平和垂直的图平方和开根后, 直接用 plt.imshow 显示, 和 GIMP 的处理方式不同

  1. out = np.sqrt(np.square(out_v) + np.square(out_h))
  2. plt.imshow(out)

简单地直接将值规范到[0, 255] 就和 GIMP 的图相似了(但不完全一样)

  1. out = np.sqrt(np.square(out_v) + np.square(out_h))
  2. out = out * 255.0 / out.max()
  3. plt.imshow(out.astype(np.uint8))

sobel_v_h

自定义卷积核来实现边缘检测

除了上面说了使用两次卷积计算, 也可以用只计算一次的卷积核, 比如:

  1. [-1, -1, -1],
  2. [ -1, 8, -1],
  3. [ -1, -1, -1]

这是对于一个通道(灰度图)来说, 如果要扩充到三个通道(RGB), 卷积核参数就是如下形式

  1. conv_rgb_core_sobel = [
  2.                         [[-1,-1,-1],[-1,8,-1], [-1,    -1,    -1],
  3.                          [0,0,0],[0,0,0], [0,0,0],
  4.                          [0,0,0],[0,0,0], [0,0,0]
  5.                         ],
  6.                         [[0,0,0],[0,0,0], [0,0,0],
  7.                          [-1,-1,-1],[-1,8,-1], [-1,    -1,    -1],
  8.                          [0,0,0],[0,0,0], [0,0,0]
  9.                         ],
  10.                         [[0,0,0],[0,0,0], [0,0,0],
  11.                          [0,0,0],[0,0,0], [0,0,0],
  12.                          [-1,-1,-1],[-1,8,-1], [-1,    -1,    -1],
  13.                         ]]

经过卷积运算后, 前后图如下:

sobel_edge

注意, 输入值范围如果为[0, 255], 输出值则范围会变化, 以图片形式查看时需要注意加以处理, 这里使用了plt.imshow(out)来显示, 这个函数会自动对图像做简单的处理, 才会看起来是黑色背景

导出成模型使用

可以将 Net 导出成 onnx 即可在其它平台使用, 就是一个简单的卷积层

部署到 V831 后的样子(使用了卷积核[-1,-1,-1],[-1,8,-1], [-1, -1, -1],):

sobel_v831

源码

  1. '''
  2.     simple sobel edge demo
  3.     visit: https://neucrack.com/p/377
  4.     @author neucrack
  5.     @license MIT
  6. '''
  8. import torch
  9. import torch.nn as nn
  10. import numpy as np
  11. import cv2
  12. from PIL import Image
  13. import matplotlib.pyplot as plt
  15. class Net(nn.Module):
  17.     def __init__(self):
  18.         super(Net, self).__init__()
  19.         self.conv1 = nn.Conv2d(3, 3, 3, padding=(0, 0), bias=False)
  21.     def forward(self, x):
  22.         x = self.conv1(x)
  23.         return x
  27. net = Net()
  29. conv_rgb_core_original = [
  30.                         [[0,0,0],[0,1,0], [0,0,0],
  31.                          [0,0,0],[0,0,0], [0,0,0],
  32.                          [0,0,0],[0,0,0], [0,0,0]
  33.                         ],
  34.                         [[0,0,0],[0,0,0], [0,0,0],
  35.                          [0,0,0],[0,1,0], [0,0,0],
  36.                          [0,0,0],[0,0,0], [0,0,0]
  37.                         ],
  38.                         [[0,0,0],[0,0,0], [0,0,0],
  39.                          [0,0,0],[0,0,0], [0,0,0],
  40.                          [0,0,0],[0,1,0], [0,0,0]
  41.                         ]]
  43. conv_rgb_core_sobel = [
  44.                         [[-1,-1,-1],[-1,8,-1], [-1,    -1,    -1],
  45.                          [0,0,0],[0,0,0], [0,0,0],
  46.                          [0,0,0],[0,0,0], [0,0,0]
  47.                         ],
  48.                         [[0,0,0],[0,0,0], [0,0,0],
  49.                          [-1,-1,-1],[-1,8,-1], [-1,    -1,    -1],
  50.                          [0,0,0],[0,0,0], [0,0,0]
  51.                         ],
  52.                         [[0,0,0],[0,0,0], [0,0,0],
  53.                          [0,0,0],[0,0,0], [0,0,0],
  54.                          [-1,-1,-1],[-1,8,-1], [-1,    -1,    -1],
  55.                         ]]
  57. conv_rgb_core_sobel_vertical = [
  58.                         [[-1,0,1],[-2,0,2], [-1,    0,    1],
  59.                          [0,0,0],[0,0,0], [0,0,0],
  60.                          [0,0,0],[0,0,0], [0,0,0]
  61.                         ],
  62.                         [[0,0,0],[0,0,0], [0,0,0],
  63.                          [-1,0,1],[-2,0,2], [-1,    0,    1],
  64.                          [0,0,0],[0,0,0], [0,0,0]
  65.                         ],
  66.                         [[0,0,0],[0,0,0], [0,0,0],
  67.                          [0,0,0],[0,0,0], [0,0,0],
  68.                          [-1,0,1],[-2,0,2], [-1,    0,    1],
  69.                         ]]
  71. conv_rgb_core_sobel_horizontal = [
  72.                         [[1,2,1],[0,0,0], [-1, -2, -1],
  73.                          [0,0,0],[0,0,0], [0,0,0],
  74.                          [0,0,0],[0,0,0], [0,0,0]
  75.                         ],
  76.                         [[0,0,0],[0,0,0], [0,0,0],
  77.                          [1,2,1],[0,0,0], [-1, -2, -1],
  78.                          [0,0,0],[0,0,0], [0,0,0]
  79.                         ],
  80.                         [[0,0,0],[0,0,0], [0,0,0],
  81.                          [0,0,0],[0,0,0], [0,0,0],
  82.                          [1,2,1],[0,0,0], [-1, -2, -1],
  83.                         ]]
  85. def sobel(net, kernel):
  86.     sobel_kernel = np.array(kernel,    dtype='float32')
  87.     sobel_kernel = sobel_kernel.reshape((3,    3,    3,    3))
  88.     net.conv1.weight.data = torch.from_numpy(sobel_kernel)
  90. params = list(net.parameters())
  93. img = cv2.imread("out/test.jpg")
  94. input_img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
  95. input_tensor = (input_img.astype(np.float32) - 127.5) / 128 # to [-1, 1]
  96. input_tensor = torch.Tensor(input_tensor).permute((2, 0, 1))
  97. print(input_tensor.shape)
  99. input_tensor = input_tensor.unsqueeze(0)
  100. print("input shape:", input_tensor.shape)
  102. sobel(net, conv_rgb_core_sobel)
  103. out = net(input_tensor).detach().numpy()[0].transpose([1,2,0])
  104. sobel(net, conv_rgb_core_sobel_vertical)
  105. out_v = net(input_tensor).detach().numpy()[0].transpose([1,2,0])
  106. sobel(net, conv_rgb_core_sobel_horizontal)
  107. out_h = net(input_tensor).detach().numpy()[0].transpose([1,2,0])
  109. print("out shape: {}, tensor:{}".format(out.shape, out))
  110. print(out.shape, out.max(), out.min())
  113. plt.figure()
  115. plt.figure()
  116. plt.subplot(1, 5, 1)
  117. input = input_tensor.numpy()[0].transpose((1,2,0))
  118. print(input.max(), input.min())
  119. plt.imshow(input_img)
  120. plt.subplot(1, 5, 2)
  121. print(out.max(), out.min())
  122. # out = np.sqrt(np.square(out))
  123. # out = out * 255.0 / out.max()
  124. # out = out.astype(np.uint8)
  125. # print(out.max(), out.min())
  126. plt.imshow(out)
  128. plt.subplot(1, 5, 3)
  129. out = np.abs(out_v)
  130. # out = out * 255.0 / out.max()
  131. # plt.imshow(out.astype(np.uint8))
  132. plt.imshow(out)
  134. plt.subplot(1, 5, 4)
  135. out = np.abs(out_h)
  136. # out = out * 255.0 / out.max()
  137. # plt.imshow(out.astype(np.uint8))
  138. plt.imshow(out)
  140. plt.subplot(1, 5, 5)
  141. out = np.sqrt(np.square(out_v) + np.square(out_h))
  142. # out = out * 255.0 / out.max()
  143. # plt.imshow(out.astype(np.uint8))
  144. plt.imshow(out)
  146. plt.show()

上述代码效果：
demo

更多

如何部署到 V831/ V833 上可以参考: 在V831上（awnn）跑 pytorch resnet18 模型

参考

文章有误？有想法想讨论？查看或者发起勘误/讨论 主题
(发起评论需要先登录 github)

/wallpaper/wallhaven-g8oepq.jpg
©Neucrack