在YOLO数据集总数中提取出固定数量的图片和对应的标签

在YOLO格式数据集中，提取出部分固定数量的图片和对应的标签进行用于模型训练，例如提取出10000张图片

import os
import random
import shutil
 
# 图片和标签文件所在的目录
image_folder = 'path_to_images_folder'
label_folder = 'path_to_labels_folder'
 
# 创建目标文件夹来保存随机选出的10000张图片和标签文件
output_image_folder = 'path_to_output_images_folder'
output_label_folder = 'path_to_output_labels_folder'
 
# 如果目标文件夹不存在，则创建它们
if not os.path.exists(output_image_folder):
    os.makedirs(output_image_folder)
 
if not os.path.exists(output_label_folder):
    os.makedirs(output_label_folder)
 
# 获取所有图片文件和标签文件列表
image_files = [f for f in os.listdir(image_folder) if f.endswith(('jpg', 'jpeg', 'png', 'bmp', 'gif'))]
 
# 检查是否有足够的图片
if len(image_files) < 10000:
    print(f"警告：数据集中的图片数量只有 {len(image_files)} 张，少于 10000 张")
    num_images_to_select = len(image_files)
else:
    num_images_to_select = 10000
 
# 随机选择10000张图片（如果有足够数量的话）
selected_images = random.sample(image_files, num_images_to_select)
 
# 复制选中的图片和对应的标签文件到目标文件夹
for image in selected_images:
    # 生成对应标签文件的文件名
    # label_file = image.replace(image.split('.')[-1], 'txt')  # 假设标签文件与图片同名，只是扩展名为 .txt
    label_file = os.path.splitext(image)[0] + '.txt' # 处理复杂情况的命名方式
    source_image_path = os.path.join(image_folder, image)
    source_label_path = os.path.join(label_folder, label_file)
 
    # 检查标签文件是否存在
    if os.path.exists(source_label_path):
        # 复制图片和标签文件
        shutil.copy(source_image_path, os.path.join(output_image_folder, image))
        shutil.copy(source_label_path, os.path.join(output_label_folder, label_file))
    else:
        print(f"警告：找不到标签文件 {label_file}，跳过该图片。")
 
print(f"成功从数据集中随机选出 {num_images_to_select} 张图片及其标签文件，并复制到目标文件夹。")