【Code】格式转换 voc2yolo()
直接上代码
import xml.etree.ElementTree as ET
import os
- 分类的集合，类别和标号的桥梁
- 文件路径
classes = ["class_0", "class_1", "class_2", "class_3", "class_4"]
file_paths = ['path_0', 'path_1', 'path_2']
- 转换成相对坐标(中心x坐标，中心y坐标，宽，高)
def convert(size, box):
dw = 1./size[0]
dh = 1./size[1]
x = (box[0] + box[1])/2.0
y = (box[2] + box[3])/2.0
w = box[1] - box[0]
h = box[3] - box[2]
x = x * dw
w = w * dw
y = y * dh
h = h * dh
return (x,y,w,h)
- 读取xml文件，新建并写入txt文件
def convert_annotation(file_path, image_id):
in_file = open('%s/%s.xml'%(file_path, image_id))
out_file = open('%s/%s.txt'%(file_path, image_id), 'w')
tree=ET.parse(in_file)
root = tree.getroot()
size = root.find('size')
w = int(size.find('width').text)
h = int(size.find('height').text)
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:
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')
- 提取出所有的图片编号，为了后期作为文件名以及循环做准备。
def img_id_set(path):
imgids = []
file_type = '.xml'
for root, dir, files in os.walk(path):
for file in files:
# 找出以'.xml'为后缀的文件，将其后缀删除，并存入返回值imgids中
if file.endswith(file_type):
f_name = file.replace(file_type, '')
imgids.append(f_name)
return imgids
- 格式转换函数
def voc2yolo(path_set):
for path in path_set:
for img_id in img_id_set(path):
convert_annotation(file_path=path, image_id=img_id)
- 程序运行
voc2yolo(path_set=file_paths)