如何向量化Numpy代码中含numpy.where的for循环以提升性能？

原代码的性能瓶颈在于对每个标签循环提取对应像素并计算众数，当标签数量达数千时，Python循环的开销会被放大。下面是完全基于Numpy的向量化实现，彻底消除Python循环：

import numpy as np
from skimage.segmentation import expand_labels  # 假设expand_labels来自skimage

def img_4c_d(lab, img):
    seg_islands = np.where(seg_4col == lab, seg, 0)
    seg_exp = expand_labels(seg_islands, distance=max(img.shape[0], img.shape[1]))

    # 提取非零标签及其对应的像素（0是无效标签，跳过）
    valid_mask = seg_islands != 0
    labels = seg_islands[valid_mask]
    if not np.any(labels):
        # 没有有效标签时直接返回全0数组
        return np.zeros_like(img)
    
    max_label = labels.max()
    img_pixels = img[valid_mask]

    # 定义向量化计算每个标签对应通道众数的函数
    def get_channel_mode(labels, channel_vals, max_label):
        # 创建计数矩阵：行=标签，列=通道值(0-255)
        count_matrix = np.zeros((max_label + 1, 256), dtype=np.int64)
        # 对每个(标签,通道值)位置累加计数
        np.add.at(count_matrix, (labels, channel_vals), 1)
        # 找到每个标签下计数最大的通道值（众数）
        return count_matrix.argmax(axis=1)
    
    # 分别计算RGB三个通道的众数
    mode_r = get_channel_mode(labels, img_pixels[..., 0], max_label)
    mode_g = get_channel_mode(labels, img_pixels[..., 1], max_label)
    mode_b = get_channel_mode(labels, img_pixels[..., 2], max_label)

    # 构建颜色字典并映射到扩展后的标签图
    color_dic = np.stack([mode_r, mode_g, mode_b], axis=1).astype(np.uint8)
    return color_dic[seg_exp]

关键优化点说明 #

• 跳过无效标签：通过valid_mask直接提取所有非零标签的像素，避免循环中反复筛选seg_islands == label
• 向量化计数：用np.add.at一次性完成所有(标签,通道值)组合的计数，替代循环中的np.bincount，底层C实现的批量操作远快于Python循环
• 批量计算众数：对计数矩阵按行取argmax，直接得到每个标签对应的通道众数，无需逐个处理标签

额外性能建议 #

如果seg_islands中的标签是连续的（没有断层），可以省略max_label的计算，直接用seg.max()提前获取最大标签值；如果标签存在大量断层，也可以用np.unique获取唯一标签后构建更小的计数矩阵，但np.add.at对稀疏标签的处理效率依然远高于Python循环。

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