图像数据增强的高效执行

一、简介

图像数据增强是一种常用的扩充训练图片数据量，从而提升模型性能的方法，一般常用的图像数据增强方法主要有随机翻转、扩大缩小图像、随机调节亮度对比度、添加椒盐噪声等。

当原始图像数据集的数量为几千张时，常规的方法处理速度就可以满足要求，但是当图像数据集的数量在几万张或者几十万张时，一般的图像处理方法耗时非常长，这时就需要使用能够高效执行的代码来扩充数据集，这也是本文的主要内容。

下面以添加椒盐噪声为例，逐步提高代码运行效率。

二、初始代码

2.1代码

添加椒盐噪声代码如下：

def add_saltpepperNoise(imgpath, prob = 0.08):image = Image.open(imgpath).convert('RGB')image_data = np.array(image)height, width = image_data.shape[:2]num_pixels = int(height * width * prob)for _ in range(num_pixels):y = random.randint(0, height - 1)x = random.randint(0, width - 1)image_data[y, x] = [0, 0, 0] if random.random() < 0.5 else [255, 255, 255]noisy_image = Image.fromarray(image_data)return noisy_image

2.2代码说明

上述代码执行逻辑如下：

1. 根据prob计算出添加噪声的像素数，0.08表示8%像素被替换；
2. 逐像素循环，对每一个像素，通过random方法生成随机位置，同样也通过random方法决定添加黑点或白点；
3. 循环结束后保存图像。

上述代码较为简单，但执行效率较低，经测试处理1万张图片需耗时约7分钟。

三、改进代码1.0

3.1代码

初始代码效率低的主要原因是逐像素循环和random的调用，这一过程可以使用numpy的random函数替代运算，代码如下：

def add_saltpepperNoise_pro(imgpath, prob=0.08):image = Image.open(imgpath).convert('RGB')image_data = np.array(image)random_matrix = np.random.rand(*image_data.shape[:2])pepper_mask = random_matrix < (prob / 2)salt_mask = random_matrix > (1 - prob / 2)image_data[pepper_mask] = [0, 0, 0]image_data[salt_mask] = [255, 255, 255]noisy_image = Image.fromarray(image_data)return noisy_image

3.2代码说明

上述代码改进点如下：

1. 去掉循环，用 np.random.rand() 一次生成所有像素位置的随机数，去掉了random.randint调用；
2. 直接对选中的像素批量赋值。

上述代码执行效率有所提升，经测试处理1万张图片需耗时约1分钟。

四、改进代码2.0

4.1代码

可以将上一节中的改进代码1.0转移到显卡上进行操作，从而提高效率，代码如下：

def add_saltpepperNoise_gpu(imgpath, prob=0.08, device='cuda'):image = Image.open(imgpath).convert('RGB')transform_to_tensor = transforms.ToTensor()transform_to_pil = transforms.ToPILImage()image_tensor = transform_to_tensor(image).to(device)random_matrix = torch.rand(image_tensor.shape[1:], device=image_tensor.device)pepper_mask = random_matrix < (prob / 2)salt_mask = random_matrix > (1 - prob / 2)noisy_image = image_tensor.clone()noisy_image[:, pepper_mask] = 0.0noisy_image[:, salt_mask] = 1.0noisy_image = transform_to_pil(noisy_image.cpu())return noisy_image

4.2代码说明

上述代码不再依赖numpy，全部用pytorch tensor运算，用布尔掩码直接替换像素，避免循环。

但实际执行上述代码发现，效率对比改进代码1.0反而下降，处理1万张图片需耗时约2.5分钟。效率不如改进代码1.0的原因有可能是因为CPU-GPU之间的I/O往返增加了代码执行时间。

五、改进代码3.0

5.1代码

可以在改进代码1.0的基础上，使用CPU多核并行的方式提高代码效率，代码如下：

def addnoise(args):imgpath, savepath = argsimage = Image.open(imgpath).convert('RGB')nimg = add_saltpepperNoise_pro(image)nimg.save(savepath)def main_process():with Pool(10) as p:for _ in tqdm(p.imap_unordered(addnoise, args_list), total=len(args_list), desc="Processing"):pass

5.2代码说明

上述代码的核心点在于用multiprocessing.Pool并行执行，处理1万张图片需耗时约0.4分钟，性能有了很大提升，代码注意事项如下：

1. with Pool(10) as p中的10为10个并行线程，应根据计算机CPU核数适当调整，该值越大虽然速度越快，但占用CPU也越高；
2. for _ in tqdm(p.imap_unordered(addnoise, args_list)中的args_list，为包含了多个(imgpath, savepath)的列表，在执行代码前需要进行构造。

六、总结

添加椒盐噪声代码处理1万张图片运行时间汇总如下表所示：

代码	处理时间	相对初始代码速度提升
初始代码	7分钟	——
改进代码1.0	1分钟	7倍
改进代码2.0	2.5分钟	2.8倍
改进代码3.0	0.4分钟	17.5倍

可以看到经过改进，代码效率有了非常明显的提升。