简介
大部分都是参考YOLO从零开始:基于YOLOv3的行人检测入门指南,不过他这个还得分几步,各种配置也比较麻烦,因此进行了优化,直接一步到位。
具体操作
原来的VOC2007和VOC2012的路径是这样的:
- 建立一个data文件夹,这样直接按照相对路径去写图片和标注文件位置,方便自己使用
- 运行下面代码
import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir
from os.path import join
import shutil
# 数据集之间的标注转换
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
return (x,y,w,h)
def convert_annotation(year, image_id, root, labels_dir):
in_file = open(os.path.join(root, 'VOC%s/Annotations/%s.xml'%(year, image_id)))
out_file = open(os.path.join(labels_dir, image_id + '.txt'), '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:
# 保留困难的数据集
if cls not in classes:
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__':
sets = [('2007', 'trainval'), ('2007', 'test'), ('2012', 'trainval')]
classes = ['person']
src_dir = '/home/common/wangsong/VOC/VOCdevkit' # 存放VOC2007 VOC2012的路径
dst_dir = '/home/wangsong/data' # 转换的YOLO数据集的路径
images_dir = os.path.join(dst_dir, 'images') # data/images
labels_dir = os.path.join(dst_dir, 'labels') # data/labels
# 创建YOLO这俩文件夹images_dir labels_dir
if not os.path.exists(images_dir):
os.makedirs(images_dir)
if not os.path.exists(labels_dir):
os.makedirs(labels_dir)
for year, image_set in sets:
# 读取VOC的Main里面的txt,得到trainval.txt/test.txt里面的图片ID
image_ids = open(os.path.join(src_dir, 'VOC%s/ImageSets/Main/%s.txt' % (year, image_set))).read().strip().split()
# YOLO需要的训练集验证集图片所在路径
# 这里先分别写,后面再根据需要汇总即可
list_file = open(os.path.join(dst_dir, '%s_%s.txt' % (year, image_set)), 'w')
for image_id in image_ids:
# 判断有木有行人目标的标志
objCount = 0
# 打开VOC的XML标注文件
in_file = open(os.path.join(src_dir, 'VOC%s/Annotations/%s.xml' % (year, image_id)))
# YOLO所需要的标注文件
out_file = open(os.path.join(labels_dir, image_id + '.txt'), 'w')
# 解析VOC标注文件XML
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:
# 保留困难的数据集
if cls not in classes:
continue
# 说明有行人目标了
objCount += 1
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 objCount > 0:
# 将这图片路径加入YOLO的训练验证txt
# list_file.write('data/images/%s.jpg\n' % image_id) # 相对路径
list_file.write( os.path.join(dst_dir, 'images/%s.jpg\n' % image_id) ) # 绝对路径
src_img = os.path.join(src_dir, 'VOC%s/JPEGImages/%s.jpg'%(year, image_id))
dst_img = os.path.join(images_dir, '%s.jpg' % image_id)
# 将这图片复制到YOLO文件夹的images_dir
shutil.copy(src_img, dst_img)
else:
# 如果这图片没有行人,那么就直接删除了这个标注文件
# 不需要复制图片和标注文件了,也不要写进训练验证txt
del out_file
os.remove(os.path.join(labels_dir, image_id + '.txt'))
list_file.close()
# os.system("cat 2007_train.txt 2007_val.txt 2012_train.txt 2012_val.txt > train.txt")
# os.system("cat 2007_train.txt 2007_val.txt 2007_test.txt 2012_train.txt 2012_val.txt > train.all.txt")