imgaug批量椒盐噪声 python_python图像扩增-imgaug

# encoding=utf-8

import os

import imageio

import random

import numpy as np

import imgaug as ia

import concurrent.futures

from imgaug import augmenters as iaa

import matplotlib.pyplot as plt

from argparse import ArgumentParser

ia.seed(4)

# %matplotlib inline

# 参考 https://github.com/aleju/imgaug

func_norm = iaa.ContrastNormalization(alpha=(0.5, 1.5), per_channel=False)

func_add = iaa.Add(value=(-30, 40), per_channel=False)

func_hue_saturation = iaa.AddToHueAndSaturation(

value=(-35, 30), per_channel=False)

func_multiply = iaa.Multiply(mul=(0.8, 1.5), per_channel=False)

func_gamma = iaa.GammaContrast(gamma=(0.8, 1.5), per_channel=False)

func_log = iaa.LogContrast(gain=(0.7, 1.2), per_channel=False)

func_sigmoid = iaa.SigmoidContrast(gain=(2, 5), cutoff=0.5, per_channel=False)

func_coarse = iaa.CoarseDropout(

p=(0.1, 0.2), size_percent=0.1, per_channel=False)

switch_func = {

0: func_norm,

1: func_add,

2: func_hue_saturation,

3: func_multiply,

4: func_gamma,

5: func_log,

6: func_sigmoid,

7: func_coarse

}

def getFilePath(root_path, file_list):

dir_or_files = os.listdir(root_path)

for dir_file in dir_or_files:

dir_file_path = os.path.join(root_path, dir_file)

if os.path.isdir(dir_file_path):

getFilePath(dir_file_path, file_list)

else:

file_list.append(dir_file_path)

def getFilePath2(root_path, file_list, folder_list):

dir_or_files = os.listdir(root_path)

for dir_file in dir_or_files:

dir_file_path = os.path.join(root_path, dir_file)

if os.path.isdir(dir_file_path):

folder_list.append(dir_file_path)

getFilePath2(dir_file_path, file_list, folder_list)

else:

file_list.append(dir_file_path)

def getFilePathFilters(root_path, file_list, dir_filters):

dir_or_files = os.listdir(root_path)

for dir_file in dir_or_files:

dir_file_path = os.path.join(root_path, dir_file)

if os.path.isdir(dir_file_path):

if dir_file in dir_filters:

getFilePathFilters(dir_file_path, file_list, dir_filters)

else:

file_list.append(dir_file_path)

def makeDir(root_path):

if not os.path.exists(root_path):

os.makedirs(root_path)

def showImage(image):

image = np.array(image)

plt.imshow(image)

plt.show()

def saveImage(image, write_path):

imageio.imwrite(write_path, image)

def augmentFunc(image, index):

return switch_func[index].augment_image(image)

# 多线程执行函数

def augExecute(file_path):

new_file_path = file_path.replace(args.root_folder, args.new_folder)

print(new_file_path, end='\n')

image = imageio.imread(file_path)

index = random.randint(0, 7)

new_image = augmentFunc(image, index)

# showImage(new_image)

saveImage(new_image, new_file_path)

def colorImageAugment(augment_list, max_workers, dir_dict):

for augment_folder in augment_list:

file_list = []

folder_list = []

getFilePath2(augment_folder, file_list, folder_list)

# 创建新目录文件夹及其子文件夹

for folder_path in folder_list:

new_folder_path = folder_path.replace(

args.root_folder, args.new_folder)

makeDir(new_folder_path)

if max_workers == 0:

# 单线程

for file_path in file_list:

augExecute(file_path)

else:

# 多线程

with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:

for _ in executor.map(augExecute, file_list):

pass

def main(args):

aug_folder_list = args.image_folder_path

dir_dict = {'root_dir': args.root_folder, 'trans_dir': args.new_folder}

max_workers = int(args.max_workers)

print("Start to augment image")

colorImageAugment(aug_folder_list, max_workers, dir_dict)

print("The end")

if __name__ == "__main__":

parser = ArgumentParser(description="Augment color image")

parser.add_argument(

"image_folder_path", nargs='+', help="Absolute path of image folder that going to augmentation")

parser.add_argument(

"root_folder", type=str, help="The folder and file under the root_folder is going to augmentation")

parser.add_argument("new_folder", type=str,

help="The folder that going to be build")

parser.add_argument("max_workers", type=int,

default=9, help="Multi thread num")

args = parser.parse_args()

main(args)

'''Usepython data_aug.py /your/path/to/augmentation/train/subdir1/ /your/path/to/augmentation/train/subdir2/ ... /your/path/to/augmentation/train/subdirn/ train train_aug第一个参数是 image_folder_path 需要处理的文件夹 要求绝对路径 可以传入多个第二个参数是 root_folder 处理文件夹的父级目录 程序会将该目录同级路径下创建 new_folder, 同时创建相同的字目录, 然后将 root_folder 下的图像处理保存到 new_folder 下相同自路径下第三个参数是 new_folder 见root_folder第四个参数是 max_workers 线程数量 0 为单线程'''

总结

以上是生活随笔为你收集整理的imgaug批量椒盐噪声 python_python图像扩增-imgaug的全部内容，希望文章能够帮你解决所遇到的问题。

如果觉得生活随笔网站内容还不错，欢迎将生活随笔推荐给好友。