CV数据增强

jh__chen

于 2024-04-28 00:18:11 发布

阅读量396

点赞数 5

文章标签： python opencv 计算机视觉

本文链接：https://blog.csdn.net/jh__chen/article/details/138261503

版权

使用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_images_in_folder(folder_path, output_folder, augmentation_params):
    os.makedirs(output_folder, exist_ok=True)
    for filename in os.listdir(folder_path):
        if filename.endswith(('.jpg', '.png', '.jpeg')):  # 只处理图像文件
            image_path = os.path.join(folder_path, filename)
            image = cv2.imread(image_path)
            augmented_image = data_augmentation(image, augmentation_params)
            output_path = os.path.join(output_folder, filename)
            cv2.imwrite(output_path, augmented_image)

# 示例参数
augmentation_params = {
    'flip': True,
    'rotate': 15,  # 角度范围为正负15度
    'scale': 0.1,  # 尺度范围为原尺寸的正负10%
    'crop': 0.1,   # 裁剪比例为原始图像的正负10%
    'brightness': 0.2,  # 亮度调整范围为正负20%
    'contrast': 0.2,    # 对比度调整范围为正负20%
    'noise': True,
    'color_shift': False,
    'elastic_deformation': False
}

# 指定输入和输出文件夹，并执行数据增强
input_folder = 'input_images_folder'
output_folder = 'output_images_folder'
augment_images_in_folder(input_folder, output_folder, augmentation_params)

翻转：

import cv2

def flip_image(image_path, flip_code):
    """
    对图像进行翻转。

    参数：
    - image_path：图像文件的路径。
    - flip_code：翻转方式，可以是以下值之一：
        * 1：水平翻转（沿 y 轴翻转）。
        * 0：垂直翻转（沿 x 轴翻转）。
        * -1：同时在水平和垂直方向翻转。
    返回值：
    - flipped_image：翻转后的图像。
    """
    # 读取图像
    image = cv2.imread(image_path)
    
    # 检查图像是否成功读取
    if image is None:
        print("无法读取图像:", image_path)
        return None
    
    # 进行翻转
    flipped_image = cv2.flip(image, flip_code)
    
    return flipped_image

# 示例使用
image_path = "example.jpg"
flip_code = 1  # 水平翻转
flipped_image = flip_image(image_path, flip_code)

# 显示翻转后的图像
cv2.imshow("Flipped Image", flipped_image)
cv2.waitKey(0)
cv2.destroyAllWindows()

这个函数接受两个参数：图像文件的路径 image_path 和翻转方式 flip_code。flip_code 的值可以是 1（水平翻转）、0（垂直翻转）或 -1（同时在水平和垂直方向翻转）。你可以根据需要自定义翻转的方式来使用这个函数。

旋转、缩放和平移

import cv2
import numpy as np

def transform_image(image_path, rotation_angle=0, scale_factor=1.0, translation=(0, 0)):
    """
    对图像进行旋转、缩放和平移。

    参数：
    - image_path：图像文件的路径。
    - rotation_angle：旋转角度（单位：度），默认为 0。
    - scale_factor：缩放因子，默认为 1.0。
    - translation：平移距离，格式为 (dx, dy)，默认为 (0, 0)。

    返回值：
    - transformed_image：旋转、缩放和平移后的图像。
    """
    # 读取图像
    image = cv2.imread(image_path)

    # 检查图像是否成功读取
    if image is None:
        print("无法读取图像:", image_path)
        return None

    # 获取图像尺寸
    height, width = image.shape[:2]

    # 构建旋转矩阵
    rotation_matrix = cv2.getRotationMatrix2D((width/2, height/2), rotation_angle, scale_factor)

    # 进行旋转、缩放和平移
    transformed_image = cv2.warpAffine(image, rotation_matrix, (width, height))
    transformed_image = cv2.warpAffine(transformed_image, np.float32([[1, 0, translation[0]], [0, 1, translation[1]]]), (width, height))

    return transformed_image

# 示例使用
image_path = "example.jpg"
rotation_angle = 45  # 旋转角度为45度
scale_factor = 1.5  # 缩放因子为1.5
translation = (50, 50)  # 平移距离为(50, 50)
transformed_image = transform_image(image_path, rotation_angle, scale_factor, translation)

# 显示处理后的图像
cv2.imshow("Transformed Image", transformed_image)
cv2.waitKey(0)
cv2.destroyAllWindows()

裁剪、亮度调整、对比度调整

import cv2

def crop_image(image_path, crop_area):
    """
    对图像进行裁剪。

    参数：
    - image_path：图像文件的路径。
    - crop_area：裁剪区域，格式为 (x, y, width, height)。

    返回值：
    - cropped_image：裁剪后的图像。
    """
    # 读取图像
    image = cv2.imread(image_path)

    # 检查图像是否成功读取
    if image is None:
        print("无法读取图像:", image_path)
        return None

    # 进行裁剪
    x, y, width, height = crop_area
    cropped_image = image[y:y+height, x:x+width]

    return cropped_image

def adjust_brightness(image_path, brightness_factor):
    """
    调整图像的亮度。

    参数：
    - image_path：图像文件的路径。
    - brightness_factor：亮度调整因子，大于1表示增加亮度，小于1表示降低亮度。

    返回值：
    - adjusted_image：调整亮度后的图像。
    """
    # 读取图像
    image = cv2.imread(image_path)

    # 检查图像是否成功读取
    if image is None:
        print("无法读取图像:", image_path)
        return None

    # 进行亮度调整
    adjusted_image = cv2.convertScaleAbs(image, alpha=brightness_factor, beta=0)

    return adjusted_image

def adjust_contrast(image_path, contrast_factor):
    """
    调整图像的对比度。

    参数：
    - image_path：图像文件的路径。
    - contrast_factor：对比度调整因子，大于1表示增加对比度，小于1表示降低对比度。

    返回值：
    - adjusted_image：调整对比度后的图像。
    """
    # 读取图像
    image = cv2.imread(image_path)

    # 检查图像是否成功读取
    if image is None:
        print("无法读取图像:", image_path)
        return None

    # 进行对比度调整
    adjusted_image = cv2.convertScaleAbs(image, alpha=contrast_factor, beta=128 * (1 - contrast_factor))

    return adjusted_image

# 示例使用
image_path = "example.jpg"

# 裁剪参数，格式为 (x, y, width, height)
crop_area = (100, 100, 300, 300)

# 亮度调整因子
brightness_factor = 1.5

# 对比度调整因子
contrast_factor = 1.5

# 对图像进行裁剪
cropped_image = crop_image(image_path, crop_area)

# 对图像进行亮度调整
adjusted_brightness_image = adjust_brightness(image_path, brightness_factor)

# 对图像进行对比度调整
adjusted_contrast_image = adjust_contrast(image_path, contrast_factor)

# 显示处理后的图像
cv2.imshow("Cropped Image", cropped_image)
cv2.imshow("Adjusted Brightness Image", adjusted_brightness_image)
cv2.imshow("Adjusted Contrast Image", adjusted_contrast_image)
cv2.waitKey(0)
cv2.destroyAllWindows()

这个示例代码中包含了三个函数：crop_image 用于裁剪图像，adjust_brightness 用于调整图像的亮度，adjust_contrast 用于调整图像的对比度。你可以根据需要自定义裁剪区域、亮度调整因子和对比度调整因子来使用这些函数。

噪声添加、颜色变换、弹性变形

import cv2
import numpy as np

def add_noise(image_path, noise_type='gaussian', noise_strength=0.1):
    """
    向图像添加噪声。

    参数：
    - image_path：图像文件的路径。
    - noise_type：噪声类型，可以是 'gaussian'（高斯噪声）或 'salt_and_pepper'（椒盐噪声），默认为 'gaussian'。
    - noise_strength：噪声强度，范围为 [0, 1]，默认为 0.1。

    返回值：
    - noisy_image：添加噪声后的图像。
    """
    # 读取图像
    image = cv2.imread(image_path)

    # 检查图像是否成功读取
    if image is None:
        print("无法读取图像:", image_path)
        return None

    # 生成噪声
    if noise_type == 'gaussian':
        noise = np.random.normal(loc=0, scale=1, size=image.shape).astype(np.uint8)
        noisy_image = cv2.addWeighted(image, 1 - noise_strength, noise, noise_strength, 0)
    elif noise_type == 'salt_and_pepper':
        salt_and_pepper = np.random.choice([0, 1], size=image.shape[:2] + (1,), p=[1 - noise_strength, noise_strength])
        noisy_image = image.copy()
        noisy_image[salt_and_pepper == 1] = [255, 255, 255]  # 白噪声
        noisy_image[salt_and_pepper == 0] = [0, 0, 0]        # 黑噪声
    else:
        print("不支持的噪声类型:", noise_type)
        return None

    return noisy_image

def color_shift(image_path, shift_factor=0.1):
    """
    对图像进行颜色变换。

    参数：
    - image_path：图像文件的路径。
    - shift_factor：颜色变换因子，范围为 [0, 1]，默认为 0.1。

    返回值：
    - shifted_image：颜色变换后的图像。
    """
    # 读取图像
    image = cv2.imread(image_path)

    # 检查图像是否成功读取
    if image is None:
        print("无法读取图像:", image_path)
        return None

    # 进行颜色变换
    hsv_image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
    hsv_image[:, :, 1] = np.clip(hsv_image[:, :, 1] * (1 + shift_factor), 0, 255).astype(np.uint8)
    shifted_image = cv2.cvtColor(hsv_image, cv2.COLOR_HSV2BGR)

    return shifted_image

def elastic_deformation(image_path, alpha=10, sigma=5):
    """
    对图像进行弹性变形。

    参数：
    - image_path：图像文件的路径。
    - alpha：弹性变形参数，控制变形的程度，默认为 10。
    - sigma：高斯滤波器的标准差，影响变形的平滑程度，默认为 5。

    返回值：
    - deformed_image：弹性变形后的图像。
    """
    # 读取图像
    image = cv2.imread(image_path)

    # 检查图像是否成功读取
    if image is None:
        print("无法读取图像:", image_path)
        return None

    # 生成随机位移场
    rows, cols = image.shape[:2]
    random_field_x = np.random.uniform(low=-1, high=1, size=(rows, cols)).astype(np.float32)
    random_field_y = np.random.uniform(low=-1, high=1, size=(rows, cols)).astype(np.float32)

    # 对位移场进行高斯滤波
    random_field_x_smoothed = cv2.GaussianBlur(random_field_x, (0, 0), sigma)
    random_field_y_smoothed = cv2.GaussianBlur(random_field_y, (0, 0), sigma)

    # 计算每个像素的新坐标
    map_x = np.arange(cols) + random_field_x_smoothed * alpha
    map_y = np.arange(rows) + random_field_y_smoothed * alpha

    # 对图像进行弹性变形
    deformed_image = cv2.remap(image, map_x, map_y, interpolation=cv2.INTER_LINEAR)

    return deformed_image

# 示例使用
image_path = "example.jpg"

# 噪声添加参数
noise_type = 'gaussian'  # 高斯噪声
noise_strength = 0.1     # 强度为 0.1

# 颜色变换参数
shift_factor = 0.1  # 颜色变换因子为 0.1

# 弹性变形参数
alpha = 10  # 弹性变形参数为 10
sigma = 5   # 高斯滤波器的标准差为 5

# 对图像添加噪声
noisy_image = add_noise(image_path, noise_type, noise_strength)

# 对图像进行颜色变换
shifted_image = color_shift(image_path, shift_factor)

# 对图像进行弹性变形
deformed_image = elastic_deformation(image_path, alpha, sigma)

# 显示处理后的图像
cv2.imshow("Noisy Image", noisy_image)
cv2.imshow("Shifted Image", shifted_image)
cv2.imshow("Deformed Image", deformed_image)
cv2.waitKey(0)
cv2.destroyAllWindows()

这个函数接受四个参数：图像文件的路径 image_path、旋转角度 rotation_angle、缩放因子 scale_factor 和平移距离 translation。你可以根据需要自定义这些参数来对图像进行旋转、缩放和平移。

创建字典调用方法

import cv2
import numpy as np
import random

def data_augmentation(image, augmentation_params):
    """
    对图像进行数据增强。

    参数：
    - image：要增强的图像。
    - augmentation_params：数据增强参数，格式为 {'flip': True, 'rotate': 15, 'scale': 0.1, 'crop': 0.1, 'brightness': 0.2, 'contrast': 0.2, 'noise': True, 'color_shift': False, 'elastic_deformation': False}。

    返回值：
    - augmented_image：增强后的图像。
    """
    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 = cv2.flip(augmented_image, flip_code)
            elif param == 'rotate':
                angle = random.uniform(-value, value)
                rows, cols = augmented_image.shape[:2]
                matrix = cv2.getRotationMatrix2D((cols/2, rows/2), angle, 1)
                augmented_image = cv2.warpAffine(augmented_image, matrix, (cols, rows))
            elif param == 'scale':
                scale_factor = random.uniform(1 - value, 1 + value)
                augmented_image = cv2.resize(augmented_image, None, fx=scale_factor, fy=scale_factor, interpolation=cv2.INTER_LINEAR)
            elif param == 'crop':
                rows, cols = augmented_image.shape[:2]
                crop_size = (int(cols * (1 - value)), int(rows * (1 - value)))
                x = random.randint(0, cols - crop_size[0])
                y = random.randint(0, rows - crop_size[1])
                augmented_image = augmented_image[y:y+crop_size[1], x:x+crop_size[0]]
            elif param == 'brightness':
                brightness_factor = random.uniform(1 - value, 1 + value)
                augmented_image = cv2.convertScaleAbs(augmented_image, alpha=brightness_factor, beta=0)
            elif param == 'contrast':
                contrast_factor = random.uniform(1 - value, 1 + value)
                augmented_image = cv2.convertScaleAbs(augmented_image, alpha=contrast_factor, beta=128 * (1 - 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

# 示例参数
augmentation_params = {
    'flip': True,
    'rotate': 15,  # 角度范围为正负15度
    'scale': 0.1,  # 尺度范围为原尺寸的正负10%
    'crop': 0.1,   # 裁剪比例为原始图像的正负10%
    'brightness': 0.2,  # 亮度调整范围为正负20%
    'contrast': 0.2,    # 对比度调整范围为正负20%
    'noise': True,
    'color_shift': False,
    'elastic_deformation': False
}

# 调用数据增强函数
augmented_image = data_augmentation(image, augmentation_params)