Yolov5(一)VOC划分数据集、VOC转YOLO数据集

代码使用方法注意修改一下路径、验证集比例、类别名称，其他均不需要改动，自动划分训练集、验证集、建好全部文件夹、一键自动生成Yolo格式数据集在当前目录下，大家可以直接修改相应的配置文件进行训练。

目录

使用方法：

 全部代码：

使用方法：

 全部代码：

import os, random, shutil

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


def convert(size, box):
    x_center = (box[0] + box[1]) / 2.0
    y_center = (box[2] + box[3]) / 2.0

    x = x_center / size[0]
    y = y_center / size[1]
    w = (box[1] - box[0]) / size[0]
    h = (box[3] - box[2]) / size[1]
    return (x, y, w, h)


def convert_annotation(xml_file_path, save_txt_file_path, classes):
    xml_file = os.listdir(xml_file_path)
    print(xml_file)
    for xml_name in xml_file:
        print(xml_file)

        xml_file = os.path.join(xml_file_path, xml_name)
        out_txt_path = os.path.join(save_txt_file_path, xml_name.split('.')[0] + ".txt")
        out_txt_f = open(out_txt_path, 'w')
        tree = ET.parse(xml_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"):
            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))

            print(w, h, b)
            bb = convert((w, h), b)
            out_txt_f.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + "\n")


def moving(fileDir, tarDir, rate=0.2):
    pathDir = os.listdir(fileDir)
    filenumber = len(pathDir)
    # 自定义抽取图片比例
    picknumber = int(filenumber * rate)  # 按照rate比例从文件夹中取一定数量的图片
    sample = random.sample(pathDir, picknumber)  # 随机选取picknumber数量的图片
    print(sample)

    for name in sample:
        shutil.move(fileDir + "/" + name, tarDir + "/" + name)
    return


def movelabel(file_list, file_label_train, file_label_val):
    for i in file_list:
        if i.endswith(".png") or i.endswith(".jpg"):
            # filename=file_label_train+"/"+i[:-4]       可以将.xml文件将.txt改成.xml文件
            filename = file_label_train + "/" + i[:-4] + ".xml"  # 可以改成xml文件将.txt改成.xml
            if os.path.exists(filename):
                shutil.move(filename, file_label_val)
                print("处理成功")


if __name__ == "__main__":
    """
    设置图片路径、label路径、验证集比例、类别
    """
    fileDir = r"./JPEGImages"  # 图片的路径
    file_label_train = r"./Annotations"  # 标签文件的路径
    rate = 0.2  # 验证集的比例
    classes1 = ['ship']

    """
    以下均不需要改动
    """

    if not os.path.exists("./JPEGImages_val"):
        # Create the folder
        os.makedirs("./JPEGImages_val")

    tarDir = r"./JPEGImages_val"
    moving(fileDir, tarDir, rate)
    file_list = os.listdir(tarDir)

    if not os.path.exists("./Annotations_val"):
        # Create the folder
        os.makedirs("./Annotations_val")

    file_label_val = r"./Annotations_val"
    movelabel(file_list, file_label_train, file_label_val)

    # VOC转Yolo格式

    # 2 voc 格式的xml标签文件
    if not os.path.exists("./val"):
        # Create the folder
        os.makedirs("./val")

    if not os.path.exists("./train"):
        # Create the folder
        os.makedirs("./train")

    xml_files1 = r"./Annotations_val"
    save_txt_files1 = r"./val"

    convert_annotation(xml_files1, save_txt_files1, classes1)

    xml_files1 = r"./Annotations"
    save_txt_files1 = r"./train"

    convert_annotation(xml_files1, save_txt_files1, classes1)

    # 创建所有文件夹
    if not os.path.exists("./images"):
        # Create the folder
        os.makedirs("./images")

    if not os.path.exists("./labels"):
        # Create the folder
        os.makedirs("./labels")

    # 将所有文件移动到最终的文件夹中

    import shutil

    # Define the source and destination folders
    source_folder = "./train"
    destination_folder = "./labels"

    # Move the files from the source folder to the destination folder
    shutil.move(source_folder, destination_folder)

    source_folder = "./val"
    destination_folder = "./labels"

    # Move the files from the source folder to the destination folder
    shutil.move(source_folder, destination_folder)

    source_folder = "./JPEGImages"
    destination_folder = "./images"

    # Move the files from the source folder to the destination folder
    shutil.move(source_folder, destination_folder)
    os.rename("./images/JPEGImages", "./images/train")

    source_folder = "./JPEGImages_val"
    destination_folder = "./images"

    # Move the files from the source folder to the destination folder
    shutil.move(source_folder, destination_folder)
    os.rename("./images/JPEGImages_val", "./images/val")

输出xml的所有类别

import os
import xml.etree.ElementTree as ET

# 定义一个函数来遍历指定目录下的所有 XML 文件
def get_classes_from_xml(xml_directory):
    # 存储所有类别
    all_classes = set()

    # 获取指定目录下的所有文件
    for file in os.listdir(xml_directory):
        if file.endswith('.xml'):  # 只处理 XML 文件
            xml_file = os.path.join(xml_directory, file)  # 获取文件的完整路径
            # 解析 XML 文件
            tree = ET.parse(xml_file)  # 确保传递的是文件路径
            root = tree.getroot()

            # 假设类别存储在 <name> 标签中
            for obj in root.findall('.//object'):  # 遍历所有 <object> 标签
                name = obj.find('name').text
                all_classes.add(name)  # 将类别添加到集合中，确保唯一

    # 输出所有类别
    print("所有类别：")
    for cls in all_classes:
        print(cls)

# 调用函数，传入 XML 文件所在的目录
xml_directory = './Annotations'  # 假设 XML 文件存储在 'Annotations' 目录下
get_classes_from_xml(xml_directory)