如何将患者文件夹整体拆分至训练、验证、测试数据集？

我有一批患者的肌肉MRI图像及掩码真值，基于U-Net做肌肉自动分割与尺寸测量。最初直接拆分所有患者的切片，发现模型可能在同一患者的数据上同时训练和测试，导致数据泄露，影响准确率。

当前数据目录结构如下：

PATIENT_FOLDERS
├───01
│       FATFRACTION_01_10_pre.png
│       ...
│       psoas_01_9_pre.png
│
└───03
        FATFRACTION_03_25_pre.png
        ...
        psoas_03_48_pre.png

我尝试用splitfolders库拆分：

splitfolders.ratio(patient_folders, output=preprocessed_dir, seed=1337, ratio=(.8, 0.1,0.1)) 

但这个方法只会拆分文件夹里的单个文件，没法把整个患者文件夹作为整体分到训练、验证、测试组。请问怎么实现按患者文件夹整体拆分？

解决方案 #

方法1：手动编写Python代码实现按患者文件夹拆分 #

核心思路是先获取所有患者文件夹的列表，随机打乱后按比例拆分，最后将整个文件夹复制到对应输出目录，完全保证同一患者的所有数据只属于一个数据集。

示例代码：

import os
import random
import shutil
from pathlib import Path

# 配置参数
patient_root = "PATIENT_FOLDERS"  # 患者文件夹根目录
output_dir = "preprocessed_dir"   # 最终输出目录
split_ratio = (0.8, 0.1, 0.1)     # 训练/验证/测试集比例
seed = 1337                       # 随机种子，保证拆分结果可复现

# 设置随机种子
random.seed(seed)

# 获取所有患者文件夹（过滤掉根目录下的非目录文件）
patient_folders = [
    folder for folder in os.listdir(patient_root)
    if os.path.isdir(os.path.join(patient_root, folder))
]
random.shuffle(patient_folders)  # 打乱患者顺序，确保拆分随机性

# 计算各数据集的患者数量
total_patients = len(patient_folders)
train_num = int(total_patients * split_ratio[0])
val_num = int(total_patients * split_ratio[1])

# 拆分患者文件夹列表
train_folders = patient_folders[:train_num]
val_folders = patient_folders[train_num:train_num+val_num]
test_folders = patient_folders[train_num+val_num:]

# 创建输出目录结构
for split_type in ["train", "val", "test"]:
    Path(os.path.join(output_dir, split_type)).mkdir(parents=True, exist_ok=True)

# 批量复制患者文件夹到对应数据集目录
def copy_patient_folders(source_folders, target_dir):
    for folder in source_folders:
        src_path = os.path.join(patient_root, folder)
        dest_path = os.path.join(target_dir, folder)
        shutil.copytree(src_path, dest_path)

copy_patient_folders(train_folders, os.path.join(output_dir, "train"))
copy_patient_folders(val_folders, os.path.join(output_dir, "val"))
copy_patient_folders(test_folders, os.path.join(output_dir, "test"))

方法2：利用splitfolders的分组参数（局限性较大） #

splitfolders的group_prefix参数可按文件名前缀分组，但仅适用于所有文件命名严格遵循同一前缀规则（比如每个患者的文件都以患者ID开头），且无法直接识别文件夹。这种方法可靠性不如手动拆分，仅作为补充方案。

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