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

一、功能介绍

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

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

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

2. 统计类别数量：

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

3. 计算分割比例：

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

4. 随机分配数据：

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

5. 最终结果：

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

二、代码

import os
import random
import shutildef 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_counttrain_category_counts[category] = train_countval_category_counts[category] = val_counttest_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_countselif rand < train_ratio + val_ratio:destination_folder = 'val'category_counts = val_category_countselse: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数据集上运行结果如下：