在Pytorch OD中使用Albumentations进行增强

import torch
from torchvision import transforms
from albumentations import (
    Compose, RandomBrightnessContrast, HorizontalFlip, ShiftScaleRotate
)
from albumentations.pytorch import ToTensorV2

# 定义增强变换
transform = Compose([
    RandomBrightnessContrast(p=0.5),
    HorizontalFlip(p=0.5),
    ShiftScaleRotate(shift_limit=0.1, scale_limit=0.2, rotate_limit=20, p=0.5),
    ToTensorV2()
])

# 假设我们有一个图像和其对应的边界框标注
image = ...  # PIL Image or numpy array
bboxes = [[x1, y1, x2, y2], ...]  # List of bounding boxes in format [x1, y1, x2, y2]

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

# 现在augmented_image和augmented_bboxes可以用于训练模型

from albumentations import BboxParams

# 定义边界框参数
bbox_params = BboxParams(format='pascal_voc', min_area=0, min_visibility=0.1, label_fields=['class_labels'])

# 在Compose中加入bbox_params
transform = Compose([
    # ... 其他变换 ...
], bbox_params=bbox_params)