yolo imgaug数据增强 标签同时也增强

原文：https://blog.csdn.net/m0_37940759/article/details/115212083

完整代码：

import xml.etree.ElementTree as ET
import pickle
import os
from os import getcwd
import numpy as np
from PIL import Image
import shutil
import matplotlib.pyplot as plt

import imgaug as ia
from imgaug import augmenters as iaa


ia.seed(1)


def read_xml_annotation(root, image_id):
    in_file = open(os.path.join(root, image_id))
    tree = ET.parse(in_file)
    root = tree.getroot()
    bndboxlist = []

    for object in root.findall('object'):  # 找到root节点下的所有country节点
        bndbox = object.find('bndbox')  # 子节点下节点rank的值

        xmin = int(bndbox.find('xmin').text)
        xmax = int(bndbox.find('xmax').text)
        ymin = int(bndbox.find('ymin').text)
        ymax = int(bndbox.find('ymax').text)
        # print(xmin,ymin,xmax,ymax)
        bndboxlist.append([xmin, ymin, xmax, ymax])
        # print(bndboxlist)

    bndbox = root.find('object').find('bndbox')
    return bndboxlist


# (506.0000, 330.0000, 528.0000, 348.0000) -> (520.4747, 381.5080, 540.5596, 398.6603)
def change_xml_annotation(root, image_id, new_target):
    new_xmin = new_target[0]
    new_ymin = new_target[1]
    new_xmax = new_target[2]
    new_ymax = new_target[3]

    in_file = open(os.path.join(root, str(image_id) + '.xml'))  # 这里root分别由两个意思
    tree = ET.parse(in_file)
    xmlroot = tree.getroot()
    object = xmlroot.find('object')
    bndbox = object.find('bndbox')
    xmin = bndbox.find('xmin')
    xmin.text = str(new_xmin)
    ymin = bndbox.find('ymin')
    ymin.text = str(new_ymin)
    xmax = bndbox.find('xmax')
    xmax.text = str(new_xmax)
    ymax = bndbox.find('ymax')
    ymax.text = str(new_ymax)
    tree.write(os.path.join(root, str("%06d" % (str(id) + '.xml'))))


def change_xml_list_annotation(root, image_id, new_target, saveroot, id):
    in_file = open(os.path.join(root, str(image_id) + '.xml'))  # 这里root分别由两个意思
    tree = ET.parse(in_file)
    #修改xml中的filename
    elem = tree.find('filename')
    elem.text = (str(id) + '.jpg')
    xmlroot = tree.getroot()
    #修改xml中的path
    elem = tree.find('path')
    elem.text = ('VOCdevkit\VOC2007\JPEGImages\\' + str(id) + '.jpg')
    xmlroot = tree.getroot()

    index = 0
    for object in xmlroot.findall('object'):  # 找到root节点下的所有country节点
        bndbox = object.find('bndbox')  # 子节点下节点rank的值

        # xmin = int(bndbox.find('xmin').text)
        # xmax = int(bndbox.find('xmax').text)
        # ymin = int(bndbox.find('ymin').text)
        # ymax = int(bndbox.find('ymax').text)

        new_xmin = new_target[index][0]
        new_ymin = new_target[index][1]
        new_xmax = new_target[index][2]
        new_ymax = new_target[index][3]

        xmin = bndbox.find('xmin')
        xmin.text = str(new_xmin)
        ymin = bndbox.find('ymin')
        ymin.text = str(new_ymin)
        xmax = bndbox.find('xmax')
        xmax.text = str(new_xmax)
        ymax = bndbox.find('ymax')
        ymax.text = str(new_ymax)

        index = index + 1

    tree.write(os.path.join(saveroot, str(id + '.xml')))


def mkdir(path):
    # 去除首位空格
    path = path.strip()
    # 去除尾部 \ 符号
    path = path.rstrip("\\")
    # 判断路径是否存在
    # 存在     True
    # 不存在   False
    isExists = os.path.exists(path)
    # 判断结果
    if not isExists:
        # 如果不存在则创建目录
        # 创建目录操作函数
        os.makedirs(path)
        print(path + ' 创建成功')
        return True
    else:
        # 如果目录存在则不创建，并提示目录已存在
        print(path + ' 目录已存在')
        return False


if __name__ == "__main__":

    IMG_DIR = "./VOCdevkit/VOC2007/JPEGImages"
    XML_DIR = "./VOCdevkit/VOC2007/Annotations"

    AUG_XML_DIR = "./AUG/Annotations"  # 存储增强后的XML文件夹路径
    try:
        shutil.rmtree(AUG_XML_DIR)
    except FileNotFoundError as e:
        a = 1
    mkdir(AUG_XML_DIR)

    AUG_IMG_DIR = "./AUG/JPEGImages"  # 存储增强后的影像文件夹路径
    try:
        shutil.rmtree(AUG_IMG_DIR)
    except FileNotFoundError as e:
        a = 1
    mkdir(AUG_IMG_DIR)

    AUGLOOP = 2  # 每张影像增强的数量

    boxes_img_aug_list = []
    new_bndbox = []
    new_bndbox_list = []

    # 影像增强
    seq = iaa.Sequential([
        iaa.Invert(0.5),  # 50%的概率反转像素的强度，即原来的强度为v那么现在的就是255-v
        iaa.Fliplr(0.5),  # 镜像    # 对50%的图像进行镜像翻转
        iaa.Multiply((1.2, 1.5)),  # change brightness, doesn't affect BBs # 像素乘上1.2或者1.5之间的数字.
        iaa.GaussianBlur(sigma=(0, 3.0)),  # iaa.GaussianBlur(0.5),
        iaa.Affine(
            translate_px={"x": 15, "y": 15},
            scale=(0.8, 0.95),
        )  # translate by 40/60px on x/y axis, and scale to 50-70%, affects BBs
    ])

    for root, sub_folders, files in os.walk(XML_DIR):

        for name in files:

            bndbox = read_xml_annotation(XML_DIR, name)
            shutil.copy(os.path.join(XML_DIR, name), AUG_XML_DIR)
            shutil.copy(os.path.join(IMG_DIR, name[:-4] + '.jpg'), AUG_IMG_DIR)

            for epoch in range(AUGLOOP):
                seq_det = seq.to_deterministic()  # 保持坐标和图像同步改变，而不是随机
                # 读取图片
                img = Image.open(os.path.join(IMG_DIR, name[:-4] + '.jpg'))
                # sp = img.size
                img = np.asarray(img)
                # bndbox 坐标增强
                for i in range(len(bndbox)):
                    bbs = ia.BoundingBoxesOnImage([
                        ia.BoundingBox(x1=bndbox[i][0], y1=bndbox[i][1], x2=bndbox[i][2], y2=bndbox[i][3]),
                    ], shape=img.shape)

                    bbs_aug = seq_det.augment_bounding_boxes([bbs])[0]
                    boxes_img_aug_list.append(bbs_aug)

                    # new_bndbox_list:[[x1,y1,x2,y2],...[],[]]
                    n_x1 = int(max(1, min(img.shape[1], bbs_aug.bounding_boxes[0].x1)))
                    n_y1 = int(max(1, min(img.shape[0], bbs_aug.bounding_boxes[0].y1)))
                    n_x2 = int(max(1, min(img.shape[1], bbs_aug.bounding_boxes[0].x2)))
                    n_y2 = int(max(1, min(img.shape[0], bbs_aug.bounding_boxes[0].y2)))
                    if n_x1 == 1 and n_x1 == n_x2:
                        n_x2 += 1
                    if n_y1 == 1 and n_y2 == n_y1:
                        n_y2 += 1
                    if n_x1 >= n_x2 or n_y1 >= n_y2:
                        print('error', name)
                    new_bndbox_list.append([n_x1, n_y1, n_x2, n_y2])
                # 存储变化后的图片
                image_aug = seq_det.augment_images([img])[0]
                path = os.path.join(AUG_IMG_DIR,
                                    str(str(name[:-4]) + '_' + str(epoch)) + '.jpg')
                image_auged = bbs.draw_on_image(image_aug, thickness=0)
                Image.fromarray(image_auged).save(path)

                # 存储变化后的XML
                change_xml_list_annotation(XML_DIR, name[:-4], new_bndbox_list, AUG_XML_DIR,
                                           str(name[:-4]) + '_' + str(epoch))
                # print(str(str(name[:-4]) + '_' + str(epoch)) + '.jpg')
                new_bndbox_list = []

Voc数据集对有标签的数据集数据增强
安装依赖库和imgaug库
Bounding Boxes实现
读取原影像bounding boxes坐标
生成变换序列
bounding box 变化后坐标计算
使用示例
数据准备
设置文件路径
设置增强次数
设置增强参数
修改xml文件中filename和path
输出
完整代码
安装依赖库和imgaug库
在训练yolo的时候，对已有数据集已经打好了标签，想要进行数据的增强（数据扩增），可以通过imgaug实现对图片和标签中的boundingbox同时变换。
imgaug使用文档
代码下载：https://github.com/CodingWZP/image_augmentation
首先，安装依赖库。

pip install six numpy scipy matplotlib scikit-image opencv-python imageio
1
安装imgaug

pip install imgaug
1
Bounding Boxes实现
读取原影像bounding boxes坐标
读取xml文件并使用ElementTree对xml文件进行解析，找到每个object的坐标值。

def change_xml_annotation(root, image_id, new_target):
new_xmin = new_target[0]
new_ymin = new_target[1]
new_xmax = new_target[2]
new_ymax = new_target[3]

in_file = open(os.path.join(root, str(image_id) + '.xml'))  # 这里root分别由两个意思
tree = ET.parse(in_file)
xmlroot = tree.getroot()
object = xmlroot.find('object')
bndbox = object.find('bndbox')
xmin = bndbox.find('xmin')
xmin.text = str(new_xmin)
ymin = bndbox.find('ymin')
ymin.text = str(new_ymin)
xmax = bndbox.find('xmax')
xmax.text = str(new_xmax)
ymax = bndbox.find('ymax')
ymax.text = str(new_ymax)
tree.write(os.path.join(root, str("%06d" % (str(id) + '.xml'))))

生成变换序列
产生一个处理图片的Sequential。

影像增强

seq = iaa.Sequential([
    iaa.Invert(0.5),
    iaa.Fliplr(0.5),  # 镜像
    iaa.Multiply((1.2, 1.5)),  # change brightness, doesn't affect BBs
    iaa.GaussianBlur(sigma=(0, 3.0)),  # iaa.GaussianBlur(0.5),
    iaa.Affine(
        translate_px={"x": 15, "y": 15},
        scale=(0.8, 0.95),
    )  # translate by 40/60px on x/y axis, and scale to 50-70%, affects BBs
])

bounding box 变化后坐标计算
先读取该影像对应xml文件，获取所有目标的bounding boxes，然后依次计算每个box变化后的坐标。

seq_det = seq.to_deterministic() # 保持坐标和图像同步改变，而不是随机

读取图片

img = Image.open(os.path.join(IMG_DIR, name[:-4] + ‘.jpg’))

sp = img.size

img = np.asarray(img)

bndbox 坐标增强

for i in range(len(bndbox)):
bbs = ia.BoundingBoxesOnImage([
ia.BoundingBox(x1=bndbox[i][0], y1=bndbox[i][1], x2=bndbox[i][2], y2=bndbox[i][3]),
], shape=img.shape)

bbs_aug = seq_det.augment_bounding_boxes([bbs])[0]
boxes_img_aug_list.append(bbs_aug)

# new_bndbox_list:[[x1,y1,x2,y2],...[],[]]
n_x1 = int(max(1, min(img.shape[1], bbs_aug.bounding_boxes[0].x1)))
n_y1 = int(max(1, min(img.shape[0], bbs_aug.bounding_boxes[0].y1)))
n_x2 = int(max(1, min(img.shape[1], bbs_aug.bounding_boxes[0].x2)))
n_y2 = int(max(1, min(img.shape[0], bbs_aug.bounding_boxes[0].y2)))
if n_x1 == 1 and n_x1 == n_x2:
    n_x2 += 1
if n_y1 == 1 and n_y2 == n_y1:
    n_y2 += 1
if n_x1 >= n_x2 or n_y1 >= n_y2:
    print('error', name)
new_bndbox_list.append([n_x1, n_y1, n_x2, n_y2])

存储变化后的图片

image_aug = seq_det.augment_images([img])[0]
path = os.path.join(AUG_IMG_DIR,
str(str(name[:-4]) + ‘_’ + str(epoch)) + ‘.jpg’)
image_auged = bbs.draw_on_image(image_aug, thickness=0)
Image.fromarray(image_auged).save(path)

存储变化后的XML

change_xml_list_annotation(XML_DIR, name[:-4], new_bndbox_list, AUG_XML_DIR,
str(name[:-4]) + ‘_’ + str(epoch))

print(str(str(name[:-4]) + ‘_’ + str(epoch)) + ‘.jpg’)

new_bndbox_list = []

使用示例
数据准备
输入数据为两个文件夹一个是需要增强的影像数据（JPEGImages），一个是对应的xml文件（Annotations）。注意：影像文件名需和xml文件名相对应！

设置文件路径
IMG_DIR = “./VOCdevkit/VOC2007/JPEGImages”
XML_DIR = “./VOCdevkit/VOC2007/Annotations”

AUG_XML_DIR = “./AUG/Annotations” # 存储增强后的XML文件夹路径
try:
shutil.rmtree(AUG_XML_DIR)
except FileNotFoundError as e:
a = 1
mkdir(AUG_XML_DIR)

AUG_IMG_DIR = “./AUG/JPEGImages” # 存储增强后的影像文件夹路径
try:
shutil.rmtree(AUG_IMG_DIR)
except FileNotFoundError as e:
a = 1
mkdir(AUG_IMG_DIR)

设置增强次数
AUGLOOP = 10 # 每张影像增强的数量
1
设置增强参数
通过修改Sequential函数参数进行设置，具体设置参考imgaug使用文档

seq = iaa.Sequential([
    iaa.Invert(0.5),
    iaa.Fliplr(0.5),  # 镜像
    iaa.Multiply((1.2, 1.5)),  # change brightness, doesn't affect BBs
    iaa.GaussianBlur(sigma=(0, 3.0)),  # iaa.GaussianBlur(0.5),
    iaa.Affine(
        translate_px={"x": 15, "y": 15},
        scale=(0.8, 0.95),
    )  # translate by 40/60px on x/y axis, and scale to 50-70%, affects BBs
])

修改xml文件中filename和path
tree = ET.parse(in_file)
#修改xml中的filename
elem = tree.find(‘filename’)
elem.text = (str(id) + ‘.jpg’)
xmlroot = tree.getroot()
#修改xml中的path
elem = tree.find(‘path’)
elem.text = (‘C:\working\yolo3\VOCdevkit\VOC2007\JPEGImages\’ + str(id) + ‘.jpg’)
xmlroot = tree.getroot()
输出
运行augmentation.py ，运行结束后即可得到增强的影像和对应的xml文件夹。
————————————————
版权声明：本文为CSDN博主「CodingWZP」的原创文章，遵循CC 4.0 BY-SA版权协议，转载请附上原文出处链接及本声明。
原文链接：https://blog.csdn.net/m0_37940759/article/details/115212083

其他案例：imgaug数据增强

import os
from imgaug import augmenters as iaa
import cv2
import sometimes
seq = iaa.Sequential(
    [
        iaa.Fliplr(0.5),  # 对50%的图像进行镜像翻转
        iaa.Flipud(0.2),  # 对20%的图像做左右翻转
 
        # sometimes(iaa.Crop(percent=(0, 0.1))),
        # 这里沿袭我们上面提到的sometimes，对随机的一部分图像做crop操作
        # crop的幅度为0到10%
 
        # sometimes(iaa.Affine(  # 对一部分图像做仿射变换
        #     scale={"x": (0.8, 1.2), "y": (0.8, 1.2)},  # 图像缩放为80%到120%之间
        #     translate_percent={"x": (-0.2, 0.2), "y": (-0.2, 0.2)},  # 平移±20%之间
        #     rotate=(-45, 45),  # 旋转±45度之间
        #     shear=(-16, 16),  # 剪切变换±16度，（矩形变平行四边形）
        #     order=[0, 1],  # 使用最邻近差值或者双线性差值
        #     cval=(0, 255),  # 全白全黑填充
        #     mode=iaa.ALL  # 定义填充图像外区域的方法
        # )),
 
        # 使用下面的0个到5个之间的方法去增强图像。注意SomeOf的用法
        iaa.SomeOf((0, 5),
                   [
                       # 将部分图像进行超像素的表示。o(╥﹏╥)o用超像素增强作者还是第一次见，比较孤陋寡闻
                       # sometimes(
                       #     iaa.Superpixels(
                       #         p_replace=(0, 1.0),
                       #         n_segments=(20, 200)
                       #     )
                       # ),
 
                       # 用高斯模糊，均值模糊，中值模糊中的一种增强。注意OneOf的用法
                       iaa.OneOf([
                           iaa.GaussianBlur((0, 3.0)),
                           iaa.AverageBlur(k=(2, 7)),  # 核大小2~7之间，k=((5, 7), (1, 3))时，核高度5~7，宽度1~3
                           iaa.MedianBlur(k=(3, 11)),
                       ]),
 
                       # 锐化处理
                       iaa.Sharpen(alpha=(0, 1.0), lightness=(0.75, 1.5)),
 
                       # 浮雕效果
                       iaa.Emboss(alpha=(0, 1.0), strength=(0, 2.0)),
 
                       # # 边缘检测，将检测到的赋值0或者255然后叠在原图上
                       # sometimes(iaa.OneOf([
                       #     iaa.EdgeDetect(alpha=(0, 0.7)),
                       #     iaa.DirectedEdgeDetect(
                       #         alpha=(0, 0.7), direction=(0.0, 1.0)
                       #     ),
                       # ])),
 
                       # 加入高斯噪声
                       iaa.AdditiveGaussianNoise(
                           loc=0, scale=(0.0, 0.05 * 255), per_channel=0.5
                       ),
 
                       # 将1%到10%的像素设置为黑色
                       # 或者将3%到15%的像素用原图大小2%到5%的黑色方块覆盖
                       iaa.OneOf([
                           iaa.Dropout((0.01, 0.1), per_channel=0.5),
                           iaa.CoarseDropout(
                               (0.03, 0.15), size_percent=(0.02, 0.05),
                               per_channel=0.2
                           ),
                       ]),
 
                       # 5%的概率反转像素的强度，即原来的强度为v那么现在的就是255-v
                       iaa.Invert(0.05, per_channel=True),
 
                       # 每个像素随机加减-10到10之间的数
                       iaa.Add((-10, 10), per_channel=0.5),
 
                       # 像素乘上0.5或者1.5之间的数字.
                       iaa.Multiply((0.5, 1.5), per_channel=0.5),
 
                       # 将整个图像的对比度变为原来的一半或者二倍
                       iaa.ContrastNormalization((0.5, 2.0), per_channel=0.5),
 
                       # 将RGB变成灰度图然后乘alpha加在原图上
                       iaa.Grayscale(alpha=(0.0, 1.0)),
 
                       # 把像素移动到周围的地方。这个方法在mnist数据集增强中有见到
                       # sometimes(
                       #     iaa.ElasticTransformation(alpha=(0.5, 3.5), sigma=0.25)
                       # ),
 
                       # 扭曲图像的局部区域
                       # sometimes(iaa.PiecewiseAffine(scale=(0.01, 0.05)))
                   ],
 
                   random_order=True  # 随机的顺序把这些操作用在图像上
                   )
    ],
    random_order=True  # 随机的顺序把这些操作用在图像上
)
 
seq2 = iaa.Sequential([
    # iaa.GaussianBlur((0, 1.0))
    # iaa.GaussianBlur(2.5)
    # iaa.Crop(px=(0, 16)),
    # iaa.ContrastNormalization((0.5, 2.0), per_channel=0.5)
    # iaa.ContrastNormalization(0.5, per_channel=0.5)
    # iaa.ElasticTransformation(alpha=(0.5, 3.5), sigma=0.25)
    iaa.Add((-10, 10), per_channel=0.5),
])
image = cv2.imread(r'..\00240_basepic.jpg')
images = []
for i in range(1):
    images.append(image)
images_aug = seq2.augment_images(images)  # 实现图像增强
for im in images_aug:
    cv2.imshow('aug', im)
    cv2.waitKey(0)