前言
在目标检测领域,无论是公开数据集还是自建的数据集,数据集的存储方式通常是一个JPEGImages图片文件夹和一个Annotations的xml标注数据文件夹。
在进行数据处理时,首先要做的工作是将其读取并写入一个txt检索文本中,方便后续的操作。
一、文件路径
import os
import xml.etree.ElementTree as ET
#获取所有类别
classes = get_classes('../model_data/dataset/voc_dataset/voc_classes.txt')
#文件保存路径
list_file = open('../model_data/dataset/voc_dataset/VOC2012.txt', 'w')
#图片和xml路径
image_path = 'G:/数据集/VOC2012/JPEGImages/'
xml_path = 'G:/数据集/VOC2012/Annotations/'
二、获取所有类别
classes获取的是类别文件,存储了所有类别的名称。
这里定义了一个函数用来获取该文件中所有的类别,并用list保存。
#获取所有类别
def get_classes(classes_path):
classes = os.path.expanduser(classes_path)
with open(classes) as f:
class_names = f.readlines()
class_names = [c.strip() for c in class_names]
return class_names
三、读取标注文件xml
打开xml文件,可以看到里面的标注信息。每一个object表示一个box框,name是框中物体所属的类别,还有一个属性difficult,表示识别这个物体的难度(0表示没有难度,1表示有难度)。在读取文件的时候可以选择是否读取(代码中注释的部分)。
#读取xml文件得到所有标注框
def convert_annotation(image_id, list_file):
in_file = open(os.path.join(xml_path,'%s.xml'%(image_id)),"rb")
tree=ET.parse(in_file)
root = tree.getroot()
for obj in root.iter('object'):
cls = obj.find('name').text
cls_id = classes.index(cls) # 获取类别id
#是否读取难识别标签
# difficult = 0
# if obj.find('difficult')!=None:
# difficult = obj.find('difficult').text
# if cls not in classes or int(difficult) == 1:
# continue
xmlbox = obj.find('bndbox')
b = (int(float(xmlbox.find('xmin').text)), int(float(xmlbox.find('ymin').text)), int(float(xmlbox.find('xmax').text)), int(float(xmlbox.find('ymax').text)))
list_file.write(" " + ",".join([str(a) for a in b]) + ',' + str(cls_id))
list_file.write('\n')
因为标注文件名和图片文件名是一一对应的,所以可以通过图片名称来读取xml标注文件,然后依次读取每个object中的类别信息、难度信息、框的坐标信息,最后将其写入txt文件中。
四、写入文件
使用for循环依次读取image_path中的图片文件,然后一行行写入txt文本中。
for file in os.listdir(image_path):
image_id = os.path.splitext(file)[0] #获取文件名
list_file.write(os.path.join(image_path, file)) #写图片完整路径
convert_annotation(image_id, list_file) #写标注框
list_file.close()
每一行代表一张图片,开头是图片路径,后面五个数据分别是框的坐标和类别id(从0开始),通过id可以在类别文件中检索其对应的类别名称。
总结
大部分公开数据集都可以用此方法进行处理。我也遇到过标注数据中只有四个坐标的情况,如果你处理过程中遇到了问题,都可以在下方进行留言和评论。