如何将同标识组的三张图片绑定后划分训练/验证/测试集

核心逻辑是以组标识n作为划分的最小单位，而非单张图片，从根源上保证同组三张图片不会被分到不同数据集。

实现步骤 #

1. 按组聚合所有图片 #

首先遍历图片存储文件夹，将同组的3张图片关联存储，这里用字典实现最方便，键为组号n，值为同组三张图片的路径列表：

import os
from collections import defaultdict
import random

img_dir = "CM_Images"
# 初始化字典存储分组，key是组号n，value是三张图路径
group_dict = defaultdict(list)

for img_name in os.listdir(img_dir):
    # 只处理png图片，过滤其他文件
    if img_name.endswith(".png"):
        # 拆分文件名得到组号n和组内序号m
        n, m = img_name.replace(".png", "").split("_")
        n = int(n)
        img_path = os.path.join(img_dir, img_name)
        group_dict[n].append(img_path)

# 可选校验：每个组必须有3张图，过滤异常数据
group_dict = {n: paths for n, paths in group_dict.items() if len(paths) == 3}
# 拿到所有组号的列表
all_groups = list(group_dict.keys())

2. 按比例划分组 #

直接对组号列表做随机打乱，再按训练/验证/测试的比例切分即可：

# 打乱组顺序，保证划分随机
random.shuffle(all_groups)
total = len(all_groups)
# 自定义划分比例，这里用常见的7:2:1
train_ratio = 0.7
val_ratio = 0.2
test_ratio = 0.1

train_end = int(total * train_ratio)
val_end = train_end + int(total * val_ratio)

train_groups = all_groups[:train_end]
val_groups = all_groups[train_end:val_end]
test_groups = all_groups[val_end:]

# 生成最终的数据集图片路径列表
train_imgs = [path for n in train_groups for path in group_dict[n]]
val_imgs = [path for n in val_groups for path in group_dict[n]]
test_imgs = [path for n in test_groups for path in group_dict[n]]

3. 可选：配合深度学习框架加载数据 #

如果需要在训练时直接拿到同组的三张图作为样本，可以自定义数据集类，以下是PyTorch示例：

from torch.utils.data import Dataset
from PIL import Image

class GroupImageDataset(Dataset):
    def __init__(self, group_list, group_dict, transform=None):
        self.group_list = group_list
        self.group_dict = group_dict
        self.transform = transform

    def __len__(self):
        return len(self.group_list)

    def __getitem__(self, idx):
        group_n = self.group_list[idx]
        # 按序号排序，保证img1是n_1, img2是n_2, img3是n_3
        img_paths = sorted(self.group_dict[group_n], key=lambda x: int(x.split("_")[-1].replace(".png","")))
        imgs = [Image.open(p).convert("RGB") for p in img_paths]
        if self.transform:
            imgs = [self.transform(img) for img in imgs]
        # 返回同组三张图和组号
        return imgs[0], imgs[1], imgs[2], group_n

现有代码优化建议 #

你已生成的Array_of_named_Tuples已经存储了所有组的标识和对应矩阵，完全可以跳过读取文件夹分组的步骤，直接对这个数组做划分，效率更高：

# 直接打乱命名元组数组
random.shuffle(Array_of_named_Tuples)
total = len(Array_of_named_Tuples)
train_end = int(total * 0.7)
val_end = train_end + int(total * 0.2)

train_sets = Array_of_named_Tuples[:train_end]
val_sets = Array_of_named_Tuples[train_end:val_end]
test_sets = Array_of_named_Tuples[val_end:]

这样你既可以拿到组的矩阵数据，也可以通过组号拼接出对应图片路径，完全满足绑定要求。

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