数据增强 imgaug 框架使用

Posted 2020-11-27 clemente

官方文档

---

https://imgaug.readthedocs.io/en/latest/source/overview/arithmetic.html#cutout

 

首先下载项目

https://github.com/aleju/imgaug

然后建立conda独立环境

conda install six numpy scipy Pillow matplotlib scikit-image imageio

conda install -c menpo opencv

 

 

代码

注意 待增强的图片需要分别放置于/abc/下单独的文件夹，以保证每个文件夹都是相关增强图片

# -*- coding: utf-8 -*-
"""
图片增强
一个目录下有多个类别

"""

import numpy as np
import imgaug as ia
import imgaug.augmenters as iaa
import matplotlib.pyplot as plt
import os
from skimage import io
import random
import time


#根目录，其中是不同类别的子目录，子目录中是图片
data_dir = r"/home/joe/Documents/abc/" 

total_count = 4 #所有子目录下文件计数
for (rootpath, dirnames, filenames) in os.walk(data_dir):
    for dirnametmp in dirnames: 
        dirnametmp = data_dir+dirnametmp+"/"  #每个子文件夹
        for (rootpath, dirnames, filenames) in os.walk(dirnametmp):
            filenamescount = len(filenames)
            imgcounttmp = 1 #当前子目录下的文件计数
            for icount in range(filenamescount): 
                imgname = filenames.pop()   #每个图片
             
                #qqq，有imgaug的表示造的图片
                file_extension = os.path.splitext(imgname)
                filetmp = file_extension[0]    #文件名不带扩展名
                extension = file_extension[1]  #扩展名
                if filetmp.endswith("_imgaug"):
                    continue
                
                imgfullnametmp = dirnametmp+imgname #每个图片全路径
                print( f"{imgcounttmp}/{total_count} {imgfullnametmp}")
                imgcounttmp += 1
                total_count += 1 
                image = io.imread(imgfullnametmp)
                
                saveimage = []
                random1 = random.uniform(-5,5)
                random2 = random.uniform(0,1)
                random3 = random.randint(1,5)
                 
                augtmp = iaa.Affine(rotate=(-25+random1, 25.5+random1))
                res = augtmp.augment_image(image)
                saveimage.append(res)
                
                augtmp = iaa.Fliplr(1.0)
                res = augtmp.augment_image(image)
                saveimage.append(res)
                
                augtmp = iaa.Crop(px=(0, random3*2), keep_size=True)
                res = augtmp.augment_image(image)
                saveimage.append(res)
                
                augtmp = iaa.Sometimes(random2,iaa.GaussianBlur(sigma=(random2, 0.5+random2)))
                res = augtmp.augment_image(image)
                saveimage.append(res)
                
                augtmp = iaa.ContrastNormalization((0.75+random2, 1.5+random2))
                res = augtmp.augment_image(image)
                saveimage.append(res)
                
                augtmp = iaa.AdditiveGaussianNoise(loc=random3, scale=(0.01+random2, 0.08*255+random2))
                res = augtmp.augment_image(image)
                saveimage.append(res)
                
                augtmp = iaa.Affine(
                    scale={"x": (0.8+random2, 1.2+random2), "y": (0.8+random2, 1.2+random2)}, #缩放
                    translate_percent={"x": (-0.2+random2, 0.2+random2), "y": (-0.2+random2, 0.2+random2)}, #平移 
                    rotate=(-20+random1, 60+random1),                        #旋转
                    shear=(-8+random1, 8+random1) )
                res = augtmp.augment_image(image)
                saveimage.append(res)
                
                augtmp = iaa.GaussianBlur(sigma=random2)
                res = augtmp.augment_image(image)
                saveimage.append(res) 
                
                augtmp= iaa.AverageBlur(k=random3)
                res = augtmp.augment_image(image)
                saveimage.append(res) 
                
                augtmp = iaa.MedianBlur(k=random3*2+1)
                res = augtmp.augment_image(image)
                saveimage.append(res) 
                
                augtmp = iaa.MotionBlur()
                res = augtmp.augment_image(image)
                saveimage.append(res) 
                
                augtmp = iaa.BilateralBlur()
                res = augtmp.augment_image(image)
                saveimage.append(res) 
                
                augtmp = iaa.WithColorspace(to_colorspace="HSV")
                res = augtmp.augment_image(image)
                saveimage.append(res) 
                
                augtmp = iaa.AddToHueAndSaturation((-20, 20), per_channel=True)
                res = augtmp.augment_image(image)
                saveimage.append(res) 
                
                augtmp = iaa.Grayscale(alpha=1.0)
                res = augtmp.augment_image(image)
                saveimage.append(res) 
                
                augtmp = iaa.CLAHE(clip_limit=(1, random3*10))
                res = augtmp.augment_image(image)
                saveimage.append(res) 
                
                augtmp = iaa.GammaContrast(gamma=2*random2+1)
                res = augtmp.augment_image(image)
                saveimage.append(res) 
                
                augtmp = iaa.SigmoidContrast(gain= random2+0.5)
                res = augtmp.augment_image(image)
                saveimage.append(res) 
                
                augtmp = iaa.LogContrast(gain=2*random2+1)
                res = augtmp.augment_image(image)
                saveimage.append(res) 
                
                augtmp = iaa.LinearContrast(alpha=random3)
                res = augtmp.augment_image(image)
                saveimage.append(res) 
                
                augtmp = iaa.LinearContrast(alpha=random3*2)
                res = augtmp.augment_image(image)
                saveimage.append(res) 
                
                aug = iaa.CLAHE(clip_limit=(1, random3*20))
                res = augtmp.augment_image(image)
                saveimage.append(res) 
                
                #保存图片
                imgid = 0
                t = time.time()
                tt = round(t*1000)
                for imgtmp in saveimage:
                    #原文件名带有扩展名，需要拆分开
                    #文件名加时间戳,imgaug的表示造的图片
                    newname = filetmp+"_" +str(tt)+str(imgid)+"_imgaug"+extension
                    newfullname = dirnametmp+newname
                    #ia.imshow(imgtmp)

                    io.imsave(newfullname,imgtmp)
                    imgid += 1