基于albumentations的数据增强与标注处理实战：解决CV训练中的样本瓶颈

1次阅读

共计 2730 个字符，预计需要花费 7 分钟才能阅读完成。

在计算机视觉任务中，数据增强是提升模型泛化能力的关键技术。传统方法如 OpenCV 在进行几何变换（旋转、缩放、裁剪等）时，常出现标注与图像不同步的问题，尤其是对于 Bounding Box 和 Keypoint 这类结构化标注数据。常见的标注错位现象包括：

旋转后的 Bounding Box 未同步调整角度，导致包含无关背景
随机裁剪时未考虑标注框的完整性，造成目标被意外截断
仿射变换后 Keypoint 坐标未经过齐次坐标矩阵运算，产生像素级偏移

这些问题在小样本场景下尤为突出，错误的增强数据反而会误导模型训练。

特性	albumentations	torchvision	imgaug
标注同步支持	原生支持	需手动实现	插件式支持
GPU 加速	通过 CuPy 支持	部分操作支持	不支持
API 设计	函数式 + 面向对象	面向对象	命令式
变换类型	120+	40+	200+
随机种子管理	全局控制	独立控制	独立控制

import albumentations as A

transform = A.Compose([A.RandomRotate90(),
    A.Flip(p=0.5),
    A.RandomBrightnessContrast(p=0.2),
    A.HueSaturationValue(hue_shift_limit=20, sat_shift_limit=30, val_shift_limit=20, p=0.5)
], bbox_params=A.BboxParams(format='coco', min_visibility=0.3))

import json

def load_coco_annotations(json_path):
    with open(json_path) as f:
        data = json.load(f)

    images = {item['id']: item for item in data['images']}
    annotations = {img_id: [] for img_id in images.keys()}

    for ann in data['annotations']:
        annotations[ann['image_id']].append(ann['bbox'] + [ann['category_id']])  # [x,y,w,h,class_id]

    return images, annotations

from torch.utils.data import Dataset
import cv2

class AlbumentationsDataset(Dataset):
    def __init__(self, image_dir, annotations, transform=None):
        self.image_dir = image_dir
        self.annotations = annotations
        self.transform = transform
        self.image_ids = list(annotations.keys())

    def __len__(self):
        return len(self.image_ids)

    def __getitem__(self, idx):
        image_id = self.image_ids[idx]
        image_info = self.annotations['images'][image_id]

        # 加载图像
        image = cv2.imread(f"{self.image_dir}/{image_info['file_name']}")
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

        # 获取标注
        bboxes = self.annotations['annotations'][image_id]
        labels = [box[-1] for box in bboxes]
        boxes = [box[:4] for box in bboxes]

        # 应用变换
        if self.transform:
            transformed = self.transform(image=image, bboxes=boxes, class_labels=labels)
            image = transformed['image']
            boxes = transformed['bboxes']
            labels = transformed['class_labels']

        # 转换为 tensor
        # ...（后续处理）return image, {'boxes': boxes, 'labels': labels}

import timeit

setup = '''
import albumentations as A
import numpy as np

transform = A.Compose([A.HorizontalFlip(p=0.5),
    A.RandomBrightnessContrast(p=0.2),
    A.ShiftScaleRotate(shift_limit=0.1, scale_limit=0.1, rotate_limit=15, p=0.5)
])
image = np.random.randint(0, 256, (512, 512, 3), dtype=np.uint8)
'''timeit.timeit('transform(image=image)', setup=setup, number=1000)