Skip to content
0

文章发布较早,内容可能过时,阅读注意甄别。

数据集划分

==以下讨论两种情况==

1、已分类

1.1、描述

假设一个data文件夹下有一堆文件夹,每个文件夹下有图片,现在需要新建一个Dataset文件夹,里面有train和test文件夹,然后将data文件夹下的文件进行随机分层采样,分成训练集和测试集,放到对应的train和test文件夹中

1.2、目录结构

1.2.1、data文件夹:

python
data/

├── class1/
│   ├── img1.jpg
│   ├── img2.jpg
│   ├── img3.jpg
│   └── ...

├── class2/
│   ├── img4.jpg
│   ├── img5.jpg
│   ├── img6.jpg
│   └── ...

└── class3/
    ├── img7.jpg
    ├── img8.jpg
    ├── img9.jpg
    └── ...

1.2.2、Dataset文件夹:

Dataset/

├── train/
│   ├── class1/
│   │   ├── img1.jpg
│   │   ├── img3.jpg
│   │   └── ...
│   ├── class2/
│   │   ├── img5.jpg
│   │   ├── img6.jpg
│   │   └── ...
│   └── class3/
│       ├── img7.jpg
│       ├── img9.jpg
│       └── ...

└── test/
    ├── class1/
    │   ├── img2.jpg
    │   └── ...
    ├── class2/
    │   ├── img4.jpg
    │   └── ...
    └── class3/
        ├── img8.jpg
        └── ...

1.3、代码实现

python
from pathlib import Path
import shutil
from sklearn.model_selection import train_test_split

# 定义路径
data_dir = Path('data')  # 原始数据文件夹
output_dir = Path('Dataset')  # 新的数据集文件夹
train_dir = output_dir / 'train'
test_dir = output_dir / 'test'

# 定义训练集与测试集的比例
train_ratio = 0.8

# 创建train和test文件夹
train_dir.mkdir(parents=True, exist_ok=True)
test_dir.mkdir(parents=True, exist_ok=True)

# 遍历data文件夹下的所有子文件夹
for folder_path in data_dir.iterdir():
    if folder_path.is_dir():  # 确保只处理文件夹
        # 获取当前文件夹中的所有图片文件
        image_files = list(folder_path.glob('*'))  # 获取所有文件路径
        image_files = [f for f in image_files if f.is_file()]  # 过滤掉非文件的路径

        # 使用train_test_split进行数据集划分
        train_images, test_images = train_test_split(image_files, train_size=train_ratio, random_state=42)

        # 创建对应的train和test子文件夹
        train_folder = train_dir / folder_path.name
        test_folder = test_dir / folder_path.name
        train_folder.mkdir(parents=True, exist_ok=True)
        test_folder.mkdir(parents=True, exist_ok=True)

        # 复制图片到对应的文件夹
        for image_path in train_images:
            shutil.copy(image_path, train_folder / image_path.name)

        for image_path in test_images:
            shutil.copy(image_path, test_folder / image_path.name)

print("数据集划分完成!")

其中数据集划分也可以用以下代码实现

python
# 随机打乱图片顺序
random.shuffle(images)

# 按比例分配图片到train和test
train_count = int(len(images) * train_ratio)
train_images = images[:train_count]
test_images = images[train_count:]

2、未分类

2.1、描述

假设有一个Data文件夹目录,里面存放着各种图片文件,现在需要将其进行随即划分,生成一个Dataset文件夹目录,里面有traintest目录,之后有10个子目录,每个子目录下存放着图片

2.2、目录结构

python
Dataset/
|-- train/
|   |-- subfolder_1/
|   |   |-- image1.jpg
|   |   |-- image2.jpg
|   |   |-- ...
|   |-- subfolder_2/
|   |   |-- image3.jpg
|   |   |-- ...
|   |-- subfolder_3/
|   |-- subfolder_4/
|   |-- subfolder_5/
|   |-- subfolder_6/
|   |-- subfolder_7/
|   |-- subfolder_8/
|   |-- subfolder_9/
|   |-- subfolder_10/
|-- test/
|   |-- subfolder_1/
|   |   |-- image4.jpg
|   |   |-- image5.jpg
|   |   |-- ...
|   |-- subfolder_2/
|   |   |-- image6.jpg
|   |   |-- ...
|   |-- subfolder_3/
|   |-- subfolder_4/
|   |-- subfolder_5/
|   |-- subfolder_6/
|   |-- subfolder_7/
|   |-- subfolder_8/
|   |-- subfolder_9/
|   |-- subfolder_10/

2.3、代码实现

python
from pathlib import Path
import random
from sklearn.model_selection import train_test_split


# 创建 10 个子文件夹,并将 train 和 test 中的文件随机分配到这些子文件夹中
def create_and_distribute_files(source_dir, num_subfolders=10):
    files = list(source_dir.glob('*'))
    random.shuffle(files)  # 打乱文件顺序
    subfolder_files = [[] for _ in range(num_subfolders)]

    for i, file in enumerate(files):
        subfolder_files[i % num_subfolders].append(file)

    for i in range(num_subfolders):
        subfolder_path = source_dir / f'subfolder_{i + 1}'
        subfolder_path.mkdir(parents=True, exist_ok=True)
        for file in subfolder_files[i]:
            destination = subfolder_path / file.name
            file.replace(destination)  # 移动文件

if __name__ == '__main__':
    # 定义路径和比例
    dataset_dir = Path('Data')  # 原始数据集文件夹路径
    train_dir = Path('./Dataset/train')  # 训练集文件夹路径
    test_dir = Path('./Dataset/test')  # 测试集文件夹路径
    split_ratio = 0.8  # 训练集占比,80%
    
    # 创建 train 和 test 文件夹
    train_dir.mkdir(parents=True, exist_ok=True)
    test_dir.mkdir(parents=True, exist_ok=True)
    
    # 列出 dataset 文件夹中的所有文件
    image_files = list(dataset_dir.glob('*'))  # 获取所有文件路径
    image_files = [f for f in image_files if f.is_file()]  # 过滤掉非文件的路径
    
    # 将图片文件按比例分为 train 和 test 两部分,可以指定随机种子 random_state,确保划分结果可复现
    train_files, test_files = train_test_split(image_files, train_size=split_ratio, random_state=42)
    
    # 将文件复制到相应的文件夹中
    for file in train_files:
        destination = train_dir / file.name
        destination.write_bytes(file.read_bytes())  # 复制文件
    
    for file in test_files:
        destination = test_dir / file.name
        destination.write_bytes(file.read_bytes())  # 复制文件
    
    # 在 train 和 test 文件夹中分别创建 10 个子文件夹并分配文件
    create_and_distribute_files(train_dir, num_subfolders=10)
    create_and_distribute_files(test_dir, num_subfolders=10)
    
    print("数据集划分完成")

其中,数据集的划分有很多方法,上面用的是,可以指定随机种子random_state,确保划分结果可复现

python
train_files, test_files = train_test_split(image_files, train_size=split_ratio, random_state=42)

还可以使用如下方法,但是没有指定随机种子:

1、

python
# 使用 random 进行划分
random.shuffle(image_files)  # 随机打乱文件顺序
split_index = int(len(image_files) * split_ratio)
train_files = image_files[:split_index]  # 训练集文件
test_files = image_files[split_index:]   # 测试集文件

2、

python
# 使用 random.sample 进行划分
num_train = int(len(image_files) * split_ratio)
train_files = random.sample(image_files, num_train)  # 随机抽取 num_train 个文件作为训练集
test_files = [f for f in image_files if f not in train_files]  # 剩下的文件作为测试集
最近更新