深度学习,数据增强批量处理代码。

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), encoding="utf-8")
    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'), encoding="utf-8")  # 这里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(image_id) + "aug" + '.xml'))


def change_xml_list_annotation(root, image_id, new_target, saveroot, idd, new_w, new_h):
    in_file = open(os.path.join(root, str(image_id) + '.xml'), encoding="utf-8")  # 这里root分别由两个意思
    tree = ET.parse(in_file)
    elem = tree.find('filename')
    elem.text = (idd + '.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.find('size').find('width').text = str(new_w)
    tree.find('size').find('height').text = str(new_h)

    tree.write(os.path.join(saveroot, idd + '.xml'), encoding="utf-8")
    in_file.close()


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 = "C:/Users/haohao/Desktop/11/all/images"
    XML_DIR = "C:/Users/haohao/Desktop/11/all/annotations"

    AUG_XML_DIR = "C:/Users/haohao/Desktop/11/output1/annotations"  # 存储增强后的XML文件夹路径
    AUG_IMG_DIR = "C:/Users/haohao/Desktop/11/output1/images"  # 存储增强后的影像文件夹路径

    try:
        shutil.rmtree(AUG_XML_DIR)
    except FileNotFoundError as e:
        a = 1
    os.makedirs(AUG_XML_DIR)

    try:
        shutil.rmtree(AUG_IMG_DIR)
    except FileNotFoundError as e:
        a = 1
    os.makedirs(AUG_IMG_DIR)

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

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

    # 影像增强
    # 产生一个处理图片的Sequential函数原型:iaa.Sequential
    seq = iaa.SomeOf(
        (1, 2),  # 1-2种变形
        [  # 选几个执行
            # iaa.CropAndPad(px=(0, 40)), # crop images from each side by 0 to 20px (randomly chosen)
            # iaa.Cutout(nb_iterations=(1, 10), size=0.02, squared=False,fill_mode="constant", cval=0),

            # 提取该通道,乘0.5和1.5之间的一个因子,然后转换回原始色彩空间。
            iaa.MultiplyBrightness((1.0, 1.8)),
            # 创建一个通过随机因素使图像亮度变差的增强器:
            iaa.pillike.EnhanceBrightness(),
            # 增亮
            iaa.WithBrightnessChannels(
                iaa.Add((-50, 50)), to_colorspace=[iaa.CSPACE_Lab, iaa.CSPACE_HSV]),
            iaa.LinearContrast((0.7, 1.3)),  ##增强或减弱图片的对比度
            # iaa.Cartoon(blur_ksize=3, segmentation_size=1.0,saturation=2.0, edge_prevalence=1.0),
            # 散焦模糊
            iaa.imgcorruptlike.DefocusBlur(severity=2),
            iaa.MotionBlur(k=6),  # 像素移动模糊
            # iaa.Rain(speed=(0.1, 0.3)),  # 加雨线
            # iaa.Affine(scale=(0.5, 1.5)),  # 将图像缩放到原始大小的 50% 到 150%:
            # iaa.Clouds(),  # 加云
            # iaa.Snowflakes(flake_size=(0.3, 0.5), speed=(0.007, 0.03)),  # 加雪点
            # iaa.Canny(alpha=(0.0, 0.3)),#边缘提取
            # iaa.Dropout((0.01, 0.1), per_channel=True), # randomly remove up to 10% of the pixels
            # iaa.Flipud(0.5),  # 垂直
            # iaa.Fliplr(0.5),  # 水平 (垂直+水平=镜像)
            # iaa.Crop(px=(1, 16), keep_size=False),
            # mask 增强
            iaa.CoarseSaltAndPepper(0.06, size_percent=(0.01, 0.1)),
            # iaa.Multiply((0.7, 1.3)),  # change brightness, doesn't affect BBs
            iaa.GaussianBlur(sigma=(0.2, 0.5)),
            # 参数sigma表示高斯模糊的强度,一般 sigma = 0 表示没有高斯模糊,sigma = 3.0 表示很强的模糊。采用 tripe 输入,表示 sigma 值在 (0, 0.03) 之间随机取。
            # iaa.Rot90(1, keep_size=False),
            # # 每个区域具有20% 相应的高度和宽度大小(对于非方形图像,这会导致要填充非方形区域)。
            # iaa.Cutout(nb_iterations=(5, 15), size=0.08, squared=False),
            # iaa.Sequential([iaa.Rot90(1, keep_size=False),#不保持原始图像尺寸
            #        iaa.CropToFixedSize(width=960, height=960)]),
            # # 将所有图像的所有组件设置50 % 为零:
            # iaa.TotalDropout(0.10),
            # iaa.Affine(
            #     translate_percent=0,
            #     scale=1.0,#不缩放
            #     rotate=90,

            #     #mode="edge"
            # ),  # 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):
                error = 0
                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)))

                    n_x1 = int(bbs_aug.bounding_boxes[0].x1)
                    n_y1 = int(bbs_aug.bounding_boxes[0].y1)
                    n_x2 = int(bbs_aug.bounding_boxes[0].x2)
                    n_y2 = int(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
                    # print(n_x1 >= n_x2,n_y1 >= n_y2)
                    if n_x1 >= n_x2 or n_y1 >= n_y2 or abs(n_x1 - n_x2) <= 20 or abs(n_y1 - n_y2) <= 20:
                        print('框不合规范舍去(error):', name)
                        error = 1
                        break
                        # new_bndbox_list.append([n_x2, n_y2, n_x1, n_y1])

                    new_bndbox_list.append([n_x1, n_y1, n_x2, n_y2])

                if error == 1:
                    continue

                # 存储变化后的图片
                image_aug = seq_det.augment_images([img])[0]
                path = os.path.join(AUG_IMG_DIR, str(name[:-4] + "_a%02d" % epoch) + '.jpg')
                print(path)
                image_auged = bbs.draw_on_image(image_aug)
                Image.fromarray(image_auged).save(path)
                print(image_auged.shape)  # (高,宽,通道)
                new_width = image_auged.shape[1]
                new_height = image_auged.shape[0]

                # 存储变化后的XML
                change_xml_list_annotation(XML_DIR, name[:-4], new_bndbox_list, AUG_XML_DIR,
                                           str(name[:-4] + "_a%02d" % epoch), new_width, new_height)
                # print(str(name[:-4] + "_a%s"%epoch ) + '.jpg')
                new_bndbox_list = []

            # break

  • 2
    点赞
  • 3
    收藏
    觉得还不错? 一键收藏
  • 1
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论 1
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值