将VOC格式标注文件转换为Yolo格式

这篇文章主要参考博客中的代码，对原博客VOC格式数据集转yolo格式代码进行一定修改、添加注释，此外还在后面添加了我自己写的一段关于对转换后的图片和标注文件进行整理的脚本代码。

关于数据集在Yolo格式和VOC格式的区别,我已经在这篇博客将Yolo格式标注文件转换为VOC格式中详细讲过，这里不再赘述，下面之间列出所用代码。

1.将VOC格式标注文件（xml文件）转换成Yolo格式标注文件（txt文件）

import xml.etree.ElementTree as ET
import os


voc_folder = r"/home/dwt/DataSets/HRSID/HRSID_JPG_voc格式/Annotations" #储存voc格式的标注文件的文件夹
yolo_folder = r"/home/dwt/DataSets/HRSID/HRSID_yolo/Annotations_yolo" #转换后的yolo格式标注文件的储存文件夹

class_id = ["ship"] #储存数据集中目标种类名称的列表，接下来的转换函数中会将该列表中种类名称对应的列表索引号作为写入yolo标注文件中该类目标的种类序号

#voc标注的目标框坐标值转换到yolo标注的目标框坐标值的函数
def convert(size, box):
    dw = 1. / size[0]
    dh = 1. / size[1]
    x = (box[0] + box[1]) / 2.0
    y = (box[2] + box[3]) / 2.0
    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)

#对单个voc标注文件进行转换成其对应的yolo文件的函数
def convert_annotation(xml_file):
    file_name = xml_file.strip(".xml")  # 这一步将所有voc格式标注文件取出后缀名“.xml”，方便接下来作为yolo格式标注文件的名称
    in_file = open(os.path.join(voc_folder,xml_file)) #打开当前转换的voc标注文件
    out_file = open(os.path.join(yolo_folder,file_name + ".txt",),'w') #创建并打开要转换成的yolo格式标注文件
    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'):
        cls = obj.find('name').text
        cls_id = class_id.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')



xml_fileList = os.listdir(voc_folder) #将所有voc格式的标注文件的名称取出存放到列表xml_fileList中
for xml_file in xml_fileList: #这里的for循环开始依次对所有voc格式标注文件生成其对应的yolo格式的标注文件
    convert_annotation(xml_file)

2.根据VOC格式中的train.txt和test.txt文件来将数据集图片和标注归类到Yolo格式下的train和val文件夹中。

因为Yolo格式的数据集是将训练数据和验证(测试)数据分成两个文件夹，而VOC格式是将所有图片放在一个图片文件夹、所有标注文件放在一个标注文件夹中。所有在上面的转换步骤完成后，我们还需要将训练数据和验证(测试)数据分别存放。下面列出整理代码：

import os
import shutil

train_txt = r"/home/dwt/DataSets/HRSID/HRSID_yolo/train.txt"
test_txt = r"/home/dwt/DataSets/HRSID/HRSID_yolo/test.txt"

img_folder = r"/home/dwt/DataSets/HRSID/HRSID_JPG/JPEGImages"
yolo_folder = r"/home/dwt/DataSets/HRSID/HRSID_yolo/Annotations_yolo"

HRSID_dir = r"/home/dwt/MyCode/object detection/yolov5/DataSets/HRSID"

file_train =  open(train_txt)
file_test = open(test_txt)
for line in file_train.readlines():
    line = line.strip()
    shutil.copyfile(os.path.join(img_folder,line + ".jpg"),os.path.join(HRSID_dir,"images","train",line + ".jpg")) #根据train.txt指示的文件名将对应的图片复制到yolo格式整理的数据文件夹中
    shutil.copyfile(os.path.join(yolo_folder,line + ".txt"),os.path.join(HRSID_dir,"labels","train",line + ".txt")) #根据train.txt指示的文件名将对应的标注文件复制到yolo格式整理的数据文件夹中

for line in file_test.readlines():
    line = line.strip()
    shutil.copyfile(os.path.join(img_folder, line + ".jpg"),os.path.join(HRSID_dir, "images", "val",line + ".jpg"))  # 根据train.txt指示的文件名将对应的图片复制到yolo格式整理的数据文件夹中
    shutil.copyfile(os.path.join(yolo_folder, line + ".txt"),os.path.join(HRSID_dir, "labels", "val",line + ".txt"))  # 根据train.txt指示的文件名将对应的标注文件复制到yolo格式整理的数据文件夹中

将VOC格式标注文件转换为Yolo格式

猜你喜欢