canny算子python代码

Canny算子是一种常用的边缘检测算法，其Python实现如下：

import cv2
import numpy as np

def canny(image, low_threshold, high_threshold):
    gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    blur = cv2.GaussianBlur(gray, (5, 5), 0)
    canny = cv2.Canny(blur, low_threshold, high_threshold)
    return canny

其中，`image`为输入的图像，`low_threshold`和`high_threshold`为双阈值滤波的低阈值和高阈值。函数首先将图像转换为灰度图像，然后进行高斯模糊，最后使用Canny算子进行边缘检测。

另外，如果需要使用PyTorch实现Canny算子，可以使用以下代码：

import torch
import torch.nn.functional as F

def canny_pytorch(image, low_threshold, high_threshold):
    gray = torch.mean(image, dim=0, keepdim=True)
    blur = F.conv2d(gray, torch.ones(1, 1, 5, 5) / 25, padding=2)
    dx = F.conv2d(blur, torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=torch.float32).unsqueeze(0).unsqueeze(0), padding=1)
    dy = F.conv2d(blur, torch.tensor([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], dtype=torch.float32).unsqueeze(0).unsqueeze(0), padding=1)
    gradient = torch.sqrt(dx ** 2 + dy ** 2)
    theta = torch.atan2(dy, dx)
    theta[theta < 0] += np.pi
    theta[theta >= np.pi * 7 / 8] = 0
    theta[(theta >= np.pi * 1 / 8) & (theta < np.pi * 3 / 8)] = np.pi / 4
    theta[(theta >= np.pi * 3 / 8) & (theta < np.pi * 5 / 8)] = np.pi / 2
    theta[(theta >= np.pi * 5 / 8) & (theta < np.pi * 7 / 8)] = np.pi * 3 / 4
    suppressed = gradient.clone()
    for i in range(1, suppressed.shape[2] - 1):
        for j in range(1, suppressed.shape[3] - 1):
            if theta[0, 0, i, j] == 0:
                if (gradient[0, 0, i, j] <= gradient[0, 0, i, j - 1]) or (gradient[0, 0, i, j] <= gradient[0, 0, i, j + 1]):
                    suppressed[0, 0, i, j] = 0
            elif theta[0, 0, i, j] == np.pi / 4:
                if (gradient[0, 0, i, j] <= gradient[0, 0, i - 1, j + 1]) or (gradient[0, 0, i, j] <= gradient[0, 0, i + 1, j - 1]):
                    suppressed[0, 0, i, j] = 0
            elif theta[0, 0, i, j] == np.pi / 2:
                if (gradient[0, 0, i, j] <= gradient[0, 0, i - 1, j]) or (gradient[0, 0, i, j] <= gradient[0, 0, i + 1, j]):
                    suppressed[0, 0, i, j] = 0
            elif theta[0, 0, i, j] == np.pi * 3 / 4:
                if (gradient[0, 0, i, j] <= gradient[0, 0, i - 1, j - 1]) or (gradient[0, 0, i, j] <= gradient[0, 0, i + 1, j + 1]):
                    suppressed[0, 0, i, j] = 0
    suppressed[gradient < low_threshold] = 0
    suppressed[gradient > high_threshold] = 255
    for i in range(1, suppressed.shape[2] - 1):
        for j in range(1, suppressed.shape[3] - 1):
            if suppressed[0, 0, i, j] == 255:
                if (suppressed[0, 0, i - 1, j] == 0) and (suppressed[0, 0, i + 1, j] == 0) and (suppressed[0, 0, i, j - 1] == 0) and (suppressed[0, 0, i, j + 1] == 0):
                    suppressed[0, 0, i, j] = 0
    return suppressed

其中，`image`为输入的图像，`low_threshold`和`high_threshold`为双阈值滤波的低阈值和高阈值。函数首先将图像转换为灰度图像，然后进行高斯模糊，接着计算图像的梯度和方向，进行非极大值抑制和双阈值滤波，最后进行边缘连接。