基于ID从多JPG图片集随机抽单图划分集合的Python代码需求

我的数据集包含大量jpg格式图片，命名规则为[ID]_[Cam]_[Frame].jpg。每个ID对应的图片数量不同，且ID并非连续编号（存在跳号情况，如示例中无ID2、3）。需编写Python代码实现：从每个ID对应的图片中随机抽取1张放入ValidationSet，剩余图片保留在TrainSet。

处理前示例 #

TrainSet包含以下文件：

• 00000000_0001_00000000.jpg
• 00000000_0001_00000001.jpg
• 00000000_0002_00000001.jpg
• 00000001_0001_00000001.jpg
• 00000001_0002_00000001.jpg
• 00000001_0002_00000002.jpg
• 00000004_0001_00000001.jpg
• 00000004_0002_00000001.jpg

处理后示例 #

TrainSet文件 #

• 00000000_0001_00000000.jpg
• 00000000_0002_00000001.jpg
• 00000001_0002_00000001.jpg
• 00000001_0001_00000001.jpg
• 00000004_0001_00000001.jpg

ValidationSet文件 #

• 00000000_0001_00000001.jpg
• 00000001_0002_00000002.jpg
• 00000004_0002_00000001.jpg

以下是实现需求的Python代码：

import os
import random
import shutil

def split_train_val(train_dir, val_dir):
    # 确保验证集目录存在
    os.makedirs(val_dir, exist_ok=True)
    
    # 按ID分组所有图片
    id_image_groups = {}
    for file_name in os.listdir(train_dir):
        if file_name.endswith(".jpg"):
            # 拆分文件名提取ID部分
            id_part = file_name.split("_")[0]
            if id_part not in id_image_groups:
                id_image_groups[id_part] = []
            id_image_groups[id_part].append(file_name)
    
    # 遍历每个ID组，随机选一张移到验证集
    for _, image_list in id_image_groups.items():
        if len(image_list) >= 1:
            selected_image = random.choice(image_list)
            # 构建文件路径并移动
            src_path = os.path.join(train_dir, selected_image)
            dst_path = os.path.join(val_dir, selected_image)
            shutil.move(src_path, dst_path)

# 替换为你的实际路径
TRAIN_SET_PATH = "./TrainSet"
VALIDATION_SET_PATH = "./ValidationSet"

split_train_val(TRAIN_SET_PATH, VALIDATION_SET_PATH)

代码说明 #

1. 分组逻辑：通过拆分文件名的第一个下划线前的内容作为ID，将同ID的图片归类到同一组
2. 随机抽取：对每个ID组，用random.choice()随机选中一张图片
3. 文件移动：使用shutil.move()完成图片从训练集到验证集的迁移，剩余图片自动留在训练集
4. 目录保障：用os.makedirs()确保验证集目录存在，避免运行报错

内容的提问来源于stack exchange，提问作者Wildan Shafa