CV数据增强

使用python写一个对一个文件夹中的图像进行批量可以自定义数据增强程度的翻转、旋转、缩放、裁剪、亮度调整、对比度调整、噪声添加、噪声添加、颜色变换、弹性变形这些方法的代码

import os
import cv2
import numpy as np
import random

# 定义数据增强方法
def flip(image, flip_code):
    return cv2.flip(image, flip_code)

def rotate(image, angle):
    rows, cols = image.shape[:2]
    matrix = cv2.getRotationMatrix2D((cols/2, rows/2), angle, 1)
    return cv2.warpAffine(image, matrix, (cols, rows))

def scale(image, scale_factor):
    return cv2.resize(image, None, fx=scale_factor, fy=scale_factor, interpolation=cv2.INTER_LINEAR)

def crop(image, crop_size):
    rows, cols = image.shape[:2]
    x = random.randint(0, cols - crop_size[0])
    y = random.randint(0, rows - crop_size[1])
    return image[y:y+crop_size[1], x:x+crop_size[0]]

def adjust_brightness(image, brightness_factor):
    return cv2.convertScaleAbs(image, alpha=brightness_factor, beta=0)

def adjust_contrast(image, contrast_factor):
    return cv2.convertScaleAbs(image, alpha=contrast_factor, beta=128 * (1 - contrast_factor))

def add_noise(image):
    noise = np.random.normal(loc=0, scale=25, size=image.shape).astype(np.uint8)
    return cv2.add(image, noise)

def color_shift(image):
    # 这里可以自定义颜色变换方法，例如转换为灰度图像或调整颜色通道
    # 以示例简单，这里只返回原图像
    return image

def elastic_deformation(image):
    # 这里可以自定义弹性变形方法
    # 以示例简单，这里只返回原图像
    return image

# 数据增强函数
def data_augmentation(image, augmentation_params):
    augmented_image = image.copy()
    for param, value in augmentation_params.items():
        if value:  # 只对设置为 True 的参数执行数据增强
            if param == 'flip':
                flip_code = random.choice([-1, 0, 1])
                augmented_image = flip(augmented_image, flip_code)
            elif param == 'rotate':
                angle = random.uniform(-value, value)
                augmented_image = rotate(augmented_image, angle)
            elif param == 'scale':
                scale_factor = random.uniform(1 - value, 1 + value)
                augmented_image = scale(augmented_image, scale_factor)
            elif param == 'crop':
                crop_size = (int(image.shape[1] * (1 - value)), int(image.shape[0] * (1 - value)))
                augmented_image = crop(augmented_image, crop_size)
            elif param == 'brightness':
                brightness_factor = random.uniform(1 - value, 1 + value)
                augmented_image = adjust_brightness(augmented_image, brightness_factor)
            elif param == 'contrast':
                contrast_factor = random.uniform(1 - value, 1 + value)
                augmented_image = adjust_contrast(augmented_image, contrast_factor)
            elif param == 'noise':
                augmented_image = add_noise(augmented_image)
            elif param == 'color_shift':
                augmented_image = color_shift(augmented_image)
            elif param == 'elastic_deformation':
                augmented_image = elastic_deformation(augmented_image)
    return augmented_image

# 遍历文件夹中的图像文件并进行数据增强
def augment_image