通过albumentation对目标检测进行数据增强（简单直接）

news2024/9/19 6:38:45

albumentation官方文档看不懂？xml文件不知道如何操作？下面只需要修改部分代码即可上手使用

要使用这个方法之前需要按照albumentation这个库还有一些辅助库,自己看着来安装就行

pip install albumentation
pip install opencv-python
pip install json
pip install xml

import os
import cv2
import json
import xml.etree.ElementTree as ET
import albumentations as A

def data_aug(source_annotated, data_aug_annotated, times=3):
    transform = A.Compose([
            A.RandomBrightnessContrast(p=0.3),  # 随机亮度对比度
            A.OneOf([A.GaussNoise(), ], p=0.5),  # 高斯噪声
            A.OneOf([
                A.MotionBlur(p=0.2),  # 运动模糊
                A.MedianBlur(blur_limit=3, p=0.5),  # 中值模糊
                A.Blur(blur_limit=3, p=0.2),  # 模糊
            ], p=0.5),

            A.RandomBrightnessContrast(p=0.5),   
            A.RandomCrop(300, 300, p=0.5),  # 随机裁剪
            A.RandomRotate90(p=0.5),  # 随机旋转90度
            A.HorizontalFlip(p=0.5),  # 水平翻转
            A.VerticalFlip(p=0.5),  # 垂直翻转
        ], bbox_params=A.BboxParams(format='pascal_voc', min_area=0., min_visibility=0., label_fields=['class_labels']))

    for xml_file in os.listdir(source_annotated):
        if xml_file.endswith(".xml"):
            for nth in range(times):
                tree = ET.parse(os.path.join(source_annotated, xml_file))
                root = tree.getroot()

                image_file = root.find('filename').text
                image_path = os.path.join(source_annotated, image_file)
                image = cv2.imread(image_path)

                bboxes = []
                class_labels = []
                for obj in root.findall('object'):
                    label = obj.find('name').text
                    bbox = obj.find('bndbox')
                    xmin = int(bbox.find('xmin').text)
                    ymin = int(bbox.find('ymin').text)
                    xmax = int(bbox.find('xmax').text)
                    ymax = int(bbox.find('ymax').text)
                    bboxes.append([xmin, ymin, xmax, ymax])
                    class_labels.append(label)

                transformed = transform(image=image, bboxes=bboxes, class_labels=class_labels)
                transformed_image = transformed['image']
                transformed_bboxes = transformed['bboxes']

                save_base_name = os.path.splitext(image_file)[0] + "_" + str(nth)
                cv2.imwrite(os.path.join(data_aug_annotated, save_base_name + ".jpg"), transformed_image)

                # 更新XML文件中的变换数据
                root.find('filename').text = save_base_name + ".jpg"
                root.find('path').text = os.path.join(data_aug_annotated, save_base_name + ".jpg")

                for obj, tb in zip(root.findall('object'), transformed_bboxes):
                    bbox = obj.find('bndbox')
                    bbox.find('xmin').text = str(int(tb[0]))
                    bbox.find('ymin').text = str(int(tb[1]))
                    bbox.find('xmax').text = str(int(tb[2]))
                    bbox.find('ymax').text = str(int(tb[3]))

                tree.write(os.path.join(data_aug_annotated, save_base_name + ".xml"))

if __name__ == "__main__":
    source_annotated = "data_annotated/"  # 原始的图片、标签路径
    data_aug_annotated = "data_aug_annotated" # 处理后的图片、标签路径
    times = 3  
    data_aug(source_annotated, data_aug_annotated, times)

1.需要增强的图片标签文件存放形式

标签和文件放在统一文件夹下，同一图片和标签命名对应相同

2.修改需要增强的要求，如旋转，切割等

            A.RandomBrightnessContrast(p=0.3),  # 随机亮度对比度
            A.OneOf([A.GaussNoise(), ], p=0.5),  # 高斯噪声
            A.OneOf([
                A.MotionBlur(p=0.2),  # 运动模糊
                A.MedianBlur(blur_limit=3, p=0.5),  # 中值模糊
                A.Blur(blur_limit=3, p=0.2),  # 模糊
            ], p=0.5),

            A.RandomBrightnessContrast(p=0.5),   
            A.RandomCrop(300, 300, p=0.5),  # 随机裁剪
            A.RandomRotate90(p=0.5),  # 随机旋转90度
            A.HorizontalFlip(p=0.5),  # 水平翻转
            A.VerticalFlip(p=0.5),  # 垂直翻转

上述已经列出常见增强方式，可通过注释方式对方法进行取舍，代码中的p代表概率，即每一次增强有多少概率使用到这个方法，可通过自己的爱好进行调增，其他方法查询一下均可容易上手使用，这里不作过多的一一讲解了。