人工智能例子汇总:AI常见的算法和例子-CSDN博客
Sobel边缘检测算法是一种用于图像处理中的边缘检测方法,它能够突出图像中灰度变化剧烈的地方,也就是边缘。该算法通过计算图像在水平方向和垂直方向上的梯度来检测边缘,梯度值越大表示图像边缘越明显
完整代码:
import torch
import torch.nn.functional as F
import cv2
import numpy as np
import matplotlib.pyplot as plt# 加载图像并转换为灰度图
image = cv2.imread('touxiang.jpg', cv2.IMREAD_GRAYSCALE)# 确保图像加载成功
if image is None:print("Error loading image")exit()# 转换为 PyTorch 张量,并进行归一化
image_tensor = torch.tensor(image, dtype=torch.float32) / 255.0
image_tensor = image_tensor.unsqueeze(0).unsqueeze(0) # 添加批次和通道维度# Sobel 卷积核
sobel_x = torch.tensor([[1, 0, -1],[2, 0, -2],[1, 0, -1]], dtype=torch.float32).unsqueeze(0).unsqueeze(0) # 形状 [1, 1, 3, 3]
sobel_y = torch.tensor([[1, 2, 1],[0, 0, 0],[-1, -2, -1]], dtype=torch.float32).unsqueeze(0).unsqueeze(0) # 形状 [1, 1, 3, 3]# 使用 PyTorch 卷积计算水平和垂直梯度
grad_x = F.conv2d(image_tensor, sobel_x, padding=1)
grad_y = F.conv2d(image_tensor, sobel_y, padding=1)# 计算梯度幅度(边缘强度)
grad_mag = torch.sqrt(grad_x**2 + grad_y**2)# 去除多余的维度,并将结果缩放至 [0, 255] 范围
grad_mag = grad_mag.squeeze().numpy() # 转为 NumPy 数组
grad_mag = np.uint8(np.clip(grad_mag * 255.0, 0, 255)) # 转为 8 位图像# 显示原图和边缘图
plt.subplot(1, 2, 1)
plt.imshow(image, cmap='gray')
plt.title("Original Image")
plt.axis('off')plt.subplot(1, 2, 2)
plt.imshow(grad_mag, cmap='gray')
plt.title("Sobel Edge Detection")
plt.axis('off')plt.show()
结果: