扩增数据集

扩增方式：

1.调节对比度

2.调节亮度

3.调节饱和度

4.调节色度

5.旋转图片

6.模糊图片

7.左右翻转图片，就跟镜像后的效果那样，下同

8.上下翻转图片

9.加入噪声，并可加入不同类型的噪声，共三种类型

10.对一张图片用上面所述的9种方法叠加处理，当然也可以仅选择其中的几种方法

11.随机裁剪图片

这11种扩增方式都可独立开关，按需选择。然后裁剪有点鸡肋，因为往往裁剪出来的图片没有要打标签的物体了，所以最好关了

把要扩增的图片放入一个文件夹里，在程序里输入此文件夹的路径，再改改2个输出文件夹的路径，开关一下扩增方式即可，默认扩增方式全开。然后这两个输出文件夹的区别在于对图片的命名不同，图一样只是对生成的图片按0001，0002，.... ，重新命名了

import os
import numpy as np
from PIL import Image, ImageFilter, ImageEnhance
import shutil


def contrast_enhancement(root_path, img_name, contrast):  # 1.调节对比度
    image = Image.open(os.path.join(root_path, img_name))
    enh_con = ImageEnhance.Contrast(image)
    image_contrasted = enh_con.enhance(contrast)
    return image_contrasted


def change_brightness(root_path, img_name, brightness):  # 2.调节亮度
    image = Image.open(os.path.join(root_path, img_name))
    enh_bri = ImageEnhance.Brightness(image)
    image_brightened = enh_bri.enhance(brightness)
    return image_brightened


def change_saturation(root_path, image, saturation_factor):  # 3.调节饱和度
    image = Image.open(os.path.join(root_path, image))
    image_hsv = image.convert("HSV")
    enhancer = ImageEnhance.Color(image_hsv)
    enhanced_image_hsv = enhancer.enhance(saturation_factor)
    enhanced_image = enhanced_image_hsv.convert("RGB")
    return enhanced_image


def change_hue(root_path, img_name, hue_factor):  # 4.调节色调
    image = Image.open(os.path.join(root_path, img_name))
    image_hsv = image.convert("HSV")
    image_array = np.array(image_hsv)
    image_array[:, :, 0] = (image_array[:, :, 0] + hue_factor) % 360
    modified_image_hsv = Image.fromarray(image_array, mode="HSV")
    modified_image_rgb = modified_image_hsv.convert("RGB")
    return modified_image_rgb


def rotation(root_path, img_name, random_angle):  # 5.旋转图片
    image = Image.open(os.path.join(root_path, img_name))
    rotation_img = image.rotate(random_angle)
    return rotation_img


def blur_image(root_path, image, fuzzyDegree):  # 6.模糊图片
    image = Image.open(os.path.join(root_path, image))
    blurred_image = image.filter(ImageFilter.GaussianBlur(fuzzyDegree))
    return blurred_image


def flip_img_left_and_right(root_path, img_name):  # 7.左右翻转图片，就跟镜像后的效果那样，下同
    image = Image.open(os.path.join(root_path, img_name))
    img_left_and_right = image.transpose(Image.FLIP_LEFT_RIGHT)
    return img_left_and_right


def flip_img_top_and_bottom(root_path, img_name):  # 8.上下翻转图片
    image = Image.open(os.path.join(root_path, img_name))
    img_top_and_bottom = image.transpose(Image.FLIP_TOP_BOTTOM)
    return img_top_and_bottom


def add_noise(root_path, img_name, noise_type, noise_factor):  # 9.加入噪声，并可加入不同类型的噪声
    image = Image.open(os.path.join(root_path, img_name))
    image_array = np.array(image)
    if noise_type == "gaussian":
        noise = np.random.normal(0, noise_factor, image_array.shape)
        noisy_image_array = image_array + noise
    elif noise_type == "salt":
        noise = np.random.randint(0, noise_factor, image_array.shape)
        noisy_image_array = np.where(noise == 0, 255, image_array)
    elif noise_type == "pepper":
        noise = np.random.randint(0, noise_factor, image_array.shape)
        noisy_image_array = np.where(noise == 0, 0, image_array)
    else:
        raise ValueError("Invalid noise type. Choose from 'gaussian', 'salt', or 'pepper'.")
    noisy_image = Image.fromarray(noisy_image_array.astype(np.uint8))
    return noisy_image


def crop_image(root_path, img_name, left, top, right, bottom):  # 10.随机裁剪图片，所以出来的图片有可能不再包含想要打标签的物体
    image = Image.open(os.path.join(root_path, img_name))
    cropped_image = image.crop((left, top, right, bottom))
    return cropped_image


def multiple_image_processing(root_path, img_name, contrastRatio, brightness, saturation, hue, angle, fuzzyDegree, noiseType, noiseFactor,
                              contrastRatioIsON, brightnesIsON, saturationIsON, hueIsON, rotationIsON, blurIsON, flipLeftAndRightIsON,
                              flipTopAndBottomIsON, addNoiseIsON):  # 11.对同一张图片进行多种处理，某种处理设为True则进行这种处理，False反之
    image = Image.open(os.path.join(root_path, img_name))
    ways = ''
    if contrastRatioIsON:
        enh_con = ImageEnhance.Contrast(image)
        image = enh_con.enhance(contrastRatio)
        ways += '-contrastRatio' + str(round(contrastRatio, 2))
    if brightnesIsON:
        enh_bri = ImageEnhance.Brightness(image)
        image = enh_bri.enhance(brightness)
        ways += '-brightness' + str(round(brightness, 2))
    if saturationIsON:
        image_hsv = image.convert("HSV")
        enhancer = ImageEnhance.Color(image_hsv)
        enhanced_image_hsv = enhancer.enhance(saturation)
        image = enhanced_image_hsv.convert("RGB")
        ways += '-saturation' + str(round(saturation, 2))
    if hueIsON:
        image_hsv = image.convert("HSV")
        image_array = np.array(image_hsv)
        image_array[:, :, 0] = (image_array[:, :, 0] + hue) % 360
        modified_image_hsv = Image.fromarray(image_array, mode="HSV")
        image = modified_image_hsv.convert("RGB")
        ways += '-hue' + str(hue)
    if rotationIsON:
        image = image.rotate(angle)
        ways += '-angle' + str(angle)
    if blurIsON:
        image = image.filter(ImageFilter.GaussianBlur(fuzzyDegree))
        ways += '-fuzzy' + str(round(fuzzyDegree, 2))
    if flipLeftAndRightIsON:
        image = image.transpose(Image.FLIP_LEFT_RIGHT)
        ways += '-flipLeftAndRight'
    if flipTopAndBottomIsON:
        image = image.transpose(Image.FLIP_TOP_BOTTOM)
        ways += '-flipTopAndBottom'
    if addNoiseIsON:
        image_array = np.array(image)
        if noiseType == "gaussian":
            noise = np.random.normal(0, noiseFactor, image_array.shape)
            noisy_image_array = image_array + noise
        elif noiseType == "salt":
            noise = np.random.randint(0, noiseFactor, image_array.shape)
            noisy_image_array = np.where(noise == 0, 255, image_array)
        elif noiseType == "pepper":
            noise = np.random.randint(0, noiseFactor, image_array.shape)
            noisy_image_array = np.where(noise == 0, 0, image_array)
        else:
            raise ValueError("Invalid noise type. Choose from 'gaussian', 'salt', or 'pepper'.")
        image = Image.fromarray(noisy_image_array.astype(np.uint8))
        ways += '-noiseTypeIs' + noiseType + 'AndNoiseFactorIs' + str(noiseFactor)
    return image, ways


def create_image(imageDir, saveDir):
    i = 0
    for name in os.listdir(imageDir):

        i = i + 1
        saveName = str(i) + '-' + os.path.splitext(name)[0] + "-originalImage.jpg"
        Image.open(os.path.join(imageDir, name)).save(os.path.join(saveDir, saveName))  # 先将读取的原图也保存下来放入生成的新数据集中，这样就不用再手动复制粘贴合并原图及其衍生出来的图片了

        turnONContrast, wayOfContrast = True, 2  # 想对图片进行对比度处理就设为True,并选择对比度处理方式。方式一（值设为1）：自定义对比度的值，有几个值就会生成几张对应的新图。方式二（值设为2）：产生随机的对比度值，自定义对比度值的个数即可，下同
        contrastRatio = [1, 1.2, 1.4, 1.6, 1.8, 2.0]
        numberOfContrastRatio = 5

        turnONBrightness, wayOfBrightness = True, 2
        brightness = [0.5, 0.6, 0.7, 0.8, 0.9, 1, 1.1, 1.2, 1.3, 1.4, 1.5]
        numberOfBrightness = 5

        turnONSaturation, wayOfSaturation = True, 2
        saturation = [0.5, 0.6, 0.7, 0.8, 0.9, 1, 1.1, 1.2, 1.3, 1.4, 1.5]
        numberOfSaturation = 5

        turnONHue, wayOfHue = True, 2
        hue = [0, 45, 90, 135, 180, 225, 270, 315, 360]
        numberOfHue = 5

        turnONAngle, wayOfAngle = True, 2
        angle = [0, 45, 90, 135, 180, 225, 270, 315, 360]
        numberOfAngle = 5

        turnONFuzzy, wayOfFuzzy = True, 2
        fuzzyDegree = [1, 1.2, 1.4, 1.8, 2.0, 2.4, 2.6, 2.8, 3.0]
        numberOfFuzzyDegree = 5

        turnONNoise, wayOfNoise = True, 2
        noiseFactor = [10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30]
        randomValue = 0  # 属于方式一，自定义噪声的类型，0、1、2分别表示加入的噪声类型为''gaussian'', ''salt'', ''pepper''
        numberOfNoise = 5

        turnONCrop = True  # 裁剪最好关了，因为往往裁剪不到需要打标签的目标
        numberOfCrop = 5

        turnONFlipLeftAndRight, turnONFlipTopAndBottom = True, True

        turnONMultiprocessing, wayOfMultiprocessing = True, 2
        contrastRatioOfMul = 2  # 调节对比度，参考值1.0~2.0
        brightnessOfMul = 1.5  # 调节亮度，参考值0.5~1.5
        saturationOfMul = 1.5  # 调节饱和度，参考值0.5~1.5
        hueOfMul = 180  # 调节色调，参考值0~360
        angleOfMul = 0  # 调节旋转角度，参考值0~360
        fuzzyDegreeOfMul = 3.0  # 调节模糊程度，参考值1.0~3.0
        noiseTypeOfMul = 'gaussian'  # 设置噪声类型，有''gaussian'', ''salt'', ''pepper''
        noiseFactorOfMul = 20  # 设置噪声的影响程度，参考值10~30

        if turnONContrast:
            if wayOfContrast == 2:
                contrastRatio = np.empty(numberOfContrastRatio)
                for k in range(numberOfContrastRatio):
                    contrastRatio[k] = 1 + np.random.random()  # [1.0, 2.0)
            for j in range(len(contrastRatio)):
                i = i + 1
                saveName = str(i) + '-' + os.path.splitext(name)[0] + "-contrastRatioIs" + str(round(contrastRatio[j], 2)) + ".jpg"
                saveImage = contrast_enhancement(imageDir, name, contrastRatio[j])
                saveImage.save(os.path.join(saveDir, saveName))

        if turnONBrightness:
            if wayOfBrightness == 2:
                brightness = np.empty(numberOfBrightness)
                for k in range(numberOfBrightness):
                    brightness[k] = 0.5 + np.random.random()
            for j in range(len(brightness)):
                i = i + 1
                saveName = str(i) + '-' + os.path.splitext(name)[0] + "-brightnessIs" + str(round(brightness[j], 2)) + ".jpg"
                saveImage = change_brightness(imageDir, name, brightness[j])
                saveImage.save(os.path.join(saveDir, saveName))

        if turnONSaturation:
            if wayOfSaturation == 2:
                saturation = np.empty(numberOfSaturation)
                for k in range(numberOfSaturation):
                    saturation[k] = 0.5 + np.random.random()
            for j in range(len(saturation)):
                i = i + 1
                saveName = str(i) + '-' + os.path.splitext(name)[0] + "-saturationIs" + str(round(saturation[j], 2)) + ".jpg"
                saveImage = change_saturation(imageDir, name, saturation[j])
                saveImage.save(os.path.join(saveDir, saveName))

        if turnONHue:
            if wayOfHue == 2:
                hue = np.empty(numberOfHue)
                for k in range(numberOfHue):
                    hue[k] = np.random.randint(0, 361)
            for j in range(len(hue)):
                i = i + 1
                saveName = str(i) + '-' + os.path.splitext(name)[0] + "-hueIs" + str(int(hue[j])) + ".jpg"
                saveImage = change_hue(imageDir, name, hue[j])
                saveImage.save(os.path.join(saveDir, saveName))

        if turnONAngle:
            if wayOfAngle == 2:
                angle = np.empty(numberOfAngle)
                for k in range(numberOfAngle):
                    angle[k] = np.random.randint(0, 360)  # 范围[0,359]，皆为整数
            for j in range(len(angle)):
                i = i + 1
                saveName = str(i) + '-' + os.path.splitext(name)[0] + "-angleIs" + str(int(angle[j])) + ".jpg"
                saveImage = rotation(imageDir, name, angle[j])
                saveImage.save(os.path.join(saveDir, saveName))

        if turnONFuzzy:
            if wayOfFuzzy == 2:
                fuzzyDegree = np.empty(numberOfFuzzyDegree)
                for k in range(numberOfFuzzyDegree):
                    fuzzyDegree[k] = 1 + 3 * np.random.random()
            for j in range(len(fuzzyDegree)):
                i = i + 1
                saveName = str(i) + '-' + os.path.splitext(name)[0] + "-fuzzyDegreeIs" + str(round(fuzzyDegree[j], 2)) + ".jpg"
                saveImage = blur_image(imageDir, name, fuzzyDegree[j])
                saveImage.save(os.path.join(saveDir, saveName))

        if turnONNoise:
            if wayOfNoise == 2:
                randomValue = np.random.randint(0, 3)
                noiseFactor = np.empty(numberOfNoise)
                for k in range(numberOfNoise):
                    noiseFactor[k] = np.random.randint(10, 31)
            noiseType = ['gaussian', 'salt', 'pepper']
            for j in range(len(noiseFactor)):
                i = i + 1
                saveName = str(i) + '-' + os.path.splitext(name)[0] + "-noiseTypeIs" + noiseType[randomValue] + "-noiseFactorIs" + str(int(noiseFactor[j])) + ".jpg"
                saveImage = add_noise(imageDir, name, noiseType[randomValue], noiseFactor[j])
                saveImage.save(os.path.join(saveDir, saveName))

        if turnONCrop:
            for j in range(numberOfCrop):
                i = i + 1
                saveName = str(i) + '-' + os.path.splitext(name)[0] + "-crop" + str(j + 1) + ".jpg"
                img = Image.open(os.path.join(imageDir, name))
                w, h = img.size
                left, top = np.random.randint(0, w), np.random.randint(0, h)
                right, bottom = left + np.random.randint(1, w), top + np.random.randint(1, h)
                saveImage = crop_image(imageDir, name, left, top, right, bottom)
                saveImage.save(os.path.join(saveDir, saveName))

        if turnONFlipLeftAndRight:
            i = i + 1
            saveName = str(i) + '-' + os.path.splitext(name)[0] + "-flipLeftAndRight.jpg"
            saveImage = flip_img_left_and_right(imageDir, name)
            saveImage.save(os.path.join(saveDir, saveName))

        if turnONFlipTopAndBottom:
            i = i + 1
            saveName = str(i) + '-' + os.path.splitext(name)[0] + "-flipTopAndBottom.jpg"
            saveImage = flip_img_top_and_bottom(imageDir, name)
            saveImage.save(os.path.join(saveDir, saveName))

        if turnONMultiprocessing:
            if wayOfMultiprocessing == 2:
                contrastRatioOfMul = 1 + np.random.random()
                brightnessOfMul = 0.5 + np.random.random()
                saturationOfMul = 0.5 + np.random.random()
                hueOfMul = np.random.randint(0, 361)
                angleOfMul = np.random.randint(0, 360)
                fuzzyDegreeOfMul = 1 + 3 * np.random.random()
                noise = ['gaussian', 'salt', 'pepper']
                noiseTypeOfMul = noise[np.random.randint(0, 3)]
                noiseFactorOfMul = np.random.randint(10, 31)
            i = i + 1
            saveImage, ways = multiple_image_processing(imageDir, name, contrastRatioOfMul, brightnessOfMul, saturationOfMul, hueOfMul, angleOfMul, fuzzyDegreeOfMul, noiseTypeOfMul, noiseFactorOfMul,
                                                        contrastRatioIsON=True, brightnesIsON=True, saturationIsON=True, hueIsON=True, rotationIsON=True,
                                                        blurIsON=True, flipLeftAndRightIsON=True, flipTopAndBottomIsON=True, addNoiseIsON=True)  # True, False
            saveName = str(i) + '-' + os.path.splitext(name)[0] + "-multiprocessing" + ways + ".jpg"
            saveImage.save(os.path.join(saveDir, saveName))


def create_and_copy_rename(source_folder, target_folder):
    if not os.path.exists(target_folder):
        os.makedirs(target_folder)
    file_list = os.listdir(target_folder)
    for file_name in file_list:
        file_path = os.path.join(target_folder, file_name)
        if os.path.isfile(file_path):
            os.remove(file_path)
    file_list = os.listdir(source_folder)
    length = len(str(len(file_list)))
    for idx, filename in enumerate(file_list):
        source_file = os.path.join(source_folder, filename)
        lengthOfIdx = len(str(idx + 1))
        string = ''
        for i in range(length - lengthOfIdx):
            string += '0'
        target_file = os.path.join(target_folder, string + str(idx + 1) + ".jpg")
        shutil.copy2(source_file, target_file)


if __name__ == '__main__':
    imageDir = "C:\\Users\\15231\\Desktop\\originalImages"  # 将所有的原图放在文件夹originalImages里，然后在此输入文件夹的路径
    saveDir = "C:\\Users\\15231\\Desktop\\generatedImages"  # 程序运行完后即可在此文件夹里得到各种处理后的图片
    if not os.path.exists(saveDir):  # 若文件夹generatedImages不存在则创建它
        os.makedirs(saveDir)
    file_list = os.listdir(saveDir)
    for file_name in file_list:
        file_path = os.path.join(saveDir, file_name)
        if os.path.isfile(file_path):
            os.remove(file_path)  # 生成并保存处理后的图片前先删除文件夹generatedImages里的所有文件，以免程序上次运行的结果图在文件夹的造成混淆
    create_image(imageDir, saveDir)
    renameDir = "C:\\Users\\15231\\Desktop\\rename"  # 对生成的图片按0001，0002，0003 ... 重命名并保存在文件夹rename里
    create_and_copy_rename(saveDir, renameDir)