YOLOV5 数据集的划分和格式转换

1 数据划分

    标注好后的数据集分为两个文件夹，一个文件夹中放置图片另一个文件夹中放置了txt文件。将数据集按照一定的比例进行划分为训练集，测试集，验证集（train、test、val），下述的代码中我按照了 8:1:1 的比例去划分，若想根据自己的需要去划分数据集，修改下述代码中的

split_ratio=(0.8, 0.1, 0.1)

 可进行划分。一般的划分数据集的比例为 7:2:1和8:1:1，因为我这里的数据集比较少所以采用了      8:1:1的比例去划分。

image_folder_path = "D:\PycharmProjects\VOCdevkit\VOC2007\JPEGImages"
txt_folder_path = "D:\PycharmProjects\VOCdevkit\VOC2007\YOLOLabels"
output_dataset_path = "D:\PycharmProjects\VOCdevkit"

image_folder_path为你保存的图片的文件夹的路径

txt_folder_path 为你保存的txt文件夹的路径

output_dataset_path 为你保存的数据集的文件夹的路径，代码会在改路径下自动生成子文件夹，分别进行测试集，训练集，验证集的存储。

import os
import random
from shutil import copyfile

def split_dataset(image_folder, txt_folder, output_folder, split_ratio=(0.8, 0.1, 0.1)):
    # Ensure output folders exist
    for dataset in ['train', 'val', 'test']:
        if not os.path.exists(os.path.join(output_folder, dataset, 'images')):
            os.makedirs(os.path.join(output_folder, dataset, 'images'))
        if not os.path.exists(os.path.join(output_folder, dataset, 'txt')):
            os.makedirs(os.path.join(output_folder, dataset, 'txt'))

    # Get list of image files
    image_files = [f for f in os.listdir(image_folder) if f.endswith(('.jpg', '.jpeg', '.png'))]
    random.shuffle(image_files)

    num_images = len(image_files)
    num_train = int(split_ratio[0] * num_images)
    num_val = int(split_ratio[1] * num_images)

    train_images = image_files[:num_train]
    val_images = image_files[num_train:num_train + num_val]
    test_images = image_files[num_train + num_val:]

    # Copy images to respective folders
    for dataset, images_list in zip(['train', 'val', 'test'], [train_images, val_images, test_images]):
        for image_file in images_list:
            image_path = os.path.join(image_folder, image_file)
            copyfile(image_path, os.path.join(output_folder, dataset, 'images', image_file))
            txt_file = os.path.splitext(image_file)[0] + '.txt'
            txt_path = os.path.join(txt_folder, txt_file)

            # Copy corresponding txt file if exists
            if os.path.exists(txt_path):
                copyfile(txt_path, os.path.join(output_folder, dataset, 'txt', txt_file))

if __name__ == "__main__":
    image_folder_path = "D:\PycharmProjects\VOCdevkit\VOC2007\JPEGImages"
    txt_folder_path = "D:\PycharmProjects\VOCdevkit\VOC2007\YOLOLabels"
    output_dataset_path = "D:\PycharmProjects\VOCdevkit"

    split_dataset(image_folder_path, txt_folder_path, output_dataset_path)

2 格式的转换

  YOLOV5训练时所需要的文件格式为txt格式。我们在labeling标注数据时保存的文件格式和在网上下载的数据集如果是VOC（xml）格式时，需要对文件格式进行转换将xml文件转化成txt格式。

input_folder_path 为存储xml 文件的位置

output_folder_path 为保存转换后txt文件的位置

import os
import xml.etree.ElementTree as ET

def convert_folder_to_yolov5(input_folder, output_folder):
    # Ensure output folder exists
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    # Loop through each XML file in the input folder
    for xml_file_name in os.listdir(input_folder):
        if xml_file_name.endswith('.xml'):
            xml_file_path = os.path.join(input_folder, xml_file_name)

            # Generate corresponding output txt file path
            txt_file_name = os.path.splitext(xml_file_name)[0] + '.txt'
            txt_file_path = os.path.join(output_folder, txt_file_name)

            # Convert XML to Yolov5 format and save to txt file
            convert_to_yolov5(xml_file_path, txt_file_path)

def convert_to_yolov5(xml_file, output_file):
    tree = ET.parse(xml_file)
    root = tree.getroot()

    with open(output_file, 'w') as f:
        for obj in root.findall('object'):
            class_name = obj.find('name').text
            if class_name == 'disease':  # Assuming 'disease' is the class of interest
                xmin = int(obj.find('bndbox/xmin').text)
                ymin = int(obj.find('bndbox/ymin').text)
                xmax = int(obj.find('bndbox/xmax').text)
                ymax = int(obj.find('bndbox/ymax').text)

                width = xmax - xmin
                height = ymax - ymin
                x_center = (xmin + xmax) / 2.0
                y_center = (ymin + ymax) / 2.0

                # Normalize coordinates and dimensions
                x_center /= int(root.find('size/width').text)
                y_center /= int(root.find('size/height').text)
                width /= int(root.find('size/width').text)
                height /= int(root.find('size/height').text)

                line = f"{0} {x_center} {y_center} {width} {height}\n"
                f.write(line)

if __name__ == "__main__":
    input_folder_path = "D:\PycharmProjects\VOCdevkit/VOC2007/Annotations"
    output_folder_path = "D:\PycharmProjects\VOCdevkit\VOC2007\YOLOLabels"

    convert_folder_to_yolov5(input_folder_path, output_folder_path)