VOC转Yolo脚本

最新推荐文章于 2024-07-05 17:43:17 发布

Learning_is_endless

最新推荐文章于 2024-07-05 17:43:17 发布

阅读量214

点赞数

CC 4.0 BY-SA版权

分类专栏：深度学习

本文链接：https://blog.csdn.net/zhangxiaoduoduo/article/details/116501828

深度学习专栏收录该内容

6 篇文章

订阅专栏

该脚本将VOC格式的XML注释文件转换为适用于Yolo目标检测模型的TXT格式。它遍历指定目录下的所有XML文件，提取类别信息并进行坐标转换。转换过程中，脚本会忽略不在预定义类别列表中的对象或标记为困难的对象。最终生成的TXT文件可直接用于Yolo模型的训练。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

VOC转Yolo脚本

import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join

classes = ["ball",
"cylinder",
"square cage",
"cube",
"circle cage",
"human body",
"metal bucket",
"tyre"
          ]

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(rootpath,xmlname):
    xmlpath = rootpath + '/labels'
    xmlfile = os.path.join(xmlpath,xmlname)
    with open(xmlfile, "r") as in_file:
      txtname = xmlname[:-4]+'.txt'
      txtpath = rootpath + '/labelYOLOs'
      if not os.path.exists(txtpath):
        os.makedirs(txtpath)
      txtfile = os.path.join(txtpath,txtname)
      with open(txtfile, "w+") as out_file:
        tree=ET.parse(in_file)
        root = tree.getroot()
        size = root.find('size')
        w = int(size.find('width').text)
        h = int(size.find('height').text)
        out_file.truncate()
        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')


if __name__ == "__main__":
    rootpath= r'D:\2021TrainAndTestSet\test_A\test_A'
    xmlpath=rootpath+'/labels'
    list=os.listdir(xmlpath)
    for i in range(0,len(list)) :
        path = os.path.join(xmlpath,list[i])
        if ('.xml' in path)or('.XML' in path):
            convert_annotation(rootpath,list[i])
            print('done', i)
        else:
            print('not xml file',i)