给出一个能够将若干图片划分为训练集和测试集的代码
时间: 2023-04-12 07:02:09 浏览: 82
以下是一个能够将若干图片划分为训练集和测试集的 Python 代码:
```python
import os
import random
import shutil
def split_data(input_dir, output_dir, split_ratio):
if not os.path.exists(output_dir):
os.makedirs(output_dir)
for class_name in os.listdir(input_dir):
class_dir = os.path.join(input_dir, class_name)
if not os.path.isdir(class_dir):
continue
train_dir = os.path.join(output_dir, 'train', class_name)
test_dir = os.path.join(output_dir, 'test', class_name)
os.makedirs(train_dir, exist_ok=True)
os.makedirs(test_dir, exist_ok=True)
file_list = os.listdir(class_dir)
random.shuffle(file_list)
split_idx = int(len(file_list) * split_ratio)
train_files = file_list[:split_idx]
test_files = file_list[split_idx:]
for file_name in train_files:
src_path = os.path.join(class_dir, file_name)
dst_path = os.path.join(train_dir, file_name)
shutil.copy(src_path, dst_path)
for file_name in test_files:
src_path = os.path.join(class_dir, file_name)
dst_path = os.path.join(test_dir, file_name)
shutil.copy(src_path, dst_path)
if __name__ == '__main__':
input_dir = '/path/to/input/dir'
output_dir = '/path/to/output/dir'
split_ratio = 0.8
split_data(input_dir, output_dir, split_ratio)
```
这个代码会将 `input_dir` 目录下的图片按照 `split_ratio` 的比例划分为训练集和测试集,并将它们保存到 `output_dir` 目录下的 `train` 和 `test` 子目录中。每个子目录下的图片会按照它们原本所在的类别分别保存在以类别命名的子目录中。
阅读全文