划分图片数据集及其对应的标签

当我们拿到一个数据进行yolo训练时，如果只有图片和标签两个文件，我们希望他们各自都有train，test两个文件。并且两个文件里面图片和标签是一一对应的关系，那么我们可以用以下代码实现：最初，本代码图片和标签的数据都在train文件，test文件为空文件。现在将划分的部分图片，对应的标签移到各自test文件里去。

import os
import random
import shutil


def move_images_and_labels(source_image_dir, source_label_dir, target_image_dir, target_label_dir, rate=0.3):
    # 获取源图片文件夹中的文件列表
    image_filenames = os.listdir(source_image_dir)
    # 计算要移动的图片文件数量
    num_files_to_move = int(len(image_filenames) * rate)
    # 随机选择要移动的图片文件
    files_to_move = random.sample(image_filenames, num_files_to_move)

    # 移动选定的图片文件到目标图片文件夹，并移动对应的标签文件到目标标签文件夹
    for filename in files_to_move:
        # 移动图片文件
        source_image_path = os.path.join(source_image_dir, filename)
        target_image_path = os.path.join(target_image_dir, filename)
        shutil.move(source_image_path, target_image_path)
        print(f"Moved image {filename} to {target_image_path}")

        # 移动对应的标签文件
        label_filename = os.path.splitext(filename)[0] + ".jpg.txt"  #十分重要，注意看清自        己的图片和对应的标签的名称的区别
        source_label_path = os.path.join(source_label_dir, label_filename)
        target_label_path = os.path.join(target_label_dir, label_filename)

        # 检查标签文件是否存在，如果存在，则移动
        if os.path.exists(source_label_path):
            shutil.move(source_label_path, target_label_path)
            print(f"Moved corresponding label file {label_filename} to {target_label_path}")
        else:
            print(f"No corresponding label file found for {filename}")


if __name__ == '__main__':
    source_image_train_dir = r"xxx"  #待划分的原图片文件
    source_label_train_dir = r"xxx"  #待划分的原标签文件
    target_image_test_dir = r"xxx"  #划分的新图片文件
    target_label_test_dir = r"xxx"  #划分的新标签文件

    move_images_and_labels(source_image_train_dir, source_label_train_dir, target_image_test_dir, target_label_test_dir,
                           rate=0.3)

需要注意的是自己的数据的图片名称和对应的标签的名称的区别十分重要，本例中图片名称为数字.jpg形式，而标签名称为对应的数字.jpg.txt形式，这里需要根据自己数据情况进行修改。