二维numpy数组多值批量重标注的向量化实现方案咨询

以下两种向量化方案都可以解决你遇到的性能问题，按适用场景选择即可：

1. 查找表法（性能最优，推荐90%以上场景使用） #

适合分割标签为非负整数、最大标签值不大的场景（绝大多数分割任务的掩码都符合这个条件）：

import numpy as np

# 构造标签查找表：索引为旧标签值，对应值为新标签值
max_old_label = image.max()
lookup_table = np.zeros(max_old_label + 1, dtype=image.dtype)
for old_val, new_val in relabels.items():
    lookup_table[old_val] = new_val

# 纯向量化索引替换，一次遍历整图完成重映射
new_image = lookup_table[image]

原理说明：该方案仅需遍历标签映射字典一次（时间消耗和标签数量强相关，和图像尺寸无关），后续重映射操作是Numpy原生向量化索引，总时间复杂度为O(N)，N为图像总像素数，比原方案快2~3个数量级。

如果需要保留未出现在映射字典中的原标签值，只需要将查找表的默认初始化改成lookup_table = np.arange(max_old_label + 1, dtype=image.dtype)即可。

2. 排序匹配法（适合稀疏大数值标签场景） #

如果你的标签值非常稀疏（比如标签值为1、10000、20000，最大值远大于标签总数），构造大尺寸查找表会浪费内存，可以用该方案：

import numpy as np

old_vals = np.array(list(relabels.keys()))
new_vals = np.array(list(relabels.values()))

# 对旧标签排序用于二分查找
sort_idx = np.argsort(old_vals)
old_sorted = old_vals[sort_idx]
new_sorted = new_vals[sort_idx]

# 二分查找匹配每个像素的标签位置
match_idx = np.searchsorted(old_sorted, image)
# 过滤未匹配到的标签
valid_mask = (match_idx < len(old_sorted)) & (old_sorted[match_idx] == image)
# 生成新图像，未匹配的标签默认置0
new_image = np.where(valid_mask, new_sorted[match_idx], 0)

原理说明：该方案时间复杂度为O(N log M)，N为图像总像素数，M为标签总数，性能也远高于原方案的O(N*M)复杂度。

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