用于划分图像分类训练集，验证集的脚本

最新推荐文章于 2024-08-24 16:15:50 发布

AIOT魔法师

最新推荐文章于 2024-08-24 16:15:50 发布

阅读量161

点赞数

文章标签：分类数据挖掘人工智能

本文链接：https://blog.csdn.net/weixin_41967600/article/details/132058895

版权

本文介绍了一个使用Python编写的脚本，它从指定的文件夹中提取所有.jpg文件，随机打乱顺序后按8:2的比例划分为训练集和测试集，并将这些图片分别存储在`dataset`文件夹下的`train`和`test`子文件夹中。

摘要由CSDN通过智能技术生成

在这里插入图片描述

import os
import shutil
from random import shuffle
from math import floor

# 创建 dataset 文件夹
dataset_path = 'classify/dataset'
os.makedirs(dataset_path, exist_ok=True)

# 这里填你的图片所在文件夹
main_directory = 'H:\Processor\output_del1'

for dirpath, dirnames, filenames in os.walk(main_directory):
    for dirname in dirnames:
        # 获取每个子文件夹中的所有 jpg 文件
        folder_path = os.path.join(dirpath, dirname)
        images = [img for img in os.listdir(folder_path) if img.endswith(".jpg")]
        shuffle(images)  # 随机打乱图片顺序

        # 按8:2的比例分配训练和测试图片
        split_index = floor(len(images) * 0.8)
        train_images = images[:split_index]
        test_images = images[split_index:]

        # 在dataset文件夹下创建对应的 train 和 test 子文件夹
        os.makedirs(os.path.join(dataset_path, 'train', dirname), exist_ok=True)
        os.makedirs(os.path.join(dataset_path, 'test', dirname), exist_ok=True)

        # 复制训练和测试图片到相应的文件夹
        for img in train_images:
            shutil.copy2(os.path.join(folder_path, img), os.path.join(dataset_path, 'train', dirname))
        for img in test_images:
            shutil.copy2(os.path.join(folder_path, img), os.path.join(dataset_path, 'test', dirname))```

划分后：
![在这里插入图片描述](https://img-blog.csdnimg.cn/8e3850221d25486186756a52a1e6cc0c.png)