如何优化Python实现Canny边缘检测器滞后阈值的运行速度？

我最近学习了Canny边缘检测器，出于兴趣从零实现了该算法。虽不指望速度媲美OpenCV等知名实现，但期望其速度能更接近一些。

解决了一系列问题后，我在滞后阈值步骤陷入瓶颈，完全不知道如何提升速度，我的代码如下：

class fnc_class:
    def __init__(self, high_im=None):
        self.pos = [0, 0]
        self.high_im = high_im
        self.shape = high_im.shape

    def update_coordinate(self):
        if (self.pos == [self.shape[0],self.shape[1]]):
            self.pos = [0,0]
        for axis in (1, 0):
            if self.pos[axis] < self.shape[axis] - 1:
                self.pos[axis] += 1
                break
            else:
                self.pos[axis] = 0

    def hysteresis_threshold(self, possible_correct):
        if (possible_correct[4] > 0.5):
            possible_correct = possible_correct.reshape(3, 3)

            if (self.high_im[(self.pos[0]-1):(self.pos[0]+2),(self.pos[1]-1):(self.pos[1]+2)].sum() > 0.5):
                self.update_coordinate()
                return possible_correct[1, 1]
        
        self.update_coordinate()
        return 0

# arr : gradient magnitudes
def hysteresis_threshold(arr, high, low):
    high_im = arr*(arr >= high)
    low_im = arr*(arr >= low)

    possible_correct = low_im - high_im
    possible_correct[possible_correct < 0] = 0

    func = fnc_class(high_im=high_im)
    possible = 0

    for i in range(0,50):
        possible = scipy.ndimage.generic_filter(possible_correct, func.hysteresis_threshold, size=(3,3))
        func.high_im = possible + high_im


    hysteresis = possible + high_im
    return hysteresis

该步骤占据了我实现的大部分计算时间。使用time.time()函数测得的耗时如下：

My time: 5.349750995635986
OpenCV time: 0.0004572868347167969

我知道用for i in range(0,50)循环来反复“扩展”high_im是一种暴力方法，但不知道还有什么其他方式可以扩展高阈值边缘。我希望尽量少用辅助库，仅使用Scipy或Numpy等必要库，请问该如何优化速度？

核心问题分析 #

你的代码速度慢的核心原因：

• 使用scipy.ndimage.generic_filter配合Python类方法做逐像素遍历，Python层面的循环和函数调用开销极大
• 手动维护pos坐标属于冗余操作，额外增加了计算负担
• 暴力循环50次的方式效率极低，滞后阈值的本质是连通性分析——只需保留低阈值区域中与高阈值区域连通的部分

高效实现方案 #

方法1：基于形态学膨胀的快速实现

利用形态学膨胀扩展高阈值边缘，直到覆盖所有连通的低阈值区域：

import numpy as np
from scipy.ndimage import binary_dilation, generate_binary_structure

def hysteresis_threshold(arr, high, low):
    # 生成高低阈值掩码
    high_mask = arr >= high
    low_mask = arr >= low
    
    # 定义8连通的膨胀结构（3x3）
    struct = generate_binary_structure(2, 2)
    
    # 迭代膨胀高阈值区域，直到不再变化
    expanded = high_mask.copy()
    while True:
        new_expanded = binary_dilation(expanded, structure=struct) & low_mask
        if np.array_equal(new_expanded, expanded):
            break
        expanded = new_expanded
    
    # 返回带梯度值的边缘结果（若只需掩码可直接返回expanded）
    return arr * expanded

方法2：基于连通域分析的实现

直接标记低阈值区域的连通分量，保留与高阈值区域重叠的分量：

import numpy as np
from scipy.ndimage import label

def hysteresis_threshold(arr, high, low):
    high_mask = arr >= high
    low_mask = arr >= low
    
    # 给低阈值区域的连通域打标签
    labeled, _ = label(low_mask)
    
    # 筛选出与高阈值区域重叠的连通域标签
    connected_labels = np.unique(labeled[high_mask])
    connected_labels = connected_labels[connected_labels != 0]  # 排除背景标签
    
    # 生成最终边缘掩码
    final_mask = np.isin(labeled, connected_labels)
    
    return arr * final_mask

速度提升说明 #

• 两种方案均依赖Scipy/Numpy的底层C实现，彻底避免了Python层面的逐像素循环
• 形态学膨胀的迭代次数远少于50次（通常仅需几次即可完成所有连通区域扩展）
• 连通域分析直接一次性完成所有连通性判断，无需循环迭代

额外优化建议 #

• 若仅需边缘掩码，可直接返回布尔数组，进一步减少计算量
• 可预定义膨胀结构，避免每次调用重复生成
• 若场景适合，也可尝试binary_closing形态学操作简化流程

内容的提问来源于stack exchange，提问作者Khoa LT