1. 功能
将voc数据集转换为YOLOv5可用的数据集。
2. 代码
import xml.etree.ElementTree as ET
import os
'''
VOC数据集转U版的YOLO数据集:voc数据集标签文件为和图片名对应的xml文件,u版yolo数据集标签文件为和图片名对应的txt文件
假设:voc数据集中xml文件和图片文件一一对应
'''
def dirCheck(checkdir): # 如果没有改文件夹,则生成该文件夹
if not os.path.exists(checkdir):
os.makedirs(checkdir)
def convert(size, box):
dw = 1./(size[0])
dh = 1./(size[1])
x = (box[0] + box[1])/2.0 - 1
y = (box[2] + box[3])/2.0 - 1
w = box[1] - box[0]
h = box[3] - box[2]
x = x*dw
w = w*dw
y = y*dh
h = h*dh
x = '%.3f' % x # 保留小数点后3位
w = '%.3f' % w
y = '%.3f' % y
h = '%.3f' % h
return (x,y,w,h)
def convert_annotation( rootpath, xmlname, classes):
xmlpath = os.path.join(rootpath,'ann') # xml所在文件夹
xmlfile = os.path.join(xmlpath,xmlname) # xml文件全路径
with open(xmlfile, "r") as in_file:
txtname = xmlname[:-4]+'.txt' # txt文件名,xml文件文件名转成txt文件名:00001.xml -> 00001.txt
txtpath = os.path.join(rootpath, 'YOLO', 'labels') # txt文件所在文件夹 D:/.../YOLO/labels
dirCheck(txtpath)
txtfile = os.path.join(txtpath,txtname) # txt文件全路径:D:/.../YOLO/labels/00001.txt
with open(txtfile, "w+") as out_file:
tree=ET.parse(in_file)
root = tree.getroot()
size = root.find('size')
w = int(size.find('width').text)
h = int(size.find('height').text)
out_file.truncate()
for obj in root.iter('object'):
difficult = obj.find('difficult').text
cls = obj.find('name').text
if cls not in classes or int(difficult)==1:
print('存在范围之外的类别,请检查xml文件',xmlfile)
print('错误类别:',cls)
continue
cls_id = classes.index(cls)
xmlbox = obj.find('bndbox')
b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text), float(xmlbox.find('ymax').text))
bb = convert((w,h), b)
out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
if __name__ == "__main__":
classes = ['class1', 'class2', 'class3', 'class4', 'class5']
voc_root_dir='D:/0416'
ann_dir = os.path.join(voc_root_dir,'ann') # 最后不要加斜杠,否则:D:/test2/VOCdevkit\Annotations/
ann_list=os.listdir(ann_dir) # Annotations文件夹下的xml文件文件名列表
for i in range(0,len(ann_list)): # 对于每一个xml文件
if ann_list[i].endswith('.xml'): # 判断是不是xml文件
# ann_path = os.path.join(ann_dir, ann_list[i]) # xml文件全路径
convert_annotation(voc_root_dir, ann_list[i], classes)
else:
print('请注意:',ann_list[i])