如何在HuggingFace DETR目标检测模型中添加数据增强？

核心注意事项 #

• 不要重复归一化：你的processor已经配置了归一化参数，自定义变换里再添加归一化会导致数据分布异常，完全没必要。
• 增强要同步处理图像和标注：目标检测的增强必须保证图像变换后，对应的 bounding box 等标注也同步调整，所以要用torchvision.transforms.v2的API（需torchvision>=0.15），它支持对图像和COCO格式标注的联合变换。

修改后的数据集类示例 #

import torchvision
import os
from torchvision.transforms import v2

class CocoDetection(torchvision.datasets.CocoDetection):
    def __init__(self, img_folder, processor, train=True):
        ann_file = os.path.join(img_folder, "custom_train.json" if train else "custom_val.json")
        super(CocoDetection, self).__init__(img_folder, ann_file)
        self.processor = processor
        self.train = train
        
        # 定义训练集增强：缩放、HSV相关变换、随机翻转
        self.train_transforms = v2.Compose([
            # 随机缩放裁剪，对应尺度增强
            v2.RandomResizedCrop(size=(800, 800), scale=(0.5, 1.0)),
            # HSV变换：调整色调、饱和度、明度
            v2.ColorJitter(hue=0.1, saturation=0.2, brightness=0.2),
            # 随机水平翻转
            v2.RandomHorizontalFlip(p=0.5),
        ])
        
        # 验证集只做必要的图像转换（不做增强）
        self.val_transforms = v2.Compose([
            v2.Resize(size=(800, 800)),
        ])

    def __getitem__(self, idx):
        # 读取原始PIL图像和COCO格式的标注
        img, target = super(CocoDetection, self).__getitem__(idx)
        
        # 转换标注格式为transforms.v2能处理的结构
        # v2要求标注是包含'boxes'、'labels'键的字典
        transformed_target = {
            'boxes': [ann['bbox'] for ann in target],
            'labels': [ann['category_id'] for ann in target],
            'image_id': self.ids[idx]
        }
        # 将bbox从[x, y, w, h]转换为[xmin, ymin, xmax, ymax]，适配v2的格式
        for i in range(len(transformed_target['boxes'])):
            x, y, w, h = transformed_target['boxes'][i]
            transformed_target['boxes'][i] = [x, y, x + w, y + h]
        
        # 应用增强（训练集才做）
        if self.train:
            img, transformed_target = self.train_transforms(img, transformed_target)
        else:
            img, transformed_target = self.val_transforms(img, transformed_target)
        
        # 转换回processor需要的COCO标注格式
        coco_target = {
            'image_id': transformed_target['image_id'],
            'annotations': []
        }
        for box, label in zip(transformed_target['boxes'], transformed_target['labels']):
            xmin, ymin, xmax, ymax = box
            coco_target['annotations'].append({
                'bbox': [xmin, ymin, xmax - xmin, ymax - ymin],
                'category_id': label
            })
        
        # 用processor做剩下的预处理（resize、归一化、格式转换）
        encoding = self.processor(images=img, annotations=coco_target, return_tensors="pt")
        pixel_values = encoding["pixel_values"].squeeze()
        target = encoding["labels"][0]

        return pixel_values, target

关键细节说明 #

1. transforms.v2的使用：它会自动同步处理图像和标注的变换，比如缩放时自动调整bbox坐标，不用手动计算。
2. 标注格式转换：原始COCO标注的bbox是[x, y, w, h]格式，而transforms.v2需要[xmin, ymin, xmax, ymax]，所以要做一次格式转换，处理完后再转回去给processor。
3. 增强仅作用于训练集：验证集只做resize，避免引入不必要的噪声影响评估结果。
4. processor的作用不变：你的preprocessor_config.json无需修改，processor仍然负责最终的归一化和格式适配，增强只是在这之前的额外步骤。

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