如何用Numpy优化黑白图像像素到最近边缘的距离计算？

已通过find_edges函数实现黑白图像边缘提取（保留与黑像素相邻的白像素，其余转为黑色），现需计算图像中每个像素到最近白边缘像素的距离。当前采用嵌套循环调用find_nearest_edge的方式速度极慢，询问能否仅用Numpy实现该功能，无需逐元素循环调用。

原代码 #

from PIL import Image
import numpy as np
from scipy.ndimage import binary_erosion

def find_nearest_edge(arr, point):
    w, h = arr.shape
    x, y = point
    xcoords, ycoords = np.meshgrid(np.arange(w), np.arange(h))

    target = np.sqrt((xcoords - x)**2 + (ycoords - y)**2)
    target[arr == 0] = np.inf

    shortest_distance = np.min(target[target > 0.0])

    return shortest_distance

def find_edges(img):
    img = img.convert('L')
    img_np = np.array(img)

    kernel = np.ones((3,3))
    edges = img_np - binary_erosion(img_np, kernel)*255

    return edges

a = Image.open('a.png')
x, y = a.size

edges = find_edges(a)

out = Image.fromarray(edges.astype('uint8'), 'L')
out.save('b.png')

dists =[]
for _x in range(x):
    for _y in range(y):
        dist = find_nearest_edge(edges,(_x,_y))
        dists.append(dist)

print(dists)

示例图像 #

• 原始图像：黑白风格的图案
• 边缘提取后图像：仅保留边缘白像素的黑白图

当然可以用Numpy的广播机制实现，彻底抛弃逐元素循环，速度能提升好几个数量级。核心思路是先一次性提取所有边缘像素的坐标，再通过广播批量计算每个像素到所有边缘像素的距离，最后取最小值即可。

纯Numpy实现代码 #

from PIL import Image
import numpy as np
from scipy.ndimage import binary_erosion

def find_edges(img):
    img = img.convert('L')
    img_np = np.array(img)
    kernel = np.ones((3,3))
    edges = img_np - binary_erosion(img_np, kernel)*255
    return edges

# 加载图像并提取边缘
a = Image.open('a.png')
edges = find_edges(a)
edges_np = np.array(edges)

# 提取所有边缘白像素的坐标（图像数组shape为(高度,宽度)，坐标格式为(y, x)）
edge_points = np.argwhere(edges_np > 0)  # 输出shape: (N, 2)，每行对应一个边缘像素的(y,x)

# 生成图像所有像素的坐标网格
y_coords, x_coords = np.meshgrid(np.arange(edges_np.shape[0]), np.arange(edges_np.shape[1]), indexing='ij')
# 将网格展开为一维坐标数组，shape变为(总像素数, 2)
all_points = np.stack([y_coords.ravel(), x_coords.ravel()], axis=1)

# 用广播批量计算距离：扩展维度后做差、平方、求和、开根号
distances = np.sqrt(np.sum((all_points[:, np.newaxis, :] - edge_points[np.newaxis, :, :])**2, axis=2))

# 取每个像素到最近边缘的最小距离
min_distances = np.min(distances, axis=1)

# 将结果还原为原图像的形状
min_distances_img = min_distances.reshape(edges_np.shape)

# 可选：将距离值归一化到0-255，保存为可视化图像
normalized_dist = ((min_distances_img - min_distances_img.min()) / (min_distances_img.max() - min_distances_img.min()) * 255).astype(np.uint8)
Image.fromarray(normalized_dist, 'L').save('distance_map.png')

print(min_distances)

核心优化说明 #

• 批量提取边缘坐标：用np.argwhere一次性获取所有边缘像素的位置，替代循环中逐个判断。
• 广播批量计算：通过np.newaxis扩展数组维度，让所有像素坐标和边缘坐标做批量运算，充分利用Numpy的C级优化，比Python循环快几十到上百倍。

更高效的替代方案（结合Scipy） #

既然你的代码已经导入了scipy.ndimage，推荐直接用distance_transform_edt函数——这是专门为这类距离计算场景设计的工具，效率比纯Numpy实现还要高，代码也更简洁：

from scipy.ndimage import distance_transform_edt

# 反转边缘图像：把原边缘白像素设为0，其余黑像素设为1
inverted_edges = (edges_np == 0).astype(np.uint8)
# 调用距离变换，直接得到每个像素到最近原边缘白像素的距离
min_distances_img = distance_transform_edt(inverted_edges)

这个函数底层做了高度优化，大图像上的性能优势会非常明显。

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