如何在2D NumPy数组上叠加带累加值的圆形并加速运算

问题背景 #

需要在二维图像上加盖圆形印记，已知：

• 圆心坐标列表pt_list可能存在重复坐标点
• 不同圆形的覆盖区域可能存在重叠
  原有基于for循环逐点生成全图mask的实现运行速度极慢，期望通过np.meshgrid等内置向量化方法提速。

原有实现代码 #

def stamp_circle(img, pt_list,r,addup):
    x = np.arange(0,img.shape[1])
    y = np.arange(0,img.shape[0])
    for pt in pt_list:
        mask = (x[np.newaxis,:]-pt[1])**2 + (y[:,np.newaxis]-pt[0])**2 < r**2
        img[mask] += addup

参数说明 #

• img：尺寸为4096×8192的2D NumPy数组，为待处理的目标图像
• pt_list：形状为(1224,2)的圆心坐标数组
• r：圆形印记的半径，单位为像素
• addup：圆形覆盖区域像素需要叠加的亮度值

优化思路 #

直接用全图np.meshgrid和所有圆心做广播向量化不可行：全图共约3355万像素，和1224个圆心广播会生成大小为(1224, 4096, 8192)的中间数组，单精度内存占用超过150GB，普通硬件根本无法承载。
更高效的实现不需要做全图计算，核心优化点有三个：

• 先对重复圆心去重，统计每个坐标的出现次数，重复点只需要盖章1次，叠加值直接乘出现次数即可，消除冗余计算
• 同半径的圆形形状完全一致，提前用np.meshgrid生成单个圆形的局部mask，所有圆心复用，不需要每个点重复计算距离
• 放弃全图mask，只处理每个圆心周围边长为2r+1的局部窗口，计算量从「点数全图像素数」降到「点数(2r+1)²」，速度提升可达数百倍

优化后代码 #

import numpy as np

def stamp_circle_fast(img, pt_list, r, addup):
    H, W = img.shape
    r_int = int(np.ceil(r))
    r_sq = r ** 2
    # 坐标去重，统计重复点个数
    pts, counts = np.unique(pt_list.astype(np.int64), axis=0, return_counts=True)
    add_per_pt = counts * addup

    # 预生成单个圆形的局部mask
    dy, dx = np.meshgrid(
        np.arange(-r_int, r_int + 1),
        np.arange(-r_int, r_int + 1),
        indexing="ij"
    )
    circle_kernel = (dx ** 2 + dy ** 2) < r_sq

    # 逐点叠加局部区域
    for (cy, cx), val in zip(pts, add_per_pt):
        # 计算原图上的局部窗口范围，自动裁剪边界
        y_img_start = max(0, cy - r_int)
        y_img_end = min(H, cy + r_int + 1)
        x_img_start = max(0, cx - r_int)
        x_img_end = min(W, cx + r_int + 1)

        # 对应到圆形kernel上的裁剪范围
        y_k_start = r_int - (cy - y_img_start)
        y_k_end = r_int + (y_img_end - cy)
        x_k_start = r_int - (cx - x_img_start)
        x_k_end = r_int + (x_img_end - cx)

        # 叠加亮度值
        img[y_img_start:y_img_end, x_img_start:x_img_end] += (
            circle_kernel[y_k_start:y_k_end, x_k_start:x_k_end] * val
        )
    return img

补充说明 #

• 上述实现自动处理图像边界场景，不会出现数组索引越界问题
• 当圆形半径r<100时，该实现比原始for循环快100~1000倍，普通PC上处理4096×8192图像+1224个点仅需几十到几百毫秒
• 如果r取值极大（如r>1000），局部窗口方案优势下降，可以替换为FFT卷积实现：先构造全图点脉冲矩阵（对应位置填叠加值），再和圆形卷积核做快速傅里叶卷积，速度会进一步提升。

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