如何优化灰度图转自定义RGBA调色板的Python代码性能？

核心优化思路：规避逐像素遍历，用矢量化+并行处理突破性能瓶颈 #

1. Numpy矢量化映射替代逐像素操作 #

逐像素查询字典是性能短板，预先生成全灰度范围的RGBA映射数组，通过Numpy矢量化索引直接完成转换，底层基于C实现，速度提升显著。

import numpy as np
from PIL import Image

# 替换为你的自定义灰度-RGB调色板
custom_palette = {
    0: (0, 0, 0),
    1: (255, 0, 0),
    50: (0, 255, 0),
    255: (0, 0, 255)
}

# 预先生成灰度值到RGBA的查找数组（索引对应灰度值）
max_gray = 255
rgba_lookup = np.zeros((max_gray + 1, 4), dtype=np.uint8)
# 默认设置alpha为255（不透明）
rgba_lookup[:, 3] = 255
# 填充调色板对应的RGB值
for gray_val, rgb in custom_palette.items():
    rgba_lookup[gray_val, :3] = rgb
# 灰度值0时设置alpha为0（完全透明）
rgba_lookup[0, 3] = 0

def convert_single(img_path):
    # 读取灰度图并转为Numpy数组
    gray_arr = np.array(Image.open(img_path).convert('L'))
    # 矢量化映射直接生成RGBA数组
    rgba_arr = rgba_lookup[gray_arr]
    # 转回PIL Image对象
    return Image.fromarray(rgba_arr, mode='RGBA')

2. 多进程批量处理 #

大量图像转换属于CPU密集型任务，用多进程充分利用多核CPU，整体处理效率可提升数倍。

from multiprocessing import Pool
import os

def process_and_save(img_path, output_dir):
    rgba_img = convert_single(img_path)
    filename = os.path.basename(img_path).replace('.png', '_rgba.png')
    rgba_img.save(os.path.join(output_dir, filename))

if __name__ == '__main__':
    input_dir = './gray_images'
    output_dir = './rgba_images'
    os.makedirs(output_dir, exist_ok=True)
    # 获取所有灰度图路径
    img_paths = [os.path.join(input_dir, f) for f in os.listdir(input_dir) if f.endswith('.png')]
    
    # 用CPU核心数创建进程池
    with Pool(processes=os.cpu_count()) as pool:
        pool.starmap(process_and_save, [(path, output_dir) for path in img_paths])

3. OpenCV加速IO与转换 #

OpenCV的图像读取和数组操作底层优化更彻底，比PIL更快，适合大规模图像处理：

import cv2

def convert_single_cv(img_path):
    # 直接读取灰度图为Numpy数组
    gray_arr = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
    # 矢量化映射生成RGBA数组
    rgba_arr = rgba_lookup[gray_arr]
    # OpenCV保存RGBA需要转为BGRA格式
    bgra_arr = rgba_arr[:, :, [2, 1, 0, 3]]
    cv2.imwrite(img_path.replace('.png', '_rgba.png'), bgra_arr)

4. 内存预加载（可选） #

如果内存足够，一次性加载所有灰度图到内存，减少磁盘IO等待时间：

# 预加载所有灰度图到内存
gray_imgs = {
    f: np.array(Image.open(os.path.join(input_dir, f)).convert('L'))
    for f in os.listdir(input_dir) if f.endswith('.png')
}

# 批量转换
for filename, gray_arr in gray_imgs.items():
    rgba_arr = rgba_lookup[gray_arr]
    Image.fromarray(rgba_arr, mode='RGBA').save(os.path.join(output_dir, filename.replace('.png', '_rgba.png')))

效果预期 #

单张转换时间可从0.3秒压缩到0.01秒以内，批量处理效率随CPU核心数线性提升。

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