VOC转YOLO格式的代码备份

buxiangshui_cd

于 2024-04-07 11:22:02 发布

阅读量338

点赞数 4

分类专栏： YOLO学习文章标签： YOLO 人工智能深度学习

本文链接：https://blog.youkuaiyun.com/buxiangshui_cd/article/details/137457200

版权

YOLO学习专栏收录该内容

1 篇文章

订阅专栏

本文介绍了如何将VOC格式的XML标注文件转换为YOLO所需的txt格式，同时提供了数据集的划分代码，包括训练、验证和测试集的创建，以支持YOLO模型的训练过程。

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

学习了一段时间YOLO，整体感觉YOLO是一个比较好的框架，从框架源码的下载、训练、验证、测试。这里备份一段VOC转YOLO的相关代码，在这里也对原作者做出感谢。

VOC2YOLO.py

import os
import xml.etree.ElementTree as ET
import io


find_path = 'D:/数据集/annotations/'
savepath = 'D:/数据集/labels/'

# classes = ['car', 'Truck', 'person', 'bicycle', 'bus']
classes = ['hat', 'person','ventilator']


class Voc_Yolo(object):
    def __init__(self, find_path):
        self.find_path = find_path

    def Make_txt(self, outfile):
        out = open(outfile, 'w')
        print("创建成功：{}".format(outfile))
        return out

    def Work(self, count):
        # 找到文件路径
        for root, dirs, files in os.walk(self.find_path):
            # 找到文件目录中每一个xml文件
            for file in files:
                # 记录处理过的文件
                count += 1
                # 输入、输出文件定义
                input_file = find_path + file
                outfile = savepath + file[:-4] + '.txt'
                # 新建txt文件，确保文件正常保存
                out = self.Make_txt(outfile)
                # 分析xml树，取出w_image、h_image
                tree = ET.parse(input_file)
                root = tree.getroot()
                size = root.find('size')
                w_image = float(size.find('width').text)
                h_image = float(size.find('height').text)
                # 继续提取有效信息来计算txt中的四个数据
                for obj in root.iter('object'):
                    # 将类型提取出来，不同目标类型不同，本文仅有一个类别->0
                    classname = obj.find('name').text
                    # 如果类别不是对应在我们预定好的class文件中，或difficult==1则跳过
                    if classname not in classes == 1:
                        continue
                    # 通过类别名称找到id
                    cls_id = classes.index(classname)
                    xmlbox = obj.find('bndbox')
                    x_min = float(xmlbox.find('xmin').text)
                    x_max = float(xmlbox.find('xmax').text)
                    y_min = float(xmlbox.find('ymin').text)
                    y_max = float(xmlbox.find('ymax').text)
                    # 计算公式
                    x_center = ((x_min + x_max) / 2 - 1) / w_image
                    y_center = ((y_min + y_max) / 2 - 1) / h_image
                    w = (x_max - x_min) / w_image
                    h = (y_max - y_min) / h_image
                    # 文件写入
                    out.write(
                        str(cls_id) + " " + str(x_center) + " " + str(y_center) + " " + str(w) + " " + str(h) + '\n')
                out.close()
        return count


if __name__ == "__main__":
    data = Voc_Yolo(find_path)
    number = data.Work(0)
    print(number)

数据集划分：

split.py

mport os
import random
import shutil

# 原数据集目录
root_dir = 'D:/hatCHV/'
# 划分比例
train_ratio = 0.8
valid_ratio = 0.1
test_ratio = 0.1

# 设置随机种子
random.seed(42)

# 拆分后数据集目录
split_dir = 'D:/hat/Helmet_split'
os.makedirs(os.path.join(split_dir, 'train/images'), exist_ok=True)
os.makedirs(os.path.join(split_dir, 'train/labels'), exist_ok=True)
os.makedirs(os.path.join(split_dir, 'valid/images'), exist_ok=True)
os.makedirs(os.path.join(split_dir, 'valid/labels'), exist_ok=True)
os.makedirs(os.path.join(split_dir, 'test/images'), exist_ok=True)
os.makedirs(os.path.join(split_dir, 'test/labels'), exist_ok=True)

# 获取图片文件列表
image_files = os.listdir(os.path.join(root_dir, 'images'))
label_files = os.listdir(os.path.join(root_dir, 'labels'))

# 随机打乱文件列表
combined_files = list(zip(image_files, label_files))
random.shuffle(combined_files)
image_files_shuffled, label_files_shuffled = zip(*combined_files)

# 根据比例计算划分的边界索引
train_bound = int(train_ratio * len(image_files_shuffled))
valid_bound = int((train_ratio + valid_ratio) * len(image_files_shuffled))

# 将图片和标签文件移动到相应的目录
for i, (image_file, label_file) in enumerate(zip(image_files_shuffled, label_files_shuffled)):
    if i < train_bound:
        shutil.move(os.path.join(root_dir, 'images', image_file), os.path.join(split_dir, 'train/images', image_file))
        shutil.move(os.path.join(root_dir, 'labels', label_file), os.path.join(split_dir, 'train/labels', label_file))
    elif i < valid_bound:
        shutil.move(os.path.join(root_dir, 'images', image_file), os.path.join(split_dir, 'valid/images', image_file))
        shutil.move(os.path.join(root_dir, 'labels', label_file), os.path.join(split_dir, 'valid/labels', label_file))
    else:
        shutil.move(os.path.join(root_dir, 'images', image_file), os.path.join(split_dir, 'test/images', image_file))
        shutil.move(os.path.join(root_dir, 'labels', label_file), os.path.join(split_dir, 'test/labels', label_file))

参考文献：

[1]LabelImg标注的VOC格式xml标签与YOLO格式txt标签相互转换

[2]手把手实现 | 使用Yolov8训练自己的数据集【环境配置-准备数据集(采集&标注&划分)-模型训练(多种方式)-模型预测-模型导出】