python读取数据集并生成txt文件

前言

在目标检测领域，无论是公开数据集还是自建的数据集，数据集的存储方式通常是一个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可以在类别文件中检索其对应的类别名称。

总结

大部分公开数据集都可以用此方法进行处理。我也遇到过标注数据中只有四个坐标的情况，如果你处理过程中遇到了问题，都可以在下方进行留言和评论。