如何加速基于字典映射的像素替换？含语义分割预处理问题

任务背景 #

我正在开展语义分割所需的.png掩码文件像素替换任务，文件尺寸为1208x1920x3。已构建原始像素与类别编号的映射字典：replacement_dic = {(xxx,yyy,zzz):(cl1,cl1,cl1),....}。

现有逐像素替换代码如下，但单文件处理耗时达20秒（Intel® Core™ i7-12700H处理器），速度过慢：

def replace_pixels(path, replacement_dict):
    label = cv2.cvtColor(cv2.imread(path), cv2.COLOR_BGR2RGB)
    height, width, _ = label.shape
    for i in range(height):
        for j in range(width):
            pixel = tuple(label[i,j])    
            if pixel in replacement_dic:
                 label[i,j] = replacement_dic[pixel]
            else:
                label[i,j] = [0,0,0]
    return label

for i in range(len(all_labels)):
    new_label = replace_pixels(all_labels[i], replacement_dic)
    cv2.imwrite(all_labels[i], new_label)

注：原代码存在拼写错误，已修正COLOR_BRG2RGB为COLOR_BGR2RGB、hight为height、in为j、renge为range、replace_pixeles为replace_pixels

问题1：有哪些加速该处理的方法？ #

• 用NumPy向量化操作替代嵌套循环：逐像素遍历是性能核心瓶颈，可将RGB像素转换为唯一整数索引，构建查找表后通过向量化操作完成批量替换，示例代码：

  import numpy as np
  import cv2
  
  def replace_pixels_fast(path, replacement_dict):
      label = cv2.cvtColor(cv2.imread(path), cv2.COLOR_BGR2RGB)
      # 将RGB三通道合并为单个24位整数
      rgb_int = label[...,0] * 65536 + label[...,1] * 256 + label[...,2]
      # 初始化查找表，默认值对应[0,0,0]的整数
      max_val = 2**24 - 1
      lut = np.zeros(max_val + 1, dtype=np.uint32)
      # 填充查找表：原像素整数映射到目标像素整数
      for rgb_tuple, target_tuple in replacement_dict.items():
          rgb_key = rgb_tuple[0] * 65536 + rgb_tuple[1] * 256 + rgb_tuple[2]
          target_val = target_tuple[0] * 65536 + target_tuple[1] * 256 + target_tuple[2]
          lut[rgb_key] = target_val
      # 批量应用查找表
      result_int = lut[rgb_int]
      # 将整数还原为RGB三通道
      result = np.stack([(result_int >> 16) & 0xFF, 
                         (result_int >> 8) & 0xFF, 
                         result_int & 0xFF], axis=-1)
      return result
  

• 多进程批量处理：用multiprocessing或concurrent.futures模块，同时处理多个文件，充分利用CPU多核性能。
• 转换为单通道类别图：既然目标是类别编号（目标像素为(cl1,cl1,cl1)），可直接生成单通道的类别索引图（每个像素值为cl1），既提升处理速度，也简化后续CNN预处理流程。
• 用OpenCV内置优化接口：OpenCV底层为C++实现并做了SIMD优化，优先使用其数组操作替代纯Python逻辑。

问题2：未匹配字典时的替换值[0,0,0]格式是否合理，是否需改为字典中使用的元组格式(0,0,0)？ #

格式合理，无需修改。label是OpenCV读取的NumPy数组，元素类型为uint8，赋值时列表[0,0,0]会自动转换为符合数组要求的格式，和元组(0,0,0)效果完全一致。如果担心兼容性，用元组也可以，两者在NumPy数组赋值场景下没有区别。

问题3：CNN预处理流水线中实现该功能的方案建议 #

TensorFlow流水线实现 #

1. 构建TensorFlow哈希查找表
   先将replacement_dic中的RGB键转换为整数，提取目标值的类别编号（因目标是(cl1,cl1,cl1)，直接取cl1即可），构建哈希表：

   import tensorflow as tf
   
   def build_hash_table(replacement_dict):
       keys = []
       values = []
       for rgb_tuple, target_tuple in replacement_dict.items():
           rgb_int = rgb_tuple[0] * 65536 + rgb_tuple[1] * 256 + rgb_tuple[2]
           keys.append(rgb_int)
           values.append(target_tuple[0])  # 提取类别编号
       keys_tensor = tf.constant(keys, dtype=tf.int32)
       values_tensor = tf.constant(values, dtype=tf.int32)
       return tf.lookup.StaticHashTable(
           tf.lookup.KeyValueTensorInitializer(keys_tensor, values_tensor),
           default_value=tf.constant(0, dtype=tf.int32)
       )
   
   hash_table = build_hash_table(replacement_dic)
   

2. 编写预处理函数并构建流水线

   def preprocess_mask(mask_path):
       # 读取并解码掩码为RGB格式
       mask = tf.io.read_file(mask_path)
       mask = tf.image.decode_png(mask, channels=3)
       mask = tf.cast(mask, tf.int32)
       # 将RGB转换为整数索引
       rgb_int = mask[...,0] * 65536 + mask[...,1] * 256 + mask[...,2]
       # 查询哈希表得到单通道类别图
       class_mask = hash_table.lookup(rgb_int)
       # 若需还原为三通道掩码，可执行以下操作
       # class_mask = tf.expand_dims(class_mask, axis=-1)
       # class_mask = tf.repeat(class_mask, 3, axis=-1)
       return class_mask
   
   # 构建数据集流水线，开启多线程加速
   mask_dataset = tf.data.Dataset.from_tensor_slices(all_labels)
   mask_dataset = mask_dataset.map(preprocess_mask, num_parallel_calls=tf.data.AUTOTUNE)
   

3. 关键细节处理
   • 确保哈希表的键、值类型与张量匹配，避免类型转换错误；
   • 训练流水线可结合shuffle、batch等操作，同时开启prefetch进一步提升效率；
   • 若需保存处理后的掩码，可在流水线中添加tf.io.write_file，将张量编码为PNG格式。

替代方案 #

• 离线预处理：若数据集规模不大，先用问题1中的加速方法批量处理所有掩码，保存为单通道类别图或处理后的三通道掩码，训练时直接读取处理后的文件，避免重复预处理；
• PyTorch流水线：类似TensorFlow，可利用torch.nn.Embedding或自定义哈希映射实现向量化替换，结合DataLoader的多进程加载功能加速。

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