【Python快速入门和实践013】Python常用脚本-目标检测之按照类别数量划分数据集

一、功能介绍

        这段代码实现了从给定的图像和标签文件夹中分割数据集为训练集、验证集和测试集的功能。以下是代码功能的总结：

1. 创建目标文件夹结构：

   • 在指定的根目录（dataset_root）下创建images和labels两个文件夹。
   • 在这两个文件夹下分别创建train、val和test三个子文件夹，用于存放不同阶段的数据。

2. 统计类别数量：

   • 遍历标签文件夹中的所有文本文件，统计每个类别在所有标签文件中出现的总次数。

3. 计算分割比例：

   • 根据给定的比例（默认为训练集80%，验证集10%，测试集10%），计算每个类别在训练集、验证集和测试集中应该有的数量。

4. 随机分配数据：

   • 遍历图像文件夹中的所有图片。
   • 对于每个图片，检查其对应的标签文件是否存在。
   • 读取标签文件，提取其中的类别信息。
   • 根据随机数决定图片属于训练集、验证集还是测试集。
   • 将图片和对应的标签文件复制到相应的文件夹中，同时更新类别数量记录。

5. 最终结果：

   • 数据集按照指定的比例被划分为训练集、验证集和测试集。
   • 每个类别在各个数据集中的分布尽量保持均衡。

二、代码

import os
import random
import shutil


def split_dataset(image_folder, label_folder, train_ratio=0.8, val_ratio=0.1, test_ratio=0.1):
    """
    将图像和标签文件按指定比例分割成训练集、验证集和测试集。

    参数:
    image_folder (str): 图像文件夹路径。
    label_folder (str): 标签文件夹路径。
    train_ratio (float): 训练集所占比例，默认为0.8。
    val_ratio (float): 验证集所占比例，默认为0.1。
    test_ratio (float): 测试集所占比例，默认为0.1。
    """

    # 创建目标文件夹
    dataset_root = r'E:\pythonProject\pythonProject\after_neu'
    os.makedirs(dataset_root, exist_ok=True)

    # 创建images和labels文件夹
    images_folder = os.path.join(dataset_root, 'images')
    labels_folder = os.path.join(dataset_root, 'labels')
    os.makedirs(images_folder, exist_ok=True)
    os.makedirs(labels_folder, exist_ok=True)

    # 创建train、val和test子文件夹
    for split in ['train', 'val', 'test']:
        os.makedirs(os.path.join(images_folder, split), exist_ok=True)
        os.makedirs(os.path.join(labels_folder, split), exist_ok=True)

    # 统计每个类别的图片数量
    category_counts = {}
    for filename in os.listdir(label_folder):
        label_path = os.path.join(label_folder, filename)
        with open(label_path, 'r') as label_file:
            lines = label_file.readlines()
            categories = [line.split()[0] for line in lines]
            for category in categories:
                category_counts[category] = category_counts.get(category, 0) + 1

    # 计算每个类别在训练集、验证集和测试集中的数量
    train_category_counts = {}
    val_category_counts = {}
    test_category_counts = {}
    for category, count in category_counts.items():
        train_count = int(count * train_ratio)
        val_count = int(count * val_ratio)
        test_count = count - train_count - val_count
        train_category_counts[category] = train_count
        val_category_counts[category] = val_count
        test_category_counts[category] = test_count

    # 遍历图片文件夹
    for filename in os.listdir(image_folder):
        image_path = os.path.join(image_folder, filename)
        label_path = os.path.join(label_folder, os.path.splitext(filename)[0] + '.txt')

        # 确保标注文件存在
        if not os.path.exists(label_path):
            continue

        # 读取标注文件获取类别信息
        with open(label_path, 'r') as label_file:
            lines = label_file.readlines()
            categories = [line.split()[0] for line in lines]

        # 确定将图片放入的集合
        rand = random.random()
        if rand < train_ratio:
            destination_folder = 'train'
            category_counts = train_category_counts
        elif rand < train_ratio + val_ratio:
            destination_folder = 'val'
            category_counts = val_category_counts
        else:
            destination_folder = 'test'
            category_counts = test_category_counts

        # 移动图片和标注文件到目标文件夹
        for category in categories:
            category_folder_images = os.path.join(images_folder, destination_folder)
            category_folder_labels = os.path.join(labels_folder, destination_folder)
            os.makedirs(category_folder_images, exist_ok=True)
            os.makedirs(category_folder_labels, exist_ok=True)
            if category_counts[category] > 0:
                shutil.copy(image_path, os.path.join(category_folder_images, filename))
                shutil.copy(label_path, os.path.join(category_folder_labels, os.path.splitext(filename)[0] + '.txt'))
                category_counts[category] -= 1


# 图片文件夹路径
image_folder = r'E:\pythonProject\pythonProject\NEU-DET\images'

# 标注文件夹路径
label_folder = r'E:\pythonProject\pythonProject\NEU-DET\txt'

# 调用函数进行数据集分割
split_dataset(image_folder, label_folder)

        这个数据集划分代码相比与其他的不是随机划分，考虑到每个类别的图片样张可能不均衡，所以按照类别去划分数据集。需要先把xml转成yolo的txt格式，然后指定图片、txt标签、保存文件夹路径即可。在NEU-DET数据集上运行结果如下：