如何在Detectron 2中自定义或使用Albumentations数据增强？

一、集成Albumentations并实现概率增强 #

Detectron2的默认增强选项确实有限，直接集成Albumentations是最高效的解决方案——它原生支持OneOf概率选择策略，也能给单个增强操作设置触发概率。

1. 基础准备 #

先确保安装Albumentations：

pip install albumentations

2. 定义带概率策略的增强流水线 #

写一个包含OneOf逻辑的增强组合，示例如下：

import albumentations as A
from albumentations.pytorch import ToTensorV2

def get_albu_transforms():
    return A.Compose([
        # 单个增强的概率控制：每个操作按设置的概率触发
        A.RandomBrightnessContrast(p=0.5),
        A.GaussianBlur(p=0.3),
        # OneOf策略：从以下3个操作中随机选1个执行，整体触发概率为0.7
        A.OneOf([
            A.HorizontalFlip(p=1.0),
            A.VerticalFlip(p=1.0),
            A.RandomRotate90(p=1.0)
        ], p=0.7),
        # 必须保留的转换：转成Detectron2兼容的Tensor格式
        ToTensorV2(p=1.0)
    ], bbox_params=A.BboxParams(
        format="coco",  # 匹配Detectron2默认的COCO标注格式
        label_fields=["category_ids"]
    ))

注意bbox_params必须指定格式和标签字段，确保增强时bbox、mask能和图像同步变换。

3. 替换Detectron2的DataMapper #

写自定义mapper替换默认的DatasetMapper，把Albumentations的增强逻辑嵌入进去：

from detectron2.data import DatasetMapper
import numpy as np
import torch

class AlbumentationsMapper(DatasetMapper):
    def __init__(self, cfg, is_train=True):
        super().__init__(cfg, is_train=is_train)
        # 加载自定义增强流水线
        self.albu_transforms = get_albu_transforms()

    def __call__(self, dataset_dict):
        # 复制原始数据，避免修改原字典
        dataset_dict = super().__call__(dataset_dict)
        # 转换图像格式为Albumentations要求的HWC
        image = dataset_dict["image"].permute(1, 2, 0).numpy()
        # 提取标注信息
        bboxes = dataset_dict["instances"].gt_boxes.tensor.numpy()
        category_ids = dataset_dict["instances"].gt_classes.numpy()
        
        # 应用增强
        transformed = self.albu_transforms(
            image=image,
            bboxes=bboxes,
            category_ids=category_ids
        )
        
        # 更新增强后的图像和标注
        dataset_dict["image"] = transformed["image"]
        # 重新构建Instances对象
        from detectron2.structures import Instances, BoxMode
        instances = Instances(dataset_dict["instances"].image_size)
        instances.gt_boxes = BoxMode.convert(
            transformed["bboxes"],
            BoxMode.XYXY_ABS,
            BoxMode.XYXY_ABS  # Albumentations输出的bbox为XYXY格式，与Detectron2一致
        )
        instances.gt_classes = torch.tensor(transformed["category_ids"], dtype=torch.int64)
        # 若数据集包含mask，同步处理（示例）
        if "gt_masks" in dataset_dict["instances"]:
            masks = dataset_dict["instances"].gt_masks.numpy()
            transformed_masks = self.albu_transforms(image=image, masks=masks)["masks"]
            instances.gt_masks = torch.tensor(transformed_masks, dtype=torch.uint8)
        dataset_dict["instances"] = instances
        
        return dataset_dict

如果你的数据集有mask，需要额外添加mask的增强处理，Albumentations会自动完成同步变换。

二、查看增强后的图像结果 #

有两种简单的调试方式：

1. 在mapper中添加保存逻辑 #

在自定义mapper的__call__方法末尾，加入保存增强后图像的代码（仅训练调试用）：

import cv2
import os

if self.is_train:
    save_dir = "./aug_debug"
    os.makedirs(save_dir, exist_ok=True)
    # 把Tensor转成CV2可处理的格式
    img = transformed["image"].permute(1,2,0).numpy() * 255
    img = img.astype(np.uint8)
    # 绘制bbox
    for bbox in transformed["bboxes"]:
        x1, y1, x2, y2 = map(int, bbox)
        cv2.rectangle(img, (x1,y1), (x2,y2), (0,255,0), 2)
    # 保存图像
    cv2.imwrite(os.path.join(save_dir, f"{dataset_dict['image_id']}_aug.jpg"), img)

运行训练脚本后，./aug_debug目录下会生成带bbox的增强图像，可直观验证效果。

2. 用Detectron2可视化工具 #

调用Detectron2的Visualizer类直接在代码中可视化：

from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog
import cv2

# 获取数据集元数据
metadata = MetadataCatalog.get(cfg.DATASETS.TRAIN[0])
# 可视化增强结果
v = Visualizer(transformed["image"], metadata=metadata)
out = v.draw_instance_predictions(dataset_dict["instances"].to("cpu"))
# 显示图像（本地环境用cv2，Jupyter可直接display）
cv2.imshow("Augmented Image", out.get_image()[:, :, ::-1])
cv2.waitKey(0)

三、验证实现是否正确 #

要确保增强逻辑无误，需检查以下几点：

• 标注同步性：查看保存的增强图像，确认bbox/mask与目标物体位置完全对应，无偏移或变形。
• 概率触发验证：生成几十张增强图像，统计各操作的触发次数，是否符合设置的p参数（比如OneOf设p=0.7，应有约70%的图像执行了翻转/旋转）。
• 格式兼容性：启动训练脚本，检查是否有格式错误（如bbox格式不匹配、Tensor类型错误），若训练能正常启动且loss稳定，说明格式兼容。
• 效果验证：用相同模型参数，分别用默认增强和自定义增强训练，对比验证集mAP等指标，若指标有合理提升，间接证明增强有效。

内容的提问来源于stack exchange，提问作者Deshwal