训练集的数据增强

此代码的作用是把像下图所示的文件数据增强之后还是能分类成三个文件夹

然后文件夹里面再是文件的名字

里面就是照片了

话不多说上代码

from PIL import Image
from PIL import ImageEnhance
import os
import cv2
import numpy as np

def flip(root_path, img_name):   # 翻转图像
    img = Image.open(os.path.join(root_path, img_name))
    filp_img = img.transpose(Image.Transpose.FLIP_LEFT_RIGHT)
    return filp_img

def rotation(root_path, img_name):
    img = Image.open(os.path.join(root_path, img_name))
    rotation_img = img.rotate(20)  # 旋转角度
    return rotation_img

def randomColor(root_path, img_name):  # 随机颜色
    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

def process_images(base_dir, save_dir):
    for subdir, dirs, files in os.walk(base_dir):
        for file in files:
            image_path = os.path.join(subdir, file)

            # 获取原始图片所在的相对路径（相对于base_dir）
            relative_path = os.path.relpath(image_path, base_dir)

            # 构造保存的文件名，保持原始目录结构
            save_name = os.path.join(relative_path[:-4] + "_flipped.jpg")
            save_path = os.path.join(save_dir, save_name)

            # 确保保存文件的目录及上级目录都已存在
            os.makedirs(os.path.dirname(save_path), exist_ok=True)

            # 进行图像处理并保存
            save_image = flip(subdir, file)
            save_image.save(save_path)

            # 其他增强方式的处理和保存，省略...

# 设置原始数据集目录和保存目录
base_directory = r"D:\Data3_split(稳妥)\train"
save_directory = r"D:\AAA"  # 可以改为原始的base_directory

# 处理图像
process_images(base_directory, save_directory)
process_images(base_directory, save_directory)