如何实现大量图像增强操作的Python多进程并行处理

问题根因梳理 #

• 第一次多线程实现效率无提升的核心问题：你给pool.map()传入的待处理任务列表只有1个元素[[all_folders,bpath]]，相当于4个线程里只有1个线程拿到了所有文件夹的处理任务，其余线程全程空闲，自然和串行效率一致。
• 第二次方案报错的核心问题：
  1. os.walk()返回的是三元组迭代器（根路径、子目录列表、文件列表），不是直接的文件夹路径列表
  2. chain(paths.list_images(folder) for folder in all_folders)返回的是生成器对象，multiprocessing进程间通信需要序列化（pickle）任务对象，生成器无法被序列化，所以抛出TypeError: cannot pickle 'generator' object错误。

可行实现方案 #

方案1：按文件夹维度并行（适合单文件夹内图片数量差异不大的场景） #

import os
import cv2
from multiprocessing import Pool
from imutils import paths

# 替换成你自己的图像增强逻辑
def augmentations(img):
    # 增强操作示例，可自行修改
    return cv2.flip(img, 1)

def do_augmentations(folder_path):
    # 单个进程处理单个文件夹下的所有图片
    all_imgs = list(paths.list_images(folder_path))
    for img_path in all_imgs:
        img = cv2.imread(img_path)
        aug_img = augmentations(img)
        cv2.imwrite(img_path, aug_img)

if __name__ == "__main__":
    root_folder = "替换为你的图像根目录路径"
    # 提前遍历拿到所有子文件夹的绝对路径，转成列表避免迭代器问题
    all_folders = []
    for root, dirs, files in os.walk(root_folder):
        for dir_name in dirs:
            all_folders.append(os.path.join(root, dir_name))
    
    # 进程数根据你的CPU核心数设置，通常和物理核心数一致即可
    pool = Pool(processes=4)
    # 把每个文件夹作为单独任务分配给不同进程
    pool.map(do_augmentations, all_folders)
    pool.close()
    pool.join()

方案2：按图片维度并行（适合图片总量大、单文件夹图片数量不均的场景） #

import os
import cv2
from itertools import chain
from multiprocessing import Pool
from imutils import paths

# 替换成你自己的图像增强逻辑
def augmentations(img):
    # 增强操作示例，可自行修改
    return cv2.flip(img, 1)

def augment_single_img(img_path):
    # 单个进程处理单张图片
    img = cv2.imread(img_path)
    aug_img = augmentations(img)
    cv2.imwrite(img_path, aug_img)

if __name__ == "__main__":
    root_folder = "替换为你的图像根目录路径"
    # 提前拿到所有图片的绝对路径，转成列表避免生成器无法序列化的问题
    all_folders = [os.path.join(root, d) for root, dirs, f in os.walk(root_folder) for d in dirs]
    all_imgs = list(chain.from_iterable(paths.list_images(f) for f in all_folders))
    
    pool = Pool(processes=4)
    # 单张图片处理耗时很短的话可以加chunksize参数，减少进程通信开销，例如chunksize=10
    pool.map(augment_single_img, all_imgs)
    pool.close()
    pool.join()

注意事项 #

• 必须加if __name__ == "__main__"入口保护，Windows系统下multiprocessing必须要有这个声明否则会报错
• 图像增强属于CPU密集型任务，不要用multiprocessing.dummy.ThreadPool（多线程），Python的GIL会限制CPU密集型任务的多线程效率，直接用multiprocessing.Pool多进程才能真正利用多核CPU

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