本代码共采用了四种数据增强，如采用其他数据增强方式，可以参考本代码，随意替换。
'''imageDir 为原数据集的存放位置saveDir 为数据增强后数据的存放位置'''from PIL import Imagefrom PIL import ImageEnhanceimport osimport cv2import numpy as npimport globdef flip(root_path): #翻转图像img = Image.open(root_path)filp_img = img.transpose(Image.FLIP_LEFT_RIGHT)return filp_imgdef rotation(root_path):img = Image.open(root_path)rotation_img = img.rotate(90) #旋转角度return rotation_imgdef brightnessEnhancement(root_path):#亮度增强image = Image.open(root_path)enh_bri = ImageEnhance.Brightness(image)brightness = 1.5image_brightened = enh_bri.enhance(brightness)return image_brighteneddef contrastEnhancement(root_path): # 对比度增强image = Image.open(root_path)enh_con = ImageEnhance.Contrast(image)contrast = 1.5image_contrasted = enh_con.enhance(contrast)return image_contrasteddef colorEnhancement(root_path):#颜色增强image = Image.open(root_path)enh_col = ImageEnhance.Color(image)color = 1.5image_colored = enh_col.enhance(color)return image_coloreddef randomColor(root_path): #随机颜色"""对图像进行颜色抖动:param image: PIL的图像image:return: 有颜色色差的图像image"""image = Image.open(root_path)random_factor = np.random.randint(0, 31) / 10. # 随机因子color_image = ImageEnhance.Color(image).enhance(random_factor) # 调整图像的饱和度random_factor = np.random.randint(10, 21) / 10. # 随机因子brightness_image = ImageEnhance.Brightness(color_image).enhance(random_factor) # 调整图像的亮度random_factor = np.random.randint(10, 21) / 10. # 随机因子contrast_image = ImageEnhance.Contrast(brightness_image).enhance(random_factor) # 调整图像对比度random_factor = np.random.randint(0, 31) / 10. # 随机因子return ImageEnhance.Sharpness(contrast_image).enhance(random_factor) # 调整图像锐度if __name__ == '__main__':imageDir = "D:/untitled/数据增强/train/" # 要改变的图片的路径文件夹saveDir = 'D:/untitled/Month/train_aug/' # 要保存的图片的路径文件夹cate = [imageDir + x for x in os.listdir(imageDir)]print(cate)for idx, folder in enumerate(cate):print("idx = ",idx,"folder = ",folder)for i in glob.glob(folder + '/*.png'):print("i = ",i)saveName= i[:-4]+"fl.jpg"saveImage = flip(i)saveImage.save(os.path.join(saveDir,saveName))saveName = i[:-4] + "ro.jpg"saveImage = rotation(i)saveImage.save(os.path.join(saveDir, saveName))saveName = i[:-4] + "be.jpg"saveImage = brightnessEnhancement(i)saveImage.save(os.path.join(saveDir, saveName))saveName = i[:-4] + "co.jpg"saveImage = contrastEnhancement(i)saveImage.save(os.path.join(saveDir, saveName))saveName = i[:-4] + "cl.jpg"saveImage = colorEnhancement(i)saveImage.save(os.path.join(saveDir, saveName))saveName = i[:-4] + "ra.jpg"saveImage = randomColor(i)saveImage.save(os.path.join(saveDir, saveName))
如果多类别文件夹，如图像分类，需要统一进行数据增强，加入了文件夹遍历，可自取
from tkinter import image_namesfrom matplotlib.animation import ImageMagickBasedef flip(root_path,img_name): #翻转图像img = Image.open(os.path.join(root_path, img_name))filp_img = img.transpose(Image.FLIP_LEFT_RIGHT)# filp_img.save(os.path.join(root_path,img_name.split('.')[0] + '_flip.jpg'))return filp_imgdef rotation(root_path, img_name):img = Image.open(os.path.join(root_path, img_name))rotation_img = img.rotate(20) #旋转角度# rotation_img.save(os.path.join(root_path,img_name.split('.')[0] + '_rotation.jpg'))return rotation_imgdef randomColor(root_path, img_name): #随机颜色"""对图像进行颜色抖动:param image: PIL的图像image:return: 有颜色色差的图像image"""image = Image.open(os.path.join(root_path, img_name))random_factor = np.random.randint(0, 31) / 10. # 随机因子color_image = ImageEnhance.Color(image).enhance(random_factor) # 调整图像的饱和度random_factor = np.random.randint(10, 21) / 10. # 随机因子brightness_image = ImageEnhance.Brightness(color_image).enhance(random_factor) # 调整图像的亮度random_factor = np.random.randint(10, 21) / 10. # 随机因子contrast_image = ImageEnhance.Contrast(brightness_image).enhance(random_factor) # 调整图像对比度random_factor = np.random.randint(0, 31) / 10. # 随机因子return ImageEnhance.Sharpness(contrast_image).enhance(random_factor) # 调整图像锐度def contrastEnhancement(root_path, img_name): # 对比度增强image = Image.open(os.path.join(root_path, img_name))enh_con = ImageEnhance.Contrast(image)contrast = 1.5image_contrasted = enh_con.enhance(contrast)return image_contrasteddef brightnessEnhancement(root_path,img_name):#亮度增强image = Image.open(os.path.join(root_path, img_name))enh_bri = ImageEnhance.Brightness(image)brightness = 1.5image_brightened = enh_bri.enhance(brightness)return image_brighteneddef colorEnhancement(root_path,img_name):#颜色增强image = Image.open(os.path.join(root_path, img_name))enh_col = ImageEnhance.Color(image)color = 1.5image_colored = enh_col.enhance(color)return image_coloredfrom PIL import Imagefrom PIL import ImageEnhanceimport osimport cv2import numpy as npimageDir="./image"#要改变的图片的路径文件夹saveDir="./save" #要保存的图片的路径文件夹path = os.path.join(imageDir)rootlist=[]dirlist=[]filelist=[]for root, dirnames, filenames in os.walk(path):rootlist.append(root)for dirname in dirnames:dir_save_path = os.path.join(saveDir,dirname)if not os.path.exists(dir_save_path):os.makedirs(dir_save_path)dirlist.append(os.path.join(dir_save_path))del rootlist[0]print(dirlist)for i in range(len(rootlist)):for r, d, f in os.walk(rootlist[i]):for name in f:filelist.append(os.path.join(name))filelist.sort(key=lambda x:int(x[:-4]))for j in range(len(filelist)):saveName= filelist[j][:-4]+"id.png"image = Image.open(os.path.join(rootlist[i], filelist[j]))image.save(os.path.join(dirlist[i],saveName))saveName= filelist[j][:-4]+"fl.png"saveImage=flip(rootlist[i], filelist[j])saveImage.save(os.path.join(dirlist[i],saveName))saveName= filelist[j][:-4]+"ro.png"saveImage=rotation(rootlist[i], filelist[j])saveImage.save(os.path.join(dirlist[i],saveName))saveName= filelist[j][:-4]+"be.png"saveImage=brightnessEnhancement(rootlist[i], filelist[j])saveImage.save(os.path.join(dirlist[i],saveName))saveName= filelist[j][:-4]+"ra.png"saveImage=randomColor(rootlist[i], filelist[j])saveImage.save(os.path.join(dirlist[i],saveName))saveName= filelist[j][:-4]+"con.png"saveImage=contrastEnhancement(rootlist[i], filelist[j])saveImage.save(os.path.join(dirlist[i],saveName)) saveName= filelist[j][:-4]+"col.png"saveImage=colorEnhancement(rootlist[i], filelist[j])saveImage.save(os.path.join(dirlist[i],saveName))filelist.clear()
如果觉得《数据集增强——图像翻转 旋转 随机颜色 对比度 亮度 颜色增强》对你有帮助，请点赞、收藏，并留下你的观点哦！
