图像增强（albumentations.Compose()）

目录

• 1. 安装和导入
• 2. 基本概念
• 3. Compose 的基本用法
• 4. 常用的增强操作
• 5. 示例代码
  • 5.1 处理没有边界框的图像
  • 5.2 处理带有边界框的图像
• 6. 高级用法
  • 6.1 组合多个 Compose
  • 6.2 处理多张图像或掩码
  • 6.3 使用自定义变换
• 7. 注意事项

albumentations.Compose() 是 albumentations 库中的一个核心函数，用于将多个图像增强（augmentation）操作组合在一起，形成一个增强流水线（pipeline）。

1. 安装和导入

在使用 albumentations 之前，需要先安装该库。可以使用 pip 进行安装：

pip install albumentations

安装完成后，可以在 Python 脚本或 Jupyter Notebook 中导入：

import albumentations as A

2. 基本概念

• 增强流水线（Pipeline）：一系列图像增强操作按顺序应用于输入图像。
• 变换（Transforms）：具体的图像增强操作，如旋转、缩放、裁剪等。
• 概率（Probability）：每个变换可以设置一个概率，决定是否应用该变换。

albumentations.Compose() 的主要作用是将多个变换组合成一个流水线，按指定的顺序和概率依次应用于图像。

3. Compose 的基本用法

Compose 接受一个变换列表（或字典列表），并按顺序应用这些变换。它还支持按随机顺序应用变换、保存额外的信息等功能。

transform = A.Compose([
    A.HorizontalFlip(p=0.5),
    A.RandomBrightnessContrast(p=0.2),
    A.Rotate(limit=40, p=0.7),
])

在上面的例子中，Compose 将依次尝试水平翻转、随机亮度对比度调整和旋转，且每个变换都有各自的应用概率。

参数详解：

Compose 函数有多个参数，以下是主要参数的详细说明：

• transforms (list)：

  • 类型：列表
  • 说明：要组合的变换列表。每个变换可以是一个 albumentations 的变换实例或字典。

• bbox_params (albumentations.BboxParams, 可选)：

  • 说明：如果数据集中包含边界框（如目标检测任务），需要指定如何处理边界框的参数。
  • 主要参数：
    • format：边界框的格式，如 pascal_voc、coco 等。
    • label_fields：与边界框对应的标签字段名称。

• keypoint_params (albumentations.KeypointParams, 可选)：

  • 说明：如果数据集中包含关键点（如人脸关键点检测），需要指定关键点的参数。

• additional_targets (dict, 可选)：

  • 说明：用于处理额外的输入，如多张图像或图像及掩码（mask）。

• p (float, 默认 1.0)：

  • 说明：整体变换的应用概率。如果设置为 0.5，则整个流水线有 50% 的概率被应用。

4. 常用的增强操作

albumentations 提供了丰富的图像增强操作，以下是一些常用的变换：https://blog.csdn.net/qq_27039891/article/details/100795846

• 几何变换：

  • HorizontalFlip：水平翻转。
  • VerticalFlip：垂直翻转。
  • Rotate：旋转。
  • RandomCrop：随机裁剪。
  • Resize：调整大小。

• 颜色变换：

  • RandomBrightnessContrast：随机调整亮度和对比度。
  • HueSaturationValue：调整色调、饱和度和值。
  • RGBShift：随机移动 RGB 通道。

• 噪声和模糊：

  • GaussNoise：添加高斯噪声。
  • Blur：模糊。
  • GaussBlur：高斯模糊。

• 高级变换：

  • Cutout：随机遮挡图像区域。
  • ElasticTransform：弹性变换。
  • GridDistortion：网格扭曲。

5. 示例代码

5.1 处理没有边界框的图像

以下是一个使用 albumentations.Compose() 的完整示例，展示了如何构建增强流水线并应用于图像：

import albumentations as A
from albumentations.pytorch import ToTensorV2
import cv2
import matplotlib.pyplot as plt

# 定义增强流水线
transform = A.Compose([
    A.HorizontalFlip(p=0.5),
    A.RandomBrightnessContrast(p=0.2),
    A.Rotate(limit=40, p=0.7),
    A.Resize(256, 256),
    A.Normalize(mean=(0.485, 0.456, 0.406), 
                std=(0.229, 0.224, 0.225)),
    ToTensorV2(),
])

# 读取图像
image = cv2.imread('path_to_image.jpg')
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

# 应用增强
augmented = transform(image=image)
augmented_image = augmented['image']

# 显示增强后的图像
plt.imshow(augmented_image.permute(1, 2, 0))
plt.axis('off')
plt.show()

5.2 处理带有边界框的图像

如果任务涉及目标检测，需要同时处理图像和边界框：

transform = A.Compose([
    A.HorizontalFlip(p=0.5),
    A.RandomBrightnessContrast(p=0.2),
    A.Rotate(limit=40, p=0.7),
    A.Resize(256, 256),
], bbox_params=A.BboxParams(format='pascal_voc', label_fields=['labels']))

# 示例边界框和标签
bboxes = [[50, 30, 200, 180]]
labels = [1]

# 应用增强
augmented = transform(image=image, bboxes=bboxes, labels=labels)
augmented_image = augmented['image']
augmented_bboxes = augmented['bboxes']
augmented_labels = augmented['labels']

6. 高级用法

6.1 组合多个 Compose

可以嵌套使用 Compose 来构建更复杂的增强流水线。例如，针对训练和验证设置不同的增强策略：

train_transform = A.Compose([
    A.HorizontalFlip(p=0.5),
    A.ShiftScaleRotate(shift_limit=0.1, scale_limit=0.1, rotate_limit=15, p=0.5),
    A.OneOf([
        A.IAAAdditiveGaussianNoise(),
        A.GaussNoise(),
    ], p=0.2),
    A.OneOf([
        A.MotionBlur(p=0.2),
        A.MedianBlur(blur_limit=3, p=0.1),
        A.Blur(blur_limit=3, p=0.1),
    ], p=0.2),
    A.Normalize(mean=(0.485, 0.456, 0.406), 
                std=(0.229, 0.224, 0.225)),
    ToTensorV2(),
])

valid_transform = A.Compose([
    A.Resize(256, 256),
    A.Normalize(mean=(0.485, 0.456, 0.406), 
                std=(0.229, 0.224, 0.225)),
    ToTensorV2(),
])

6.2 处理多张图像或掩码

additional_targets 参数允许在同一个 Compose 中处理多张图像或图像及其对应的掩码：

transform = A.Compose([
    A.HorizontalFlip(p=0.5),
    A.Rotate(limit=40, p=0.7),
], additional_targets={'mask': 'mask'})

# 示例掩码
mask = cv2.imread('path_to_mask.png', cv2.IMREAD_GRAYSCALE)

# 应用增强
augmented = transform(image=image, mask=mask)
augmented_image = augmented['image']
augmented_mask = augmented['mask']

6.3 使用自定义变换

自定义变换并将其集成到 Compose 中：

from albumentations.core.transforms_interface import ImageOnlyTransform

class AddGaussianBlur(ImageOnlyTransform):
    def __init__(self, blur_limit=(3,7), always_apply=False, p=0.5):
        super().__init__(always_apply, p)
        self.blur_limit = blur_limit

    def apply(self, image, **params):
        return cv2.GaussianBlur(image, (self.blur_limit[0], self.blur_limit[1]), 0)

transform = A.Compose([
    AddGaussianBlur(blur_limit=(5,5), p=1.0),
    A.HorizontalFlip(p=0.5),
])

7. 注意事项

1. 顺序重要：Compose 中变换的顺序会影响最终效果。例如，先裁剪再调整大小与先调整大小再裁剪，结果可能不同。
2. 边界框和关键点：如果处理带有边界框或关键点的图像，必须正确设置 bbox_params 或 keypoint_params，以确保变换后这些信息的准确性。
3. 数据类型：albumentations 支持 NumPy 数组作为输入。确保图像在传入 Compose 之前已正确读取并转换为 RGB 格式（如使用 cv2.cvtColor）。
4. 性能优化：albumentations 以 C++ 实现，性能优越，但在大规模数据增强时，仍需注意内存和计算资源的使用。