按指定比例划分数据集为训练集、验证集、测试集,并将结果保存在三个文件夹中
时间: 2024-05-14 13:13:57 浏览: 8
以下是一个示例代码,可以将数据集按照 6:2:2 的比例划分为训练集、验证集和测试集,并将结果保存在三个文件夹中。
```python
import os
import shutil
# 设置数据集路径和划分比例
data_dir = '/path/to/data'
train_ratio = 0.6
val_ratio = 0.2
test_ratio = 0.2
# 创建保存训练集、验证集和测试集的文件夹
train_dir = os.path.join(data_dir, 'train')
val_dir = os.path.join(data_dir, 'val')
test_dir = os.path.join(data_dir, 'test')
os.makedirs(train_dir, exist_ok=True)
os.makedirs(val_dir, exist_ok=True)
os.makedirs(test_dir, exist_ok=True)
# 遍历数据集中的每个类别文件夹
for class_name in os.listdir(data_dir):
class_dir = os.path.join(data_dir, class_name)
if os.path.isdir(class_dir):
# 计算每个类别的样本数和划分后的样本数
num_samples = len(os.listdir(class_dir))
num_train = int(num_samples * train_ratio)
num_val = int(num_samples * val_ratio)
num_test = num_samples - num_train - num_val
# 随机打乱每个类别的样本顺序
samples = os.listdir(class_dir)
samples.sort()
samples = [os.path.join(class_dir, s) for s in samples]
shuffled_samples = samples.copy()
random.shuffle(shuffled_samples)
# 将划分后的样本复制到对应的文件夹中
train_samples = shuffled_samples[:num_train]
val_samples = shuffled_samples[num_train:num_train+num_val]
test_samples = shuffled_samples[num_train+num_val:]
for sample in train_samples:
dst = os.path.join(train_dir, class_name, os.path.basename(sample))
os.makedirs(os.path.dirname(dst), exist_ok=True)
shutil.copyfile(sample, dst)
for sample in val_samples:
dst = os.path.join(val_dir, class_name, os.path.basename(sample))
os.makedirs(os.path.dirname(dst), exist_ok=True)
shutil.copyfile(sample, dst)
for sample in test_samples:
dst = os.path.join(test_dir, class_name, os.path.basename(sample))
os.makedirs(os.path.dirname(dst), exist_ok=True)
shutil.copyfile(sample, dst)
```
上述代码假设每个类别的样本存储在以类别名称命名的文件夹中,如下所示:
```
data/
├── class1/
│ ├── sample1.jpg
│ ├── sample2.jpg
│ └── ...
├── class2/
│ ├── sample1.jpg
│ ├── sample2.jpg
│ └── ...
└── ...
```
划分后的数据集将存储在如下所示的文件夹中:
```
data/
├── train/
│ ├── class1/
│ │ ├── sample1.jpg
│ │ ├── sample2.jpg
│ │ └── ...
│ ├── class2/
│ │ ├── sample1.jpg
│ │ ├── sample2.jpg
│ │ └── ...
│ └── ...
├── val/
│ ├── class1/
│ │ ├── sample1.jpg
│ │ ├── sample2.jpg
│ │ └── ...
│ ├── class2/
│ │ ├── sample1.jpg
│ │ ├── sample2.jpg
│ │ └── ...
│ └── ...
└── test/
├── class1/
│ ├── sample1.jpg
│ ├── sample2.jpg
│ └── ...
├── class2/
│ ├── sample1.jpg
│ ├── sample2.jpg
│ └── ...
└── ...
```