参考：(41条消息) voc数据集对有标签的数据集数据增强_CodingWZP的博客-CSDN博客_voc数据集增强
在博主原有的基础上，考虑自身数据集实际情况，对其添加了天气状况，增强数据集鲁棒性
核心已注释
voc和yolo转换直接抄某博主的，找不到链接，见谅
数据增强'''
Author: CodingWZP
Email: codingwzp@gmail.com
Date: 2021-08-06 10:51:35
LastEditTime: 2021-08-09 10:53:43
Description: Image augmentation with label.
'''
import xml.etree.ElementTree as ET
import os
import imgaug as ia
import numpy as np
import shutil
from tqdm import tqdm
from PIL import Image
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')
if elem != None:
elem.text = (saveroot + str(id) + '.jpg')
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 = "E:/data/data_augemnt/images/val/"#图片路径
XML_DIR = "E:/data/data_augemnt/labels/voa_val/"#标签路径，voc格式
AUG_XML_DIR = "E:/data/label_ad/" # 存储增强后的XML文件夹路径
try:
shutil.rmtree(AUG_XML_DIR)
except FileNotFoundError as e:
a = 1
mkdir(AUG_XML_DIR)
AUG_IMG_DIR = "E:/data/image_ad" # 存储增强后的影像文件夹路径
try:
shutil.rmtree(AUG_IMG_DIR)
except FileNotFoundError as e:
a = 1
mkdir(AUG_IMG_DIR)
AUGLOOP = 1 # 每张影像增强的数量
boxes_img_aug_list = []
new_bndbox = []
new_bndbox_list = []
# 影像增强
seq = iaa.Sequential([
#OneOf选取一种天气
iaa.OneOf([iaa.Rain(drop_size=(0.2, 0.4), speed=(0.01, 0.02)),#下雨
iaa.imgcorruptlike.Snow(severity=1),#下雪
iaa.imgcorruptlike.Fog(severity=1)]),#雾化
#选取3-6种增强手段
iaa.SomeOf((3,6),[
iaa.Invert(0.5),
iaa.CoarsePepper(size_percent=0.05),#随机增加黑框遮挡
iaa.Multiply((0.5,1.5)),#曝光
iaa.Fliplr(0.5),#镜像
iaa.GaussianBlur(sigma=(1.0, 3.0)),#高斯模糊
iaa.Affine(
#rotate(-15,15) #旋转，不建议使用，voc数据集有概率会出现错误
translate_px={"x": 15, "y": 15},
scale=(0.5, 0.8))])#缩小
# change brightness, doesn't affect BBs
])
for name in tqdm(os.listdir(XML_DIR), desc='Processing'):
bndbox = read_xml_annotation(XML_DIR, name)
# 保存原xml文件
#shutil.copy(os.path.join(XML_DIR, name), AUG_XML_DIR)
# 保存原图
og_img = Image.open(IMG_DIR + '/' + name[:-4] + '.jpg')
#og_img.convert('RGB').save(AUG_IMG_DIR + name[:-4] + '.jpg', 'JPEG') #原图和影像增强图保存在一起
og_xml = open(os.path.join(XML_DIR, name))
tree = ET.parse(og_xml)
# 修改增强后的xml文件中的filename
elem = tree.find('filename')
elem.text = (name[:-4] + '.jpg')
# tree.write(os.path.join(AUG_XML_DIR, name)) #原xml与增强后xml保存同一文件
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, size=0)
Image.fromarray(image_auged).convert('RGB').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 = []
print('Finish!')
yolo转voc数据集
注意修改class类名，方便后续操作
import os
import shutil
import cv2
from lxml import etree
def VOC2Yolo(class_num, voc_img_path, voc_xml_path, yolo_txt_save_path, yolo_img_save_path=None):
xmls = os.listdir(voc_xml_path)
xmls = [x for x in xmls if x.endswith('.xml')]
if yolo_img_save_path is not None:
if not os.path.exists(yolo_img_save_path):
os.mkdir(yolo_img_save_path)
if not os.path.exists(yolo_txt_save_path):
os.mkdir(yolo_txt_save_path)
all_xmls = len(xmls)
for idx, one_xml in enumerate(xmls):
xl = etree.parse(os.path.join(voc_xml_path, one_xml))
root = xl.getroot()
objects = root.findall('object')
img_size = root.find('size')
img_w = 0
img_h = 0
if img_size:
img_width = img_size.find('width')
if img_width is not None:
img_w = int(img_width.text)
img_height = img_size.find('height')
if img_height is not None:
img_h = int(img_height.text)
label_lines = []
for ob in objects:
one_annotation = {}
label = ob.find('name').text
one_annotation['tag'] = label
one_annotation['flag'] = False
bbox = ob.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)
if img_w == 0 or img_h == 0:
img = cv2.imread(os.path.join(voc_img_path, one_xml.replace('.xml', '.jpg')))
img_h, img_w = img.shape[:2]
bbox_w = (xmax - xmin) / img_w
bbox_h = (ymax - ymin) / img_h
bbox_cx = (xmin + xmax) / 2 / img_w
bbox_cy = (ymin + ymax) / 2 / img_h
try:
bbox_label = class_num[label]
label_lines.append(f'{bbox_label} {bbox_cx} {bbox_cy} {bbox_w} {bbox_h}' + 'n')
except Exception as e:
print("not find number label in class_num ", e, one_xml)
label_lines = []
break
if len(label_lines):
with open(os.path.join(yolo_txt_save_path, one_xml.replace('.xml', '.txt')), 'w') as fp:
fp.writelines(label_lines)
if yolo_img_save_path is not None:
shutil.copy(os.path.join(voc_img_path, one_xml.replace('.xml', '.jpg')),
os.path.join(yolo_img_save_path))
print(f"processing: {idx}/{all_xmls}")
if __name__ == '__main__':
VOC2Yolo(
class_num={'battle': 0}, # 标签种类
voc_img_path='E:/data/image_ad', # 数据集图片文件夹存储路径
voc_xml_path='E:/data/label_ad', # 标签xml文件夹存储路径
yolo_txt_save_path='E:/data/data_augemnt/labels/train' # 将要生成的txt文件夹存储路径
)
voc转yolo
__Author__ = "Shliang"
__Email__ = "shliang0603@gmail.com"
import os
import xml.etree.ElementTree as ET
from xml.dom.minidom import Document
import cv2
'''
import xml
xml.dom.minidom.Document().writexml()
def writexml(self,
writer: Any,
indent: str = "",
addindent: str = "",
newl: str = "",
encoding: Any = None) -> None
'''
class YOLO2VOCConvert:
def __init__(self, txts_path, xmls_path, imgs_path):
self.txts_path = txts_path # 标注的yolo格式标签文件路径
self.xmls_path = xmls_path # 转化为voc格式标签之后保存路径
self.imgs_path = imgs_path # 读取读片的路径个图片名字，存储到xml标签文件中
self.classes = ["battle"]
# 从所有的txt文件中提取出所有的类别， yolo格式的标签格式类别为数字 0,1,...
# writer为True时，把提取的类别保存到'./Annotations/classes.txt'文件中
def search_all_classes(self, writer=False):
# 读取每一个txt标签文件，取出每个目标的标注信息
all_names = set()
txts = os.listdir(self.txts_path)
# 使用列表生成式过滤出只有后缀名为txt的标签文件
txts = [txt for txt in txts if txt.split('.')[-1] == 'txt']
print(len(txts), txts)
# 11 ['0002030.txt', '0002031.txt', ... '0002039.txt', '0002040.txt']
for txt in txts:
txt_file = os.path.join(self.txts_path, txt)
with open(txt_file, 'r') as f:
objects = f.readlines()
for object in objects:
object = object.strip().split(' ')
print(object) # ['2', '0.506667', '0.553333', '0.490667', '0.658667']
all_names.add(int(object[0]))
# print(objects) # ['2 0.506667 0.553333 0.490667 0.658667n', '0 0.496000 0.285333 0.133333 0.096000n', '8 0.501333 0.412000 0.074667 0.237333n']
print("所有的类别标签：", all_names, "共标注数据集：%d张" % len(txts))
# 把从xmls标签文件中提取的类别写入到'./Annotations/classes.txt'文件中
# if writer:
# with open('./Annotations/classes.txt', 'w') as f:
# for label in all_names:
# f.write(label + 'n')
return list(all_names)
def yolo2voc(self):
# 创建一个保存xml标签文件的文件夹
if not os.path.exists(self.xmls_path):
os.mkdir(self.xmls_path)
# # 读取每张图片，获取图片的尺寸信息（shape）
# imgs = os.listdir(self.imgs_path)
# for img_name in imgs:
# img = cv2.imread(os.path.join(self.imgs_path, img_name))
# height, width, depth = img.shape
# # print(height, width, depth) # h 就是多少行（对应图片的高度）， w就是多少列（对应图片的宽度）
#
# # 读取每一个txt标签文件，取出每个目标的标注信息
# all_names = set()
# txts = os.listdir(self.txts_path)
# # 使用列表生成式过滤出只有后缀名为txt的标签文件
# txts = [txt for txt in txts if txt.split('.')[-1] == 'txt']
# print(len(txts), txts)
# # 11 ['0002030.txt', '0002031.txt', ... '0002039.txt', '0002040.txt']
# for txt_name in txts:
# txt_file = os.path.join(self.txts_path, txt_name)
# with open(txt_file, 'r') as f:
# objects = f.readlines()
# for object in objects:
# object = object.strip().split(' ')
# print(object) # ['2', '0.506667', '0.553333', '0.490667', '0.658667']
# 把上面的两个循环改写成为一个循环：
imgs = os.listdir(self.imgs_path)
txts = os.listdir(self.txts_path)
txts = [txt for txt in txts if not txt.split('.')[0] == "classes"] # 过滤掉classes.txt文件
print(txts)
# 注意，这里保持图片的数量和标签txt文件数量相等，且要保证名字是一一对应的 (后面改进，通过判断txt文件名是否在imgs中即可)
if len(imgs) == len(txts): # 注意：./Annotation_txt 不要把classes.txt文件放进去
map_imgs_txts = [(img, txt) for img, txt in zip(imgs, txts)]
txts = [txt for txt in txts if txt.split('.')[-1] == 'txt']
print(len(txts), txts)
for img_name, txt_name in map_imgs_txts:
# 读取图片的尺度信息
print("读取图片：", img_name)
img = cv2.imread(os.path.join(self.imgs_path, img_name))
height_img, width_img, depth_img = img.shape
print(height_img, width_img, depth_img) # h 就是多少行（对应图片的高度）， w就是多少列（对应图片的宽度）
# 获取标注文件txt中的标注信息
all_objects = []
txt_file = os.path.join(self.txts_path, txt_name)
with open(txt_file, 'r') as f:
objects = f.readlines()
for object in objects:
object = object.strip().split(' ')
all_objects.append(object)
print(object) # ['2', '0.506667', '0.553333', '0.490667', '0.658667']
# 创建xml标签文件中的标签
xmlBuilder = Document()
# 创建annotation标签，也是根标签
annotation = xmlBuilder.createElement("annotation")
# 给标签annotation添加一个子标签
xmlBuilder.appendChild(annotation)
# 创建子标签folder
folder = xmlBuilder.createElement("folder")
# 给子标签folder中存入内容，folder标签中的内容是存放图片的文件夹，例如：JPEGImages
folderContent = xmlBuilder.createTextNode(self.imgs_path.split('/')[-1]) # 标签内存
folder.appendChild(folderContent) # 把内容存入标签
annotation.appendChild(folder) # 把存好内容的folder标签放到 annotation根标签下
# 创建子标签filename
filename = xmlBuilder.createElement("filename")
# 给子标签filename中存入内容，filename标签中的内容是图片的名字，例如：000250.jpg
filenameContent = xmlBuilder.createTextNode(txt_name.split('.')[0] + '.jpg') # 标签内容
filename.appendChild(filenameContent)
annotation.appendChild(filename)
# 把图片的shape存入xml标签中
size = xmlBuilder.createElement("size")
# 给size标签创建子标签width
width = xmlBuilder.createElement("width") # size子标签width
widthContent = xmlBuilder.createTextNode(str(width_img))
width.appendChild(widthContent)
size.appendChild(width) # 把width添加为size的子标签
# 给size标签创建子标签height
height = xmlBuilder.createElement("height") # size子标签height
heightContent = xmlBuilder.createTextNode(str(height_img)) # xml标签中存入的内容都是字符串
height.appendChild(heightContent)
size.appendChild(height) # 把width添加为size的子标签
# 给size标签创建子标签depth
depth = xmlBuilder.createElement("depth") # size子标签width
depthContent = xmlBuilder.createTextNode(str(depth_img))
depth.appendChild(depthContent)
size.appendChild(depth) # 把width添加为size的子标签
annotation.appendChild(size) # 把size添加为annotation的子标签
# 每一个object中存储的都是['2', '0.506667', '0.553333', '0.490667', '0.658667']一个标注目标
for object_info in all_objects:
# 开始创建标注目标的label信息的标签
object = xmlBuilder.createElement("object") # 创建object标签
# 创建label类别标签
# 创建name标签
imgName = xmlBuilder.createElement("name") # 创建name标签
imgNameContent = xmlBuilder.createTextNode(self.classes[int(object_info[0])])
imgName.appendChild(imgNameContent)
object.appendChild(imgName) # 把name添加为object的子标签
# 创建pose标签
pose = xmlBuilder.createElement("pose")
poseContent = xmlBuilder.createTextNode("Unspecified")
pose.appendChild(poseContent)
object.appendChild(pose) # 把pose添加为object的标签
# 创建truncated标签
truncated = xmlBuilder.createElement("truncated")
truncatedContent = xmlBuilder.createTextNode("0")
truncated.appendChild(truncatedContent)
object.appendChild(truncated)
# 创建difficult标签
difficult = xmlBuilder.createElement("difficult")
difficultContent = xmlBuilder.createTextNode("0")
difficult.appendChild(difficultContent)
object.appendChild(difficult)
# 先转换一下坐标
# (objx_center, objy_center, obj_width, obj_height)->(xmin，ymin, xmax,ymax)
x_center = float(object_info[1])*width_img + 1
y_center = float(object_info[2])*height_img + 1
xminVal = int(x_center - 0.5*float(object_info[3])*width_img) # object_info列表中的元素都是字符串类型
yminVal = int(y_center - 0.5*float(object_info[4])*height_img)
xmaxVal = int(x_center + 0.5*float(object_info[3])*width_img)
ymaxVal = int(y_center + 0.5*float(object_info[4])*height_img)
# 创建bndbox标签(三级标签)
bndbox = xmlBuilder.createElement("bndbox")
# 在bndbox标签下再创建四个子标签(xmin，ymin, xmax,ymax) 即标注物体的坐标和宽高信息
# 在voc格式中，标注信息：左上角坐标（xmin, ymin） （xmax, ymax）右下角坐标
# 1、创建xmin标签
xmin = xmlBuilder.createElement("xmin") # 创建xmin标签（四级标签）
xminContent = xmlBuilder.createTextNode(str(xminVal))
xmin.appendChild(xminContent)
bndbox.appendChild(xmin)
# 2、创建ymin标签
ymin = xmlBuilder.createElement("ymin") # 创建ymin标签（四级标签）
yminContent = xmlBuilder.createTextNode(str(yminVal))
ymin.appendChild(yminContent)
bndbox.appendChild(ymin)
# 3、创建xmax标签
xmax = xmlBuilder.createElement("xmax") # 创建xmax标签（四级标签）
xmaxContent = xmlBuilder.createTextNode(str(xmaxVal))
xmax.appendChild(xmaxContent)
bndbox.appendChild(xmax)
# 4、创建ymax标签
ymax = xmlBuilder.createElement("ymax") # 创建ymax标签（四级标签）
ymaxContent = xmlBuilder.createTextNode(str(ymaxVal))
ymax.appendChild(ymaxContent)
bndbox.appendChild(ymax)
object.appendChild(bndbox)
annotation.appendChild(object) # 把object添加为annotation的子标签
f = open(os.path.join(self.xmls_path, txt_name.split('.')[0]+'.xml'), 'w')
xmlBuilder.writexml(f, indent='t', newl='n', addindent='t', encoding='utf-8')
f.close()
if __name__ == '__main__':
txts_path1 = './data_augemnt/labels/val' #yolo存储文件
xmls_path1 = './data_augemnt/labels/voa_val' #voc存储文件
imgs_path1 = './data_augemnt/images/val' #图片存储文件
yolo2voc_obj1 = YOLO2VOCConvert(txts_path1, xmls_path1, imgs_path1)
labels = yolo2voc_obj1.search_all_classes()
print('labels: ', labels)
yolo2voc_obj1.yolo2voc()
