读取文件夹中的图片形式的数据集,随机划分为训练集和测试集比例为9:1,将训练集中图片的名字放入一个阵列中,然后按新的顺序命名后再写入训练集的文件夹,
时间: 2023-12-03 16:44:44 浏览: 106
给定比例随机划分训练集、验证集和测试集
可以使用Python中的os和random库来实现这个功能。具体步骤如下:
1. 导入os和random库
```python
import os
import random
```
2. 设置数据集路径和训练集、测试集比例
```python
data_dir = 'path/to/dataset'
train_ratio = 0.9
```
3. 读取数据集中的图片文件名并打乱顺序
```python
file_names = os.listdir(data_dir)
random.shuffle(file_names)
```
4. 计算训练集和测试集的数量
```python
num_train = int(len(file_names) * train_ratio)
num_test = len(file_names) - num_train
```
5. 将训练集图片名字放入一个列表中
```python
train_file_names = file_names[:num_train]
```
6. 按新顺序命名并将训练集图片复制到新的文件夹中
```python
train_dir = 'path/to/train_dir'
for i, file_name in enumerate(train_file_names):
src_path = os.path.join(data_dir, file_name)
dst_name = f"{i:06d}.jpg" # 新的文件名,例如000001.jpg
dst_path = os.path.join(train_dir, dst_name)
shutil.copy(src_path, dst_path)
```
7. 将测试集图片复制到新的文件夹中
```python
test_dir = 'path/to/test_dir'
for i, file_name in enumerate(file_names[num_train:]):
src_path = os.path.join(data_dir, file_name)
dst_name = f"{i:06d}.jpg" # 新的文件名,例如000001.jpg
dst_path = os.path.join(test_dir, dst_name)
shutil.copy(src_path, dst_path)
```
完整代码如下:
```python
import os
import random
import shutil
data_dir = 'path/to/dataset'
train_dir = 'path/to/train_dir'
test_dir = 'path/to/test_dir'
train_ratio = 0.9
file_names = os.listdir(data_dir)
random.shuffle(file_names)
num_train = int(len(file_names) * train_ratio)
num_test = len(file_names) - num_train
train_file_names = file_names[:num_train]
test_file_names = file_names[num_train:]
if not os.path.exists(train_dir):
os.makedirs(train_dir)
if not os.path.exists(test_dir):
os.makedirs(test_dir)
for i, file_name in enumerate(train_file_names):
src_path = os.path.join(data_dir, file_name)
dst_name = f"{i:06d}.jpg"
dst_path = os.path.join(train_dir, dst_name)
shutil.copy(src_path, dst_path)
for i, file_name in enumerate(test_file_names):
src_path = os.path.join(data_dir, file_name)
dst_name = f"{i:06d}.jpg"
dst_path = os.path.join(test_dir, dst_name)
shutil.copy(src_path, dst_path)
```
阅读全文