# 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')
python进行批量数据增强,90,180,270度旋转
最新推荐文章于 2023-06-12 14:33:46 发布