我有两个文件夹,一个文件夹有500张猫的图像,另外一个文件夹有500张狗的图像。请帮我用python写出将这1000张图像按照8:1:1的比例划分成训练集、验证集和测试集,并且生成训练集、验证集和测试集三个文件夹,每个文件夹下有猫和狗两个子文件夹。
时间: 2024-03-05 15:47:59 浏览: 50
使用PyTorch将文件夹下的图片分为训练集和验证集实例
3星 · 编辑精心推荐
好的,以下是代码实现:
```python
import os
import random
import shutil
# 设置数据集路径
data_dir = '/path/to/dataset'
train_dir = '/path/to/train'
val_dir = '/path/to/val'
test_dir = '/path/to/test'
# 创建训练集、验证集和测试集文件夹
if not os.path.exists(train_dir):
os.makedirs(train_dir)
if not os.path.exists(val_dir):
os.makedirs(val_dir)
if not os.path.exists(test_dir):
os.makedirs(test_dir)
# 将猫和狗的图像分别放入两个列表中
cat_images = [os.path.join(data_dir, 'cat', img) for img in os.listdir(os.path.join(data_dir, 'cat'))]
dog_images = [os.path.join(data_dir, 'dog', img) for img in os.listdir(os.path.join(data_dir, 'dog'))]
# 将列表中的图像随机打乱
random.shuffle(cat_images)
random.shuffle(dog_images)
# 计算训练集、验证集和测试集的数量
total_images = len(cat_images) + len(dog_images)
train_num = int(total_images * 0.8)
val_num = int(total_images * 0.1)
# 将猫和狗的图像分别划分到训练集、验证集和测试集中
train_cat = cat_images[:int(train_num/2)]
train_dog = dog_images[:int(train_num/2)]
val_cat = cat_images[int(train_num/2):int(train_num/2)+int(val_num/2)]
val_dog = dog_images[int(train_num/2):int(train_num/2)+int(val_num/2)]
test_cat = cat_images[int(train_num/2)+int(val_num/2):]
test_dog = dog_images[int(train_num/2)+int(val_num/2):]
# 将训练集、验证集和测试集的猫和狗图像分别放入对应的文件夹中
for img in train_cat:
shutil.copy(img, os.path.join(train_dir, 'cat'))
for img in train_dog:
shutil.copy(img, os.path.join(train_dir, 'dog'))
for img in val_cat:
shutil.copy(img, os.path.join(val_dir, 'cat'))
for img in val_dog:
shutil.copy(img, os.path.join(val_dir, 'dog'))
for img in test_cat:
shutil.copy(img, os.path.join(test_dir, 'cat'))
for img in test_dog:
shutil.copy(img, os.path.join(test_dir, 'dog'))
```
请将 `data_dir` 替换成你数据集的路径,将 `train_dir`、`val_dir` 和 `test_dir` 替换成你想要保存训练集、验证集和测试集的路径。这段代码会将原始数据集中的猫和狗图像按照8:1:1的比例划分为训练集、验证集和测试集,并将它们分别保存在对应的文件夹中。
阅读全文