使用albumentations图像增强后如何更新YOLO边界框坐标？

你遇到的标注未更新问题，核心原因有两个：

• 未将边界框参数与变换绑定：Albumentations的单个变换如果要同步处理边界框，必须通过A.Compose将变换与bbox_params关联，否则变换只会修改图像，完全忽略边界框。你当前直接调用单个变换时未传入bbox_params，导致Resize、RandomScale等操作不会更新标注坐标。
• 边界框格式错误：你将类别ID（class_id）包含在了bboxes列表中，但Albumentations要求YOLO格式的bboxes仅包含坐标信息（x_center, y_center, width, height），类别ID需要单独放在labels参数里。这种格式错误会导致变换无法正确解析和更新边界框。

修正后的代码 #

import os
import cv2
import albumentations as A

# 定义输入输出目录
input_image_dir = r'C:\Users\LEVEL51PC\Desktop\image_dataset\images'
input_label_dir = r'C:\Users\LEVEL51PC\Desktop\image_dataset\labels'
output_image_dir = r'C:\Users\LEVEL51PC\Desktop\augmented_dataset\images'
output_label_dir = r'C:\Users\LEVEL51PC\Desktop\augmented_dataset\labels'

# 创建输出目录（如果不存在）
os.makedirs(output_image_dir, exist_ok=True)
os.makedirs(output_label_dir, exist_ok=True)

# 定义边界框参数（YOLO格式）
bbox_params = A.BboxParams(
    format='yolo',
    label_fields=['labels']
)

# 用Compose包装每个变换，绑定bbox_params确保同步处理边界框
augmentation_techniques = [
    A.Compose([A.Resize(width=512, height=512)], bbox_params=bbox_params),
    A.Compose([A.RandomScale(scale_limit=(0.5, 2.0), p=0.5)], bbox_params=bbox_params),
    A.Compose([A.RandomRain(p=0.5)], bbox_params=bbox_params),
    A.Compose([A.ImageCompression(quality_lower=50, quality_upper=100, p=0.5)], bbox_params=bbox_params),
]

# 过滤仅保留图像文件
image_files = [f for f in os.listdir(input_image_dir) if f.lower().endswith(('.jpg', '.jpeg', '.png'))]

for image_file in image_files:
    # 加载图像
    image_path = os.path.join(input_image_dir, image_file)
    image = cv2.imread(image_path)
    if image is None:
        print(f"跳过无法加载的图像: {image_path}")
        continue

    # 加载对应的YOLO标注文件
    annotation_file = os.path.splitext(image_file)[0] + '.txt'
    annotation_path = os.path.join(input_label_dir, annotation_file)
    if not os.path.exists(annotation_path):
        print(f"跳过无标注的图像: {image_file}")
        continue

    # 解析标注：分离类别ID与边界框坐标
    boxes = []
    labels = []
    with open(annotation_path, 'r') as f:
        for line in f.readlines():
            line = line.strip()
            if not line:
                continue
            parts = list(map(float, line.split()))
            labels.append(int(parts[0]))
            boxes.append(parts[1:])  # 仅保留坐标信息

    # 逐个应用增强变换
    for i, transform in enumerate(augmentation_techniques):
        augmented = transform(image=image, bboxes=boxes, labels=labels)
        augmented_image = augmented['image']
        augmented_boxes = augmented['bboxes']
        augmented_labels = augmented['labels']

        # 生成增强后的文件名
        base_name = os.path.splitext(image_file)[0]
        augmented_filename = f"{base_name}_{i}_aug"

        # 保存增强图像
        output_image_path = os.path.join(output_image_dir, f"{augmented_filename}.jpg")
        cv2.imwrite(output_image_path, augmented_image)

        # 保存更新后的标注
        augmented_annotation_path = os.path.join(output_label_dir, f"{augmented_filename}.txt")
        with open(augmented_annotation_path, 'w') as f:
            for label, box in zip(augmented_labels, augmented_boxes):
                x_center, y_center, width, height = box
                f.write(f"{label} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}\n")

print("增强完成")

关键修改说明 #

1. 用Compose绑定变换与bbox_params：每个增强变换都通过A.Compose包装，确保变换在处理图像的同时，按照YOLO格式更新边界框。
2. 分离labels与bboxes：将类别ID单独存入labels列表，bboxes仅保留坐标信息，符合Albumentations的输入要求。
3. 增加文件存在性检查：避免因图像或标注文件缺失导致程序崩溃。
4. 保留坐标精度：保存标注时保留6位小数，防止坐标信息丢失。

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