### 使用 PyTorch 实现 Canny 边缘检测 尽管传统的 Canny 边缘检测通常通过 OpenCV 来实现[^1]，但在某些情况下可能需要利用深度学习框架如 PyTorch 进行类似的处理。需要注意的是，Canny 边缘检测本身并不是基于深度学习的方法，而是一个经典的图像处理技术[^4]。然而，在 PyTorch 中可以通过卷积操作模拟 Sobel 或其他梯度算子来近似部分功能。 以下是在 PyTorch 中实现 Canny 边缘检测的主要思路： #### 1. 高斯模糊 高斯滤波用于减少噪声并平滑图像。这一步可以使用 `torch.nn.functional.conv2d` 和自定义的高斯核完成。 ```python import torch import torch.nn as nn import torch.nn.functional as F def gaussian_kernel(size, sigma=1.0): x_coord = torch.arange(size) y_coord = torch.arange(size) grid_x, grid_y = torch.meshgrid(x_coord, y_coord, indexing='ij') mean = (size - 1) / 2. variance = sigma ** 2. gaussian_filter = torch.exp(-((grid_x - mean)**2 + (grid_y - mean)**2) / (2 * variance)) gaussian_filter /= torch.sum(gaussian_filter) return gaussian_filter.unsqueeze(0).unsqueeze(0) def apply_gaussian_blur(image_tensor, kernel_size=5, sigma=1.0): batch_size, channels, height, width = image_tensor.shape # 创建高斯核 kernel = gaussian_kernel(kernel_size, sigma=sigma).to(image_tensor.device) kernel = kernel.expand(channels, 1, kernel_size, kernel_size) # 应用卷积 blurred_image = F.conv2d(image_tensor.view(batch_size*channels, 1, height, width), kernel, padding=(kernel_size//2)) blurred_image = blurred_image.view(batch_size, channels, height, width) return blurred_image ``` #### 2. 计算梯度幅值和方向 此步骤涉及计算水平和垂直方向上的梯度，并进一步求得梯度的方向和大小。 ```python def sobel_filters(image_tensor): batch_size, channels, height, width = image_tensor.shape # 定义 Sobel 算子 sobel_x = torch.tensor([[-1., 0., 1.], [-2., 0., 2.], [-1., 0., 1.]]).view(1, 1, 3, 3).to(image_tensor.device) sobel_y = torch.tensor([[1., 2., 1.], [0., 0., 0.], [-1., -2., -1.]]).view(1, 1, 3, 3).to(image_tensor.device) # 扩展通道维度以适应多通道输入 sobel_x = sobel_x.expand(channels, 1, 3, 3) sobel_y = sobel_y.expand(channels, 1, 3, 3) grad_x = F.conv2d(image_tensor.view(batch_size*channels, 1, height, width), sobel_x, padding=1) grad_y = F.conv2d(image_tensor.view(batch_size*channels, 1, height, width), sobel_y, padding=1) grad_x = grad_x.view(batch_size, channels, height, width) grad_y = grad_y.view(batch_size, channels, height, width) gradient_magnitude = torch.sqrt(grad_x**2 + grad_y**2) gradient_direction = torch.atan2(grad_y, grad_x) * (180 / torch.pi) % 180 return gradient_magnitude, gradient_direction ``` #### 3. 非极大值抑制 该过程保留局部最大值作为候选边缘像素。 ```python def non_max_suppression(magnitude, direction): _, _, h, w = magnitude.size() result = torch.zeros_like(magnitude) for i in range(1, h-1): for j in range(1, w-1): angle = direction[:, :, i, j].item() % 180 if (0 <= angle < 22.5 or 157.5 <= angle <= 180): # 方向接近 0°/180° q = magnitude[:, :, i, j+1] r = magnitude[:, :, i, j-1] elif (22.5 <= angle < 67.5): # 方向接近 45° q = magnitude[:, :, i+1, j-1] r = magnitude[:, :, i-1, j+1] elif (67.5 <= angle < 112.5): # 方向接近 90° q = magnitude[:, :, i+1, j] r = magnitude[:, :, i-1, j] elif (112.5 <= angle < 157.5): # 方向接近 135° q = magnitude[:, :, i-1, j-1] r = magnitude[:, :, i+1, j+1] if (magnitude[:, :, i, j] >= q and magnitude[:, :, i, j] >= r): result[:, :, i, j] = magnitude[:, :, i, j] return result ``` #### 4. 双阈值检测与滞后连接 最后一步是应用双阈值策略筛选强弱边缘，并通过滞后方法连接它们。 ```python def double_threshold_hysteresis(edge_map, low_ratio=0.05, high_ratio=0.15): max_val = edge_map.max().item() weak_edge = edge_map.new_full(edge_map.size(), fill_value=low_ratio * max_val) strong_edge = edge_map.new_full(edge_map.size(), fill_value=high_ratio * max_val) thresholded_edges = torch.where( (edge_map >= strong_edge), edge_map, torch.where( (edge_map >= weak_edge), weak_edge, torch.zeros_like(edge_map) ) ) final_edges = torch.zeros_like(thresholded_edges) indices_strong = (thresholded_edges == strong_edge.item()) queue = [(i,j,k,l) for l in range(indices_strong.size(0)) for k in range(indices_strong.size(1)) \ for j in range(indices_strong.size(2)) for i in range(indices_strong.size(3)) if indices_strong[l][k][j][i]] while len(queue) > 0: x,y,z,w = queue.pop(0) neighbors = [ (x-1, y, z, w),(x+1, y, z, w),(x, y-1, z, w),(x, y+1, z, w), (x-1, y-1, z, w),(x-1, y+1, z, w),(x+1, y-1, z, w),(x+1, y+1, z, w) ] for nx, ny, nz, nw in neighbors: if not (nx>=0 and nx<thresholded_edges.size(3)): continue if not (ny>=0 and ny<thresholded_edges.size(2)): continue if not (nz>=0 and nz<thresholded_edges.size(1)): continue if not (nw>=0 and nw<thresholded_edges.size(0)): continue if thresholded_edges[nw,nz,ny,nx]==weak_edge.item(): thresholded_edges[nw,nz,ny,nx]=strong_edge.item() queue.append((nx,ny,nz,nw)) final_edges[thresholded_edges==strong_edge.item()] = 1. return final_edges ``` --- ### 总结 上述代码展示了如何在 PyTorch 中逐步实现类似于 Canny 的边缘检测流程。虽然这些函数能够模仿传统算法的行为，但由于缺乏优化，性能上可能会逊色于专门设计的库（如 OpenCV）。因此建议仅将其视为一种实验性的替代方案[^3]。