数据增强，自动调整标签文件

2024-07-10 18:51| 来源: 网络整理| 查看: 265

现在框架都带有一些数据增强类，但不是太强大，这里记录用albumentations库实现多种增强，这个增强库的使用介绍已经很多了，所以我在这共享下我增强后自动调整标签的步骤。对带有box的数据，色彩增强不会改变box位置，但是旋转、翻转、裁切类操作，会对box数据产生影响，所以需要对标签文件自动修正。以pascal voc格式为例，在增强后自动调整xml文件，当然，coco格式也是支持的，需要自己去修改了。直接上代码，底下再说明：

import cv2 import xml.etree.ElementTree as ET import os,sys from albumentations import HorizontalFlip, IAAPerspective, ShiftScaleRotate, CLAHE, \ RandomRotate90, Transpose, ShiftScaleRotate, Blur, CenterCrop, RandomCrop, \ OpticalDistortion, GridDistortion, HueSaturationValue, \ IAAAdditiveGaussianNoise, GaussNoise, MotionBlur, MedianBlur, \ IAAPiecewiseAffine, IAASharpen, IAAEmboss, RandomContrast, \ RandomBrightness, Flip, OneOf, VerticalFlip, Resize, Rotate, Compose import numpy as np def pretty_xml(element, indent = '\t', newline = '\n', level=0): if element: # 判断element是否有子元素 if (element.text is None) or element.text.isspace(): # 如果element的text没有内容 element.text = newline + indent * (level + 1) else: element.text = newline + indent * (level + 1) + element.text.strip() + newline + indent * (level + 1) # else: # 此处两行如果把注释去掉，Element的text也会另起一行 # element.text = newline + indent * (level + 1) + element.text.strip() + newline + indent * level temp = list(element) # 将element转成list for subelement in temp: if temp.index(subelement) 0: 'fish'} aug_ver = get_aug([VerticalFlip(p = 1)]) #垂直方向翻转 aug_hor = get_aug([HorizontalFlip(p=1)]) #水平方向翻转 aug_res = get_aug([Resize(p=1, height=256, width=256)]) #resize aug_cen = get_aug([CenterCrop(p=1, height=200, width=200)], min_area=4000) aug_cen = get_aug([CenterCrop(p=1, height=100, width=100)], min_visibility=0.3) # 只返回变换后可见性大于 threshold 的 boxes aug_ran = get_aug([RandomCrop(p=1, height=100, width=100)]) aug_SCR =get_aug([ShiftScaleRotate(shift_limit=0.0625, scale_limit=1, rotate_limit=45, p=1)]) #旋转、裁切 aug_rot = get_aug([Rotate(limit=60, p =1.0)]) aug_list = [aug_ver, aug_hor, aug_rot] #想用哪个，就添加在找个list里 #-------------------------- 读取xml，解析，增强图像，修改box信息，写入xml -----------------------------# if __name__ == '__main__': jpgPath = 'JPEGImages' xmlPath = 'Annotations/' xmls = os.listdir(xmlPath) for xml in xmls: xmlName = xml.split('.')[0] imgName = xmlName + '.jpg' try: tree = ET.parse(os.path.join(xmlPath, xml)) root = tree.getroot() except Exception as e: print('prase ' + xml + ' failed!') sys.exit() else: image = cv2.imread(os.path.join(jpgPath, imgName)) for width in root.iter('width'): if int(width.text) == 0: width.text = str(image.shape[1]) for height in root.iter('height'): if int(height.text) == 0: height.text = str(image.shape[0]) tree.write(os.path.join(xmlPath, xmlName + '.xml')) bboxes = [] for object in root.findall('object'): for box in object.findall('bndbox'): x_min = int(box.find('xmin').text) x_max = int(box.find('xmax').text) y_min = int(box.find('ymin').text) y_max = int(box.find('ymax').text) root.remove(object) bboxes.append([x_min, y_min, x_max, y_max]) category_id = np.zeros(len(bboxes)) annotations = {'image': image, 'bboxes': bboxes, 'category_id': category_id} for i, aug in enumerate(aug_list): aug_type = str(aug).split('(')[1][4:] augmented = aug(**annotations) for iter in range(len(augmented['bboxes'])): x_min, y_min, x_max, y_max = augmented['bboxes'][iter] x_min, x_max, y_min, y_max = int(x_min), int(x_max), int(y_min), int(y_max) insert_object(root, x_min, x_max, y_min, y_max) for filename in root.iter('filename'): name = filename.text.split('.')[0] filename.text = name + aug_type + '.jpg' for path in root.iter('path'): pathname = path.text.split('.')[0] path.text = pathname + aug_type + '.jpg' for width in root.iter('width'): width.text = str(image.shape[1]) for height in root.iter('height'): height.text = str(image.shape[0]) if len(augmented['bboxes']) > 0: cv2.imwrite(os.path.join(jpgPath, xmlName + aug_type +'.jpg'), augmented['image']) pretty_xml(root) tree.write(os.path.join(xmlPath, xmlName + aug_type +'.xml')) for object in root.findall('object'): root.remove(object) # #centerCrop # aug = get_aug([CenterCrop(p=1, height=100, width=100)]) # augmented = aug(**annotations) # visualize(augmented, category_id_to_name) # #certerCrop，并限定最小box面积 # aug = get_aug([CenterCrop(p=1, height=200, width=200)], min_area=4000) # augmented = aug(**annotations) # visualize(augmented, category_id_to_name) # # 只返回变换后可见性大于 threshold 的 boxes # aug = get_aug([CenterCrop(p=1, height=100, width=100)], min_visibility=0.3) # augmented = aug(**annotations) # visualize(augmented, category_id_to_name) # aug = get_aug([RandomCrop(p=1, height=100, width=100)]) # augmented = aug(**annotations) # visualize(augmented, category_id_to_name) # #旋转、裁切 # aug =get_aug([ShiftScaleRotate(shift_limit=0.0625, # scale_limit=1, # rotate_limit=45, p=1)]) # augmented = aug(**annotations) # visualize(augmented, category_id_to_name) # #旋转 # aug = get_aug([Rotate(limit=60, p = 0.7)]) # augmented = aug(**annotations) # visualize(augmented, category_id_to_name) #多种增强混合，同时使用 # def augment_flips_color(p=.5): # return Compose([ # # CLAHE(), # Transpose(), # ShiftScaleRotate(shift_limit=0.0625, # scale_limit=1, # rotate_limit=45, p=.75), # # Blur(blur_limit=3), # # OpticalDistortion(), # # GridDistortion(), # # HueSaturationValue() # ], p=p) # # aug = augment_flips_color(p=1) # augmented = aug(**annotations) # visualize(augmented, category_id_to_name)

简单介绍下过程和用法：首先导入albumentations库，创建一个pretty_xml函数，这个主要用来修改xml格式的，因为xml.etree.ElementTree写入xml文件没有缩进，比较丑，当然，这个函数是抄来的，原作者博客找不到地址了。接着是xml中插入object函数、两个可视化函数。我自己这个项目里，类别只有鱼，所以category_id_to_name只有一个。可使用的增强方式很多，随便创建了几个见代码，放入aug_list列表中，供后面调用。原有xml和jpg按照voc格式要求放入两个文件夹，接下来读取xml，解析，增强图像，修改box信息，写入xml。注意三点：1、增强后box数量可能会减少，相应调整xml；2、旋转增强会导致box位置不准确，后面需要用labelimg软件手动调整一下，提高精度；3、我的lebelimg产生的xml中，存在大量width、height数值为0，导致我后面训练出错，所以我在这一并修改了原xml文件。运行完成后，会在JPEGImages文件夹下生成大量增强图像，以增强方式命名，在Annotations文件夹下生成jpg同名xml文件。后面注释部分代码，是单个使用增强的示例，包括一个混合增强的示例。

【本文地址】

公司简介

联系我们