VOC数据集转YOLO数据集

VOC数据集样例

<annotation>
        <folder>VOC2007</folder>
        <filename>009947.jpg</filename>
        <source>
                <database>The VOC2007 Database</database>
                <annotation>PASCAL VOC2007</annotation>
                <image>flickr</image>
                <flickrid>317223454</flickrid>
        </source>
        <owner>
                <flickrid>vaio_2002</flickrid>
                <name>Jonathan</name>
        </owner>
        <size>
                <width>332</width>
                <height>500</height>
                <depth>3</depth>
        </size>
        <segmented>1</segmented>
        <object>
                <name>boat</name>
                <pose>Left</pose>
                <truncated>0</truncated>
                <difficult>0</difficult>
                <bndbox>
                        <xmin>115</xmin>
                        <ymin>92</ymin>
                        <xmax>250</xmax>
                        <ymax>321</ymax>
                </bndbox>
        </object>
        <object>
                <name>person</name>
                <pose>Unspecified</pose>
                <truncated>1</truncated>
                <difficult>1</difficult>
                <bndbox>
                        <xmin>165</xmin>
                        <ymin>288</ymin>
                        <xmax>186</xmax>
                        <ymax>312</ymax>
                </bndbox>
        </object>
</annotation>

对应YOLO数据集样例

格式说明：

• 每行对应一个box，每个box分别有五个参数

  1. label index
  2. center_x
  3. center_y
  4. width
  5. height

10 0.5496987951807228 0.413 0.4066265060240964 0.458
5 0.5286144578313253 0.6 0.06325301204819277 0.048

代码实现

import xmltodict
import os
from progressbar import *

xml_dir='./labels_voc' #原xml路径
txt_dir='./labels'     #转换后txt文件存放路径 

# 所有待检测的labels
class_names = ['aeroplane', 'cat', 'car', 'dog', 'chair', 'person', 'horse', 'bird',
          'tvmonitor', 'bus', 'boat', 'diningtable', 'bicycle', 'bottle', 'sofa',
          'pottedplant', 'motorbike', 'cow', 'train', 'sheep']


# 将voc xml中的object转化为对应的一条yolo数据
def get_yolo_data(obj,img_width,img_height):
    # 获取voc格式的数据信息 
    name = obj['name']
    xmin = float(obj['bndbox']['xmin'])
    xmax = float(obj['bndbox']['xmax'])
    ymin = float(obj['bndbox']['ymin'])
    ymax = float(obj['bndbox']['ymax'])

    # 计算yolo格式的数据信息
    class_idx = class_names.index(name)
    x_center = (xmin + xmax) / 2 / img_width
    y_center = (ymin + ymax) / 2 / img_height
    box_width = (xmax - xmin) / img_width
    box_height = (ymax - ymin) / img_height

    yolo_data = "{} {} {} {} {}\n".format(class_idx,x_center,y_center,box_width,box_height)
    
    return yolo_data
    

# 逐一处理xml文件，转换为YOLO所需的格式

#进度条支持
count = 0 #计数器
widgets = ['VOC2YOLO: ',Percentage(), ' ', Bar('#'),' ', Timer(),' ', ETA()]
pbar = ProgressBar(widgets=widgets, maxval=len(os.listdir(xml_dir))).start()

for xml_file in os.listdir(xml_dir):
    xml_file_path = os.path.join(xml_dir,xml_file)
    txt_file_path = os.path.join(txt_dir,xml_file[:-4]+".txt")
    yolo_data = ""
    
    with open(xml_file_path) as f:
        xml_str = f.read()
    
     # 转为字典
    xml_dic = xmltodict.parse(xml_str)
    
    # 获取图片的width、height
    img_width = float(xml_dic["annotation"]["size"]["width"])
    img_height = float(xml_dic["annotation"]["size"]["height"])
    
    # 获取xml文件中的object
    objects = xml_dic["annotation"]["object"]
    
    if isinstance(objects,list): # xml文件中包含多个object
        for obj in objects:
            yolo_data += get_yolo_data(obj,img_width,img_height)   
    else: # xml文件中包含1个object
        obj = objects
        yolo_data += get_yolo_data(obj,img_width,img_height)
    
    with open(txt_file_path,'w') as f:
        f.write(yolo_data)
    
    #更新进度条
    count += 1
    pbar.update(count)

pbar.finish()

实现结果

原始数据

labels_voc
├── 000005.xml
├── 000007.xml
├── 000009.xml
├── 000012.xml
├── 000016.xml
├── 000017.xml
├── 000019.xml
├── 000020.xml
├── 000021.xml

生成数据

labels
├── 000005.txt
├── 000007.txt
├── 000009.txt
├── 000012.txt
├── 000016.txt
├── 000017.txt
├── 000019.txt
├── 000020.txt
├── 000021.txt

参考资料

1. PASCAL VOC 数据集转化为yolo数据集格式：https://blog.csdn.net/guo_python/article/details/107984940