网站首页 > 人工智能 > 正文

人工智能深度学习图像处理数据集增强——图像翻转、旋转、随机颜色、对比度、亮度、颜色增强

信息系统项目管理高级人工智能 2024-01-05 17 0

本代码共采用了四种数据增强，如采用其他数据增强方式，可以参考本代码，随意替换。

'''

imageDir 为原数据集的存放位置

saveDir 为数据增强后数据的存放位置

'''

from PIL import Image

from PIL import ImageEnhance

import os

import cv2

import numpy as np

import glob

def flip(root_path): #翻转图像

img = Image.open(root_path)

filp_img = img.transpose(Image.FLIP_LEFT_RIGHT)

return filp_img

def rotation(root_path):

img = Image.open(root_path)

rotation_img = img.rotate(90) #旋转角度

return rotation_img

def brightnessEnhancement(root_path):#亮度增强

image = Image.open(root_path)

enh_bri = ImageEnhance.Brightness(image)

brightness = 1.5

image_brightened = enh_bri.enhance(brightness)

return image_brightened

def contrastEnhancement(root_path): # 对比度增强

image = Image.open(root_path)

enh_con = ImageEnhance.Contrast(image)

contrast = 1.5

image_contrasted = enh_con.enhance(contrast)

return image_contrasted

def colorEnhancement(root_path):#颜色增强

image = Image.open(root_path)

enh_col = ImageEnhance.Color(image)

color = 1.5

image_colored = enh_col.enhance(color)

return image_colored

def 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_names

from matplotlib.animation import ImageMagickBase

def 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_img

def 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_img

def 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.5

image_contrasted = enh_con.enhance(contrast)

return image_contrasted

def brightnessEnhancement(root_path,img_name):#亮度增强

image = Image.open(os.path.join(root_path, img_name))

enh_bri = ImageEnhance.Brightness(image)

brightness = 1.5

image_brightened = enh_bri.enhance(brightness)

return image_brightened

def colorEnhancement(root_path,img_name):#颜色增强

image = Image.open(os.path.join(root_path, img_name))

enh_col = ImageEnhance.Color(image)

color = 1.5

image_colored = enh_col.enhance(color)

return image_colored

from PIL import Image

from PIL import ImageEnhance

import os

import cv2

import numpy as np

imageDir="./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])