使用albumentations进行图像数据集扩充方法

原创已于 2025-06-26 15:04:10 修改 · 703 阅读

11 ·

CC 4.0 BY-SA版权

文章标签：

#数据增广

于 2025-06-25 15:17:03 首次发布

一、背景

训练机器学习模型，尤其是计算机视觉模型，通常需要大量数据。这些模型从所见的样本中学习复杂的模式。通常，这些样本越多样化、数量越多，模型在处理新的、未见过的数据时的表现就越好。

然而，收集和标记海量数据集可能是一个主要的瓶颈。这通常成本高昂、耗时长，有时甚至根本无法收集到足够的数据。

数据增强（Data Augmentation）技术应运而生。它是一套通过生成现有数据点的修改副本来人为增加训练数据数量和多样性的技术。可以将其想象成通过向模型展示相同的内容（只是方式略有不同）来更稳健地训练模型。

我这里采用albumentations方法来进行图片数据的扩充。

albumentations官网：https://github.com/albumentations-team/albumentations?tab=readme-ov-file#getting-started

二、albumentations环境配置

电脑：机械臂组 4090 工作站
内存：64G
处理器：Intel i7-13700KF
显卡：4090
系统：ubuntu22.04
文件路径：/home/wd/DejaVu/albumentations

1、创建名为albumentations的虚拟环境，使用python3.9

conda create -n albumentations python=3.9

2、安装albumentations

# 激活conda环境
conda activate albumentations

# 安装albumentations
pip install -U albumentations

三、albumentations使用

遇到的坑：

之前python文件的名字为：download_collection.py，执行下面的命令后，报错：python3: can't open file '/home/wd/DejaVu/albumentations/download_collection': [Errno 2] No such file or directory

python3 download_collection

我以为是没有权限的原因，所以使用命令chmod +x download_collection加权限，但是依然报错chmod: 无法访问 'download_collection': 没有那个文件或目录。

解决办法：python文件的名称应该是不能允许有特殊符号，将文件download_collection.py改为augData.py即可成功运行。

四、颜色及亮度增广

这里记录一下使用Albumentations库的Python程序，用于扩充数据集，通过对图片应用颜色和亮度增强等变换。

电脑：机械臂组 4090 工作站
文件路径：/home/wd/DejaVu/albumentations
文件名称：augmentDatasetColor.py
完整程序：

import os
import cv2
import albumentations as A
from pathlib import Path
import numpy as np

def create_augmentation_pipeline():
    """定义数据增强流水线，包含颜色和亮度变换"""
    return A.Compose([
        A.RandomBrightnessContrast(
            brightness_limit=0.2,  # 亮度调整范围 [-0.2, 0.2]
            contrast_limit=0.2,    # 对比度调整范围 [-0.2, 0.2]
            p=0.5                  # 应用概率
        ),
        A.HueSaturationValue(
            hue_shift_limit=20,    # 色调调整范围
            sat_shift_limit=30,    # 饱和度调整范围
            val_shift_limit=20,    # 明度调整范围
            p=0.5
        ),
        A.ColorJitter(
            brightness=0.2,
            contrast=0.2,
            saturation=0.2,
            hue=0.1,
            p=0.5
        ),
        A.CLAHE(
            clip_limit=4.0,        # 对比度增强限制
            tile_grid_size=(8, 8), # 网格大小
            p=0.3
        )
    ])

def augment_images(input_dir, output_dir, num_augmentations=3):
    """
    对指定目录中的图片进行数据增强
    Args:
        input_dir: 输入图片目录
        output_dir: 增强后图片保存目录
        num_augmentations: 每张图片生成的增强版本数量
    """
    # 确保输出目录存在
    Path(output_dir).mkdir(parents=True, exist_ok=True)
    
    # 创建增强流水线
    transform = create_augmentation_pipeline()
    
    # 支持的图片扩展名
    valid_extensions = ('.jpg', '.jpeg', '.png')
    
    # 遍历输入目录中的所有图片
    for img_name in os.listdir(input_dir):
        if not img_name.lower().endswith(valid_extensions):
            continue
            
        # 读取图片
        img_path = os.path.join(input_dir, img_name)
        image = cv2.imread(img_path)
        if image is None:
            print(f"警告: 无法读取图片 {img_path}")
            continue
            
        # 转换为RGB（Albumentations需要RGB格式）
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        
        # 为每张图片生成指定数量的增强版本
        for i in range(num_augmentations):
            # 应用增强
            augmented = transform(image=image)
            aug_image = augmented['image']
            
            # 转换回BGR以保存
            aug_image = cv2.cvtColor(aug_image, cv2.COLOR_RGB2BGR)
            
            # 生成输出文件名
            base_name = Path(img_name).stem
            output_name = f"{base_name}_aug_{i+1}{Path(img_name).suffix}"
            output_path = os.path.join(output_dir, output_name)
            
            # 保存增强后的图片
            cv2.imwrite(output_path, aug_image)
            print(f"已保存增强图片: {output_path}")

def main():
    # 配置输入和输出目录
    input_directory = "/home/wd/DejaVu/albumentations/HumanIntention_data6Augmentation/pullAug"  # 替换为您的输入图片目录
    output_directory = "/home/wd/DejaVu/albumentations/HumanIntention_data6Augmentation/pullAugA"  # 替换为您的输出图片目录
    
    # 每张图片生成3个增强版本
    augment_images(input_directory, output_directory, num_augmentations=3)

if __name__ == "__main__":
    main()

缺点：原始数据和增广后的保存至同一个文件夹下了

五、旋转图片增广

使用Albumentations库，对图片进行小幅度旋转（±10度），并为每张原始图片生成4张增强版本。

电脑：机械臂组 4090 工作站
文件路径：/home/wd/DejaVu/albumentations
文件名称：augmentDatasetRotation.py
完整程序：

import os
import cv2
import albumentations as A
from pathlib import Path
import numpy as np

def create_augmentation_pipeline():
    """定义数据增强流水线，包含小幅度旋转变换"""
    return A.Compose([
        A.Rotate(
            limit=10,              # 旋转角度范围 [-10, 10] 度
            interpolation=cv2.INTER_LINEAR,  # 插值方法
            border_mode=cv2.BORDER_REFLECT,  # 边界填充方式
            p=1.0                  # 始终应用旋转
        )
    ])

def augment_images(input_dir, output_dir, num_augmentations=4):
    """
    对指定目录中的图片进行数据增强，生成旋转版本
    Args:
        input_dir: 输入图片目录
        output_dir: 增强后图片保存目录
        num_augmentations: 每张图片生成的增强版本数量
    """
    # 确保输出目录存在
    Path(output_dir).mkdir(parents=True, exist_ok=True)
    
    # 创建增强流水线
    transform = create_augmentation_pipeline()
    
    # 支持的图片扩展名
    valid_extensions = ('.jpg', '.jpeg', '.png')
    
    # 遍历输入目录中的所有图片
    for img_name in os.listdir(input_dir):
        if not img_name.lower().endswith(valid_extensions):
            continue
            
        # 读取图片
        img_path = os.path.join(input_dir, img_name)
        image = cv2.imread(img_path)
        if image is None:
            print(f"警告: 无法读取图片 {img_path}")
            continue
            
        # 转换为RGB（Albumentations需要RGB格式）
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        
        # 为每张图片生成4个增强版本
        for i in range(num_augmentations):
            # 应用增强
            augmented = transform(image=image)
            aug_image = augmented['image']
            
            # 转换回BGR以保存
            aug_image = cv2.cvtColor(aug_image, cv2.COLOR_RGB2BGR)
            
            # 生成输出文件名
            base_name = Path(img_name).stem
            output_name = f"{base_name}_rot_{i+1}{Path(img_name).suffix}"
            output_path = os.path.join(output_dir, output_name)
            
            # 保存增强后的图片
            cv2.imwrite(output_path, aug_image)
            print(f"已保存增强图片: {output_path}")

def main():
    # 配置输入和输出目录
    input_directory = "/home/wd/DejaVu/albumentations/HumanIntention_data6Augmentation/pullAug"  # 替换为您的输入图片目录
    output_directory = "/home/wd/DejaVu/albumentations/HumanIntention_data6Augmentation/pullAugRotation"  # 替换为您的输出图片目录
    
    # 每张图片生成4个增强版本
    augment_images(input_directory, output_directory, num_augmentations=4)

if __name__ == "__main__":
    main()

六、添加噪声增广

为原始图像添加噪声以扩充数据集，这里使用GaussNoise和ISONoise方法为每张图片生成4个包含不同噪声的增强版本。

电脑：机械臂组 4090 工作站
文件路径：/home/wd/DejaVu/albumentations
文件名称：augmentDatasetNoise.py
完整程序：

import os
import cv2
import albumentations as A
from pathlib import Path
import numpy as np

def create_augmentation_pipeline():
    """定义数据增强流水线，包含噪声变换"""
    return A.Compose([
        A.OneOf([
            A.GaussNoise(
                var_limit=(10.0, 50.0),  # 噪声方差范围
                mean=0,                  # 噪声均值
                p=1.0
            ),
            A.ISONoise(
                intensity=(0.1, 0.5),   # 噪声强度范围
                color_shift=(0.01, 0.3), # 颜色偏移范围
                p=1.0
            )
        ], p=1.0)
    ])

def augment_images(input_dir, output_dir, num_augmentations=4):
    """
    对指定目录中的图片进行数据增强，添加噪声
    Args:
        input_dir: 输入图片目录
        output_dir: 增强后图片保存目录
        num_augmentations: 每张图片生成的增强版本数量
    """
    # 确保输出目录存在
    Path(output_dir).mkdir(parents=True, exist_ok=True)
    
    # 创建增强流水线
    transform = create_augmentation_pipeline()
    
    # 支持的图片扩展名
    valid_extensions = ('.jpg', '.jpeg', '.png')
    
    # 遍历输入目录中的所有图片
    for img_name in os.listdir(input_dir):
        if not img_name.lower().endswith(valid_extensions):
            continue
            
        # 读取图片
        img_path = os.path.join(input_dir, img_name)
        image = cv2.imread(img_path)
        if image is None:
            print(f"警告: 无法读取图片 {img_path}")
            continue
            
        # 转换为RGB（Albumentations需要RGB格式）
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        
        # 为每张图片生成4个增强版本
        for i in range(num_augmentations):
            # 应用增强
            augmented = transform(image=image)
            aug_image = augmented['image']
            
            # 转换回BGR以保存
            aug_image = cv2.cvtColor(aug_image, cv2.COLOR_RGB2BGR)
            
            # 生成输出文件名
            base_name = Path(img_name).stem
            output_name = f"{base_name}_noise_{i+1}{Path(img_name).suffix}"
            output_path = os.path.join(output_dir, output_name)
            
            # 保存增强后的图片
            cv2.imwrite(output_path, aug_image)
            print(f"已保存增强图片: {output_path}")

def main():
    # 配置输入和输出目录
    input_directory = "/home/wd/DejaVu/albumentations/HumanIntention_data6Augmentation/pullAug"  # 替换为您的输入图片目录
    output_directory = "/home/wd/DejaVu/albumentations/HumanIntention_data6Augmentation/pullAugNoise"  # 替换为您的输出图片目录
    
    # 每张图片生成4个增强版本
    augment_images(input_directory, output_directory, num_augmentations=4)

if __name__ == "__main__":
    main()

缺点：添加高斯噪声后，增广后的图片太模糊，如下图所示。所以第七部分我不再使用高斯噪声，而是选择了模拟数码相机在高 ISO 感光度设置下拍摄时产生的噪声。
在这里插入图片描述

七、添加ISONoise噪声增广

电脑：机械臂组 4090 工作站
文件路径：/home/wd/DejaVu/albumentations
文件名称：augmentDatasetNoiseV2.py

ISONoise 是 Albumentations 库中的一种图像增强方法，模拟数码相机在高 ISO 感光度设置下拍摄时产生的噪声。这种噪声通常出现在低光环境下，相机通过提高 ISO 感光度来增强亮度，但会引入随机像素扰动（噪声）和颜色偏移。这种增强方法有助于提高模型对真实世界中低质量或高 ISO 图像的鲁棒性，广泛应用于计算机视觉任务（如图像分类、目标检测）。

具体来说，ISONoise 的特点包括：

噪声类型：主要模拟颗粒状噪声（类似高斯噪声，但更贴近相机传感器噪声）以及轻微的颜色偏移。
应用场景：适合需要模拟现实世界中相机成像缺陷的场景，例如夜间拍摄、室内低光环境或设备质量较低的情况。
与高斯噪声的区别：GaussNoise 仅添加随机像素值扰动（通常是单色噪声），而 ISONoise 同时引入噪声和颜色偏移，更真实地模拟相机噪声。

代码解析：

a、intensity=(0.05, 0.2)：

含义：控制噪声的强度，表示噪声对图像像素的扰动程度。值越大，噪声越明显。
范围：intensity 是一个元组 (min, max)，表示噪声强度的随机范围。在每次增强时，Albumentations 会从 [0.05, 0.2] 中均匀随机选择一个值。
效果：
0.05：非常微弱的噪声，图像变化几乎不可察觉，适合需要最小化变化的场景。
0.2：适中的噪声，图像会出现轻微颗粒感，但不会显著影响内容识别。
相比原始代码中的 (0.1, 0.5)，这里降低了强度，使噪声更微妙。
技术细节：噪声强度影响像素值的随机扰动幅度，通常以标准化方式应用于图像的像素值（0-255范围）。较低的强度（如 0.05）意味着像素值变化较小，视觉上接近原始图像。

b、color_shift=(0.005, 0.1)：

含义：控制颜色偏移的程度，模拟高 ISO 拍摄时颜色失真的现象（如轻微偏红、偏蓝等）。
范围：color_shift 是一个元组 (min, max)，表示颜色偏移的随机范围。Albumentations 从 [0.005, 0.1] 中随机选择一个值。
效果：
0.005：极微小的颜色偏移，肉眼几乎无法察觉，适合微调增强。
0.1：轻微的颜色变化，可能表现为图像整体色调略偏暖或偏冷，但不会显著改变图像内容。
相比原始代码中的 (0.01, 0.3)，这里大幅降低了颜色偏移，保持图像与原始版本高度相似。
技术细节：颜色偏移通过调整图像的 RGB 通道值实现，可能涉及对每个通道应用独立的随机偏移。值较小（如 0.005）时，偏移幅度仅为像素值的微小百分比。

c、p=1.0：

含义：表示应用 ISONoise 变换的概率。p=1.0 意味着每次增强时都会应用该变换。
效果：由于 p=1.0，只要 ISONoise 在 A.OneOf 中被选中（由 A.OneOf 的逻辑决定），它将始终生效。
上下文：在您的代码中，ISONoise 是 A.OneOf 变换列表中的一项，A.OneOf 本身有 p=1.0，意味着每次增强会从 GaussNoise、ISONoise、RandomBrightnessContrast 和 HueSaturationValue 中随机选择一种变换，且所选变换（如 ISONoise）会以 100% 概率应用。

完整代码如下：

import os
import cv2
import albumentations as A
from pathlib import Path
import numpy as np

def create_augmentation_pipeline():
    """定义数据增强流水线，包含低强度噪声和微调变换"""
    return A.Compose([
	A.ISONoise(
                intensity=(0.05, 0.2),  # 降低噪声强度
                color_shift=(0.005, 0.1),  # 降低颜色偏移
                p=1.0
            ),
        A.OneOf([
            A.RandomBrightnessContrast(
                brightness_limit=0.1,  # 微小亮度调整
                contrast_limit=0.1,    # 微小对比度调整
                p=1.0
            ),
            A.HueSaturationValue(
                hue_shift_limit=10,      # 微小色调调整
                sat_shift_limit=10,      # 微小饱和度调整
                val_shift_limit=10,      # 微小明度调整
                p=1.0
            )
        ], p=1.0)
    ])

def augment_images(input_dir, output_dir, num_augmentations=4):
    """
    对指定目录中的图片进行数据增强，添加低强度变换
    Args:
        input_dir: 输入图片目录
        output_dir: 增强后图片保存目录
        num_augmentations: 每张图片生成的增强版本数量
    """
    # 确保输出目录存在
    Path(output_dir).mkdir(parents=True, exist_ok=True)
    
    # 创建增强流水线
    transform = create_augmentation_pipeline()
    
    # 支持的图片扩展名
    valid_extensions = ('.jpg', '.jpeg', '.png')
    
    # 遍历输入目录中的所有图片
    for img_name in os.listdir(input_dir):
        if not img_name.lower().endswith(valid_extensions):
            continue
            
        # 读取图片
        img_path = os.path.join(input_dir, img_name)
        image = cv2.imread(img_path)
        if image is None:
            print(f"警告: 无法读取图片 {img_path}")
            continue
            
        # 转换为RGB（Albumentations需要RGB格式）
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        
        # 为每张图片生成4个增强版本
        for i in range(num_augmentations):
            # 应用增强
            augmented = transform(image=image)
            aug_image = augmented['image']
            
            # 转换回BGR以保存
            aug_image = cv2.cvtColor(aug_image, cv2.COLOR_RGB2BGR)
            
            # 生成输出文件名
            base_name = Path(img_name).stem
            output_name = f"{base_name}_minimal_{i+1}{Path(img_name).suffix}"
            output_path = os.path.join(output_dir, output_name)
            
            # 保存增强后的图片
            cv2.imwrite(output_path, aug_image)
            print(f"已保存增强图片: {output_path}")

def main():
    # 配置输入和输出目录
    input_directory = "/home/wd/DejaVu/albumentations/HumanIntention_data6Augmentation/pullAug"
    output_directory = "/home/wd/DejaVu/albumentations/HumanIntention_data6Augmentation/pullAugNoise"
    
    # 每张图片生成4个增强版本
    augment_images(input_directory, output_directory, num_augmentations=4)

if __name__ == "__main__":
    main()

弃用版本

扩容数据集，修改颜色的程序，只能一张一张的修改，所以这个是弃用版本

import albumentations as A
import cv2
import os

# Define the output directory
output_dir = "/home/wd/DejaVu/albumentations/HumanIntention_data6Augmentation/downAug/"

# Ensure the output directory exists
os.makedirs(output_dir, exist_ok=True)

# Declare an augmentation pipeline
transform = A.Compose([
    A.OneOf([
            A.CLAHE(clip_limit=5, p=0.5),
            A.RandomBrightnessContrast(p=0.5),
            A.HueSaturationValue(hue_shift_limit=30, sat_shift_limit=30, val_shift_limit=20, always_apply=False, p=0.3)
        ], p=0.5)
])

# Read an image with OpenCV and convert it to the RGB colorspace
image = cv2.imread("/home/wd/DejaVu/albumentations/HumanIntention_data6Augmentation/downAug/frame_screenshot_08.11.20247.png")
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

# Augment an image
transformed = transform(image=image)
transformed_image = transformed["image"]

# Convert back to BGR for saving with OpenCV
transformed_image_bgr = cv2.cvtColor(transformed_image, cv2.COLOR_RGB2BGR)

# Define the output file path
output_path = os.path.join(output_dir, "frame_screenshot_08.11.20247c1.png")

# Save the augmented image
cv2.imwrite(output_path, transformed_image_bgr)

实验备份：

颜色增广1：3
旋转增广1：4
噪声增广1：4

训练集trainAug：

路径：/home/wd/DejaVu/albumentations/HumanIntention_data6Augmentation/trainAug
原始向上意图upAug：95张
原始向上意图的颜色增广upAugColor：285张
原始向上意图的噪声增广upAugNoise：380张
原始向上意图的颜色增广upAugRotation：380张
实验收集+增广向上意图共：1140张

原始向下意图downAug：92张
原始向下意图的颜色增广downAugColor：276张
原始向下意图的噪声增广downAugNoise：368张
原始向下意图的颜色增广downAugRotation：368张
实验收集+增广向下意图共：1104张

原始向左意图leftAug：92张
原始向左意图的颜色增广leftAugColor：276张
原始向左意图的噪声增广leftAugNoise：368张
原始向左意图的颜色增广leftAugRotation：368张
实验收集+增广向左意图共：1104张

原始向右意图rightAug：92张
原始向右意图的颜色增广rightAugColor：276张
原始向右意图的噪声增广rightAugNoise：368张
原始向右意图的颜色增广rightAugRotation：368张
实验收集+增广向右意图共：1104张

原始拉意图pullAug：185张
原始拉意图的颜色增广pullAugColor：555张
原始拉意图的噪声增广pullAugNoise：740张
原始拉意图的颜色增广pullAugRotation：740张
实验收集+增广拉意图共：2220张

原始推意图pushAug：97张
原始推意图的颜色增广pushAugColor：291张
原始推意图的噪声增广pushAugNoise：388张
原始推意图的颜色增广pushAugRotation：388张
实验收集+增广推意图共：1164张

测试集valAug：

路径：/home/wd/DejaVu/albumentations/HumanIntention_data6Augmentation/valAug
原始向上意图upAug：10张
原始向上意图的颜色增广upAugColor：30张
原始向上意图的噪声增广upAugNoise：40张
原始向上意图的颜色增广upAugRotation：40张
实验收集+增广向上意图共：120张

原始向下意图downAug：10张
原始向下意图的颜色增广downAugColor：30张
原始向下意图的噪声增广downAugNoise：40张
原始向下意图的颜色增广downAugRotation：40张
实验收集+增广向下意图共：120张

原始向左意图leftAug：10张
原始向左意图的颜色增广leftAugColor：30张
原始向左意图的噪声增广leftAugNoise：40张
原始向左意图的颜色增广leftAugRotation：40张
实验收集+增广向左意图共：120张

原始向右意图rightAug：10张
原始向右意图的颜色增广rightAugColor：30张
原始向右意图的噪声增广rightAugNoise：40张
原始向右意图的颜色增广rightAugRotation：40张
实验收集+增广向右意图共：120张

原始拉意图pullAug：20张
原始拉意图的颜色增广pullAugColor：60张
原始拉意图的噪声增广pullAugNoise：80张
原始拉意图的颜色增广pullAugRotation：80张
实验收集+增广拉意图共：240张

原始推意图pushAug：10张
原始推意图的颜色增广pushAugColor：30张
原始推意图的噪声增广pushAugNoise：40张
原始推意图的颜色增广pushAugRotation：40张
实验收集+增广推意图共：120张