利用imgaug库对yoloDetect数据集和Pose数据集进行数据增强

引言

因为利用yolov8做实验的过程中想对数据集进行数据增强，找到了这个imgaug库。

这个库集成了非常多的常用的图像增强算法，包括平移、旋转等仿射变换，以及高斯模糊、边缘检测、cutout、dropout等操作，还支持雨、雪、雾等天气的模拟。使用起来非常方便，官方文档写的也简单易懂。

最方便的是，这个库的算法不仅支持对图像进行变换，而且还支持对选框Box以及关键点Keypoint进行变换。这大大节省了我们自己对已有的标签进行相应变换的工作量。

但是想使用这个库对Yolo格式的图片和标签进行变换，还是需要手动编写一些代码，不是很方便。所以我写了几个工具类，输入图像和标签的文件地址，就能自动输出变换后的图像和标签并自动保存。大大简化了yolo格式的目标检测数据集和Pose关键点数据集的数据增强过程。

代码与介绍

首先是图像增强助手的父类，定义了一些公共方法和一些工具函数，这个类不需要实例化，重点是后面的子类。

关于这些类的使用方法会放在第三部分工具类的使用。

import imgaug as ia
import imgaug.augmenters as iaa
from imgaug.augmentables.bbs import BoundingBox, BoundingBoxesOnImage
from imgaug.augmentables.kps import Keypoint, KeypointsOnImage
import os
import cv2 as cv

class DataAugHelper:
    def __init__(self, seq: iaa.Sequential):
        self.seq = seq

    def __call__(self):
        self.augData()

    def setSeq(self, seq: iaa.Sequential):
        self.seq = seq

    def setImgAndLabel(self, imgPath: str, labelPath: str):
        self.__checkFile(imgPath)
        self.__checkFile(labelPath)


    def augData(self):
        pass


    def __checkFile(self, filePath):
        if not os.path.exists(filePath):
            raise FileNotFoundError(f'{filePath} is not found.')

    def showOrigon(self):
        pass

    def showCompare(self):
        pass

    def showAug(self):
        pass

    def saveToFile(self, imgPath: str, labelPath: str):
        pass

---

首先是yolo目标检测数据集的数据增强助手，代码如下，没有增加什么新函数，主要是对父类的函数的重写。

from dataAugHelper import DataAugHelper
import imgaug as ia
import imgaug.augmenters as iaa
from imgaug.augmentables.bbs import BoundingBox, BoundingBoxesOnImage
from imgaug.augmentables.kps import Keypoint, KeypointsOnImage
import cv2 as cv


class YoloDetectDataAugHelper(DataAugHelper):
    def __init__(self, seq: iaa.Sequential):
        super().__init__(seq)

    def setImgAndLabel(self, imgPath: str, labelPath: str):
        super().setImgAndLabel(imgPath, labelPath)
        img = cv.imread