语义分割离线数据增强——Albumentations实现

该脚本实现了一个图像及其对应掩模的离线数据增强流程，旨在通过引入多种变换增强图像多样性，以提高深度学习模型的泛化能力。主要技术要点如下：

1.使用 Albumentations 库进行数据增强

：Albumentations 是一个强大的 Python 库，专门用于图像预处理和数据增强。它提供了丰富的图像变换操作，支持多种深度学习框架。在本脚本中，通过定义 A.Compose 对象 aug 集成了多种增强操作。

2.定义增强配置：

A.HorizontalFlip(p=0.8) 和 A.VerticalFlip(p=0.7)：分别以 80% 和 70% 的概率水平翻转和垂直翻转图像及掩模，增加图像的方向性变化。
A.Resize(height=512, width=512, interpolation=cv.INTER_CUBIC, always_apply=False, p=1)：始终将图像和掩模统一缩放到 512x512 大小，使用三次插值法保持图像质量。
A.ShiftScaleRotate(shift_limit=0.0625, scale_limit=0.2, rotate_limit=15, p=0.8)：以 80% 的概率对图像和掩模进行随机平移（最大移动图像尺寸的 1/16）、缩放（最大放大20%或缩小至80%）和旋转（最大±15度），模拟真实世界中的视角变化和物体位置、大小的差异。
A.Transpose(always_apply=False, p=1)：始终对图像和掩模进行转置（水平与垂直方向交换），增加图像的空间布局变化。
A.GridDistortion(num_steps=10, distort_limit=0.3, border_mode=cv.BORDER_CONSTANT, always_apply=False, p=1)：始终对图像和掩模应用网格失真变形，模拟镜头畸变或物体表面不规则情况。

3.数据集准备与目录结构：

image_dir、mask_dir 存储原始图像和掩模文件，分别对应图像和其对应的彩色掩模。
aug_image_dir、aug_mask_dir 分别为增强后图像和掩模的输出目录，确保其存在并创建（若不存在）。

4.遍历图像与掩模文件：

从 image_dir 和 mask_dir 中获取相同数量的图像和掩模文件列表，确保二者一一对应。
使用 zip 函数同时遍历 image_files 和 mask_files，对每一对图像和掩模进行后续处理。

5.生成增强版本：

读取指定路径下的原始图像和彩色掩模。
对每对图像-掩模组合应用增强配置 aug 进行100次随机增强。
从增强结果中提取增强后的图像 image_aug 和掩模 mask_aug。

6.保存增强结果：

构建增强后图像和掩模的新文件名，格式为原始文件名（去除扩展名）+ 序号（0~99）+ 扩展名（.jpg 或 .png）。将增强后的图像和掩模保存到对应的输出目录 aug_image_dir 和 aug_mask_dir。

总之，此脚本利用 Albumentations 库对给定图像及其彩色掩模进行一系列几何变换和像素级失真操作，生成大量增强后的图像-掩模对，旨在提升模型训练时的数据多样性，进而增强模型的泛化能力。每个原始图像-掩模对将产生100个不同的增强版本，增强了数据集的规模和复杂性。

import os
import cv2 as cv
import albumentations as A

# 定义增强配置
aug = A.Compose([
    A.HorizontalFlip(p=0.8),
    A.VerticalFlip(p=0.7),
    A.Resize(height=512, width=512, interpolation=cv.INTER_CUBIC, always_apply=False, p=1),
    A.ShiftScaleRotate(shift_limit=0.0625, scale_limit=0.2, rotate_limit=15, p=0.8),
    A.Transpose(always_apply=False, p=1),
    A.GridDistortion(num_steps=10, distort_limit=0.3, border_mode=cv.BORDER_CONSTANT, always_apply=False, p=1),
])

# load dataset
image_dir = './dataset/image'
mask_dir = './dataset/mask'
aug_image_dir = './aug/image'
aug_mask_dir = './aug/mask'

# 确保输出目录存在
os.makedirs(aug_image_dir, exist_ok=True)
os.makedirs(aug_mask_dir, exist_ok=True)

# 获取图像和掩模文件列表
image_files = [f for f in os.listdir(image_dir) if f.endswith('.jpg') or f.endswith('.png')]
mask_files = [f for f in os.listdir(mask_dir) if f.endswith('.jpg') or f.endswith('.png')]

assert len(image_files) == len(mask_files), "Number of images and masks don't match"

for image_file, mask_file in zip(image_files, mask_files):
    # 读取原始图像和掩模
    image_path = os.path.join(image_dir, image_file)
    mask_path = os.path.join(mask_dir, mask_file)
    image = cv.imread(image_path)
    mask = cv.imread(mask_path, cv.IMREAD_COLOR)  # 读取彩色掩模

    # 为当前图像和掩模生成多个增强版本（共100个）
    for i in range(100):
        augmented = aug(image=image, mask=mask)
        image_aug = augmented["image"]
        mask_aug = augmented["mask"]

        # 为增强结果构建新的文件名
        output_image_name = f'{image_file[:-4]}_{i}.jpg'  # 假设原始图像为 .jpg 格式
        output_mask_name = f'{mask_file[:-4]}_{i}.png'  # 假设原始掩模为 .png 格式

        # 保存增强后的图像和掩模
        output_image_path = os.path.join(aug_image_dir, output_image_name)
        output_mask_path = os.path.join(aug_mask_dir, output_mask_name)
        cv.imwrite(output_image_path, image_aug)
        cv.imwrite(output_mask_path, mask_aug)