【数据集划分】支持 YOLO 系列模型训练的数据随机划分实现

1. 数据准备

• 按照 yolo 格式进行数据标注，完成后数据目录结构为：

  my_datasets/
  ├─images
  │	├──xxx.jpg		
  │	└──xxx.jpg		
  └──labels
  	├──xxx.txt	
  	└──xxx.txt
  

2. 方法一：划分数据集后生成 train.txt 和 val.txt

2.1 划分数据集

• 运行下述脚本即可将数据集随机划分为 train 和 val.

  '''
  实现功能：
      对输入数据集进行随机划分
  实现方式：
      @param: data_path       数据集路径
      @param: output_path     数据集输出路径
      @param: split_ratio     数据集划分比例
  Author:
      Lily-20231123
  '''
  import os
  import random
  import shutil
  import argparse
  
  def split_dataset(input_folder, output_folder, split_ratio=0.8):
      # 创建输出文件夹
      train_img_folder = os.path.join(output_folder, 'images', 'train')
      test_img_folder = os.path.join(output_folder, 'images', 'val')
      train_lab_folder = os.path.join(output_folder, 'labels', 'train')
      test_lab_folder = os.path.join(output_folder, 'labels', 'val')
      os.makedirs(train_img_folder, exist_ok=True)
      os.makedirs(test_img_folder, exist_ok=True)
      os.makedirs(train_lab_folder, exist_ok=True)
      os.makedirs(test_lab_folder, exist_ok=True)
  
      # 图像文件夹路径
      img_folder = os.path.join(input_folder, 'images')
      # 标签文件夹路径
      lab_folder = os.path.join(input_folder, 'labels')
  
      # 获取图像文件和标签文件列表
      img_files = os.listdir(img_folder)
      lab_files = os.listdir(lab_folder)
  
      # 计算训练集的数量
      train_size = int(len(img_files) * split_ratio)
  
      # 随机选择训练集和测试集
      train_img_files = random.sample(img_files, train_size)
      test_img_files = [img_file for img_file in img_files if img_file not in train_img_files]
  
      # 将图像和标签文件 copy 到相应的文件夹
      for img_file in train_img_files:
          shutil.copy(os.path.join(img_folder, img_file), os.path.join(train_img_folder, img_file))  
          print('train_img: ', os.path.join(train_img_folder, img_file))
          if os.path.exists(os.path.join(lab_folder, img_file[:-3]+"txt")):
              shutil.copy(os.path.join(lab_folder, img_file[:-3]+"txt"), os.path.join(train_lab_folder, img_file[:-3]+"txt"))
              print('train_lab: ', os.path.join(train_lab_folder, img_file[:-3]+"txt"))
  
      for img_file in test_img_files:
          shutil.copy(os.path.join(img_folder, img_file), os.path.join(test_img_folder, img_file))
          print('val_img: ', os.path.join(test_img_folder, img_file))
          if os.path.exists(os.path.join(lab_folder, img_file[:-3]+"txt")):
              shutil.copy(os.path.join(lab_folder, img_file[:-3]+"txt"), os.path.join(test_lab_folder, img_file[:-3]+"txt"))
              print('val_lab: ', os.path.join(test_lab_folder, img_file[:-3]+"txt"))
      
      print("train nums: ", train_size)
      print("val nums: ", int(len(img_files))-train_size)
  
  #TODO 数据集随机划分
  if __name__ == "__main__":
      
      parser = argparse.ArgumentParser(description='split datasets.')
      parser.add_argument('--data_path', type=str, default=r'D:\domain_apdation\smoke\images',
                          help="数据集输入路径")
      parser.add_argument('--output_path', type=str, default=r'D:\domain_apdation\smoke\images_split',
                          help="数据集输出路径")
      parser.add_argument('--split_ratio', type=float, default=0.8,
                          help='数据集划分比例')
      args = parser.parse_args()                             
      
      if not os.path.exists(args.output_path):
          os.makedirs(args.output_path)
          
      split_dataset(args.data_path, args.output_path, args.split_ratio)
  

• 完成划分后数据集目录结构为：

  my_datasets/
  ├─images
  │	├──train
  │	│	├──xxx.jpg		
  │	│	└──xxx.jpg		
  │	└──val
  │		├──xxx.jpg		
  │		└──xxx.jpg	
  └──labels
  	├──train
  	│	├──xxx.txt	
  	│	└──xxx.txt			
  	└──val
  		├──xxx.txt			
  		└──xxx.txt	
  

2.2 生成 train.txt 和 val.txt

• 若需要进一步生成 train.txt 和 val.txt，在数据集目录下运行下述 linux 指令即可。

  # 使用 ./images/train 生成为图片相对路径，若需生成绝对路径，修改该路径为绝对路径即可
  find ./images/train -name *.jpg >> train.txt
  find ./images/val -name *.jpg >> val.txt
  

• train.txt 文件展示
  

3. 方法二：直接生成 train.txt 和 val.txt

• 运行下述脚本即可将数据集直接划分为 train.txt 和 val.txt.

  """
  实现功能：
      将数据集按给定的比例随机划分成训练数据集 train.txt 和测试数据集 val.txt
  参数说明：
      @param path: 数据集路径
      @param ratio: 划分比例
  使用方法： 
      修改参数配置后使用 python split_train_val_txt.py 运行
  Author:
      Lily-20240223
  """
  import os
  import random
  import argparse
  
  # ! 划分数据集，输出 train.txt 和 val.txt
  def split_train_val_txt(path, ratio_train):
      # TODO 划分数据集的 txt 保存路径, 保存在数据集目录下
      trainFile_Path = os.path.join(os.path.dirname(path), 'train.txt')
      valFile_Path = os.path.join(os.path.dirname(path), 'val.txt')
      
      # TODO 获得文件夹下的所有文件名
      file_list = os.listdir(path)
      all_files_num = len(file_list)
  
      train_num = int(all_files_num * ratio_train)
      test_num = all_files_num - train_num
  
      # TODO 打乱文件顺序
      random.shuffle(file_list)
  
      # TODO 划分 train.txt
      train_write = open(trainFile_Path, 'w')
      for i in range(train_num):
          train_file_name = file_list[i]
          train_file_path = os.path.join(path, train_file_name)
          train_write.writelines(train_file_path)
          train_write.writelines('\n')
      train_write.close()
      
      # TODO 划分 val.txt
      val_write = open(valFile_Path, 'w')
      for i in range(test_num):
          val_file_name = file_list[i + train_num]
          val_file_path = os.path.join(path, val_file_name)
          val_write.writelines(val_file_path)
          val_write.writelines('\n')
      val_write.close()
      print("successful split train.txt and val.txt")
      return 
  
  
  if __name__ == '__main__':
      parser = argparse.ArgumentParser(description='split train.txt and test.txt.')
      parser.add_argument('--path', type=str, default=r'D:\domain_apdation\smoke\images',
                          help="数据集路径.")
      parser.add_argument('--ratio', type=float, default=0.8,
                          help='训练样本占比.')
      args = parser.parse_args()
      
      split_train_val_txt(args.path, args.ratio)
  

• 运行上述脚本后数据集目录结构为：

  my_datasets/
  ├─train.txt
  ├─val.txt
  ├─images
  │	├──xxx.jpg		
  │	└──xxx.jpg		
  └──labels
  	├──xxx.txt	
  	└──xxx.txt			
  

• train.txt 文件内容展示。