python 多进程实现快速批量图像数据增强

python多进程实现快速批量图像数据增强（旋转及翻转）

在计算机视觉中，经常要用到数据增强（旋转、翻转、加噪声、改变亮度等）等操作。因此，我编写了一个简单实用的python程序，以实现对图片的旋转(90,180,270)以及翻转操作。单进程的运算速度对于处理分辨率较高的图像速度太慢，因此采用多进程加速的方法提高运算速度。

使用之前安装：pillow和tqdm包，python=3.+
将代码保存在txt文件中，在终端运行该txt文件
并输入图片路径，以及是否需要旋转、翻转操作
程序运行结束后，会自动在输入路径下新建一个文件夹(./augment)保存数据增强后的图片

# coding: utf-8

from PIL import Image
import os
import os.path
import copy
from tqdm import tqdm
import time
from datetime import datetime
import multiprocessing
import math


class Pathroot():

    def __init__(self, root, sort_methods):
        ''' 初始化文件路径 '''
        self.root = root
        self.sort_methods = sort_methods

    def get_path(self):
        ''' 输入路径，获取输入路径下所有文件的完整列表,并以字符串排序 '''
        path1 = []
        for parent, dirnames, filenames in os.walk(self.root):
            for filename in filenames:
                currentPath = os.path.join(parent, filename)
                path1.append(currentPath)
        if self.sort_methods == 'str':
            path1.sort(key=lambda x: str(os.path.basename(x).split('/')[-1][:-4]))
        elif self.sort_methods == 'int':
            path1.sort(key=lambda x: int(os.path.basename(x).split('/')[-1][:-4]))
        else:
            print('Please choose correct sort methods!!! str or int')
        return path1


def rotate_and_flip(old_root, new_root, rotato_info, flip_info):
    ''' 对图片进行旋转90，180，270，以及翻转 '''
    img = Image.open(old_root)
    if rotato_info == '1':
        for i in [0, 90, 180, 270]:
            img1 = img.rotate(i)
            new_root_1 = [new_root, '/rotato_', str(i), '_', str(old_root.split('/')[-1])]
            new_root_1 = ''.join(new_root_1)
            img1.save(new_root_1)
    if flip_info == '1':
        for i in [0, 90, 180, 270]:
            img2 = img.rotate(i).transpose(Image.FLIP_LEFT_RIGHT)
            new_root_2 = [new_root, '/flip_', str(i) + '_', str(old_root.split('/')[-1])]
            new_root_2 = ''.join(new_root_2)
            img2.save(new_root_2)


if __name__ == '__main__':

    # 打印当前时间
    starttime = datetime.now()
    print('\n', '--' * 10, 'Begin_time:  ', time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()), '--' * 10, '\n')

    # 输入图片路径、以及数据增强的方式
    input_path = input('Please enter the path to the folder :  ')
    rotato_info = input('\nWhether the image rotation is required: 1 or 0 (1 = YES ,0 = NO)  ')
    flip_info = input('\nWhether the image should be flipped: 1 or 0 (1 = YES ,0 = NO)  ')

    # 获取输入路径下全部图片的列表
    path_root = Pathroot(input_path, 'str')
    the_total_root = path_root.get_path()
    aug_root = '/'.join(input_path.split('/')[:-1]) + '/augment'
    os.mkdir(aug_root)
    print('\nThe number of images under the path : {:^10}'.format(str(len(the_total_root))))

    # 创建多进程
    pool = multiprocessing.Pool(multiprocessing.cpu_count())

    # 批量进行图片旋转及翻转
    with tqdm(total=len(the_total_root)) as pbar:
        for i in the_total_root:
            pool.apply_async(rotate_and_flip, (i, aug_root, rotato_info, flip_info,), callback=lambda _: pbar.update(1),
                             error_callback=lambda _: pbar.update(1), )
        pool.close()
        pool.join()

    # 程序运行结束，打印结束时间
    endtime = datetime.now()
    print('\nThe program has ended, please check the results！！！')
    list = [i for i in str((endtime - starttime)).split(':')]
    print("\nRunTime: {} h --{} m --{} s".format(list[0], list[1], list[2]))
    print('\n', '--' * 10, 'Stop_time:  ', time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()), '--' * 10, '\n')

终端程序运行结果如下：

如果需要对文件重命名的话，可以使用下面链接中的重命名代码!!!

https://blog.csdn.net/wh1213148/article/details/113872202