数据集划分 train val

m0_50888396

于 2022-10-19 13:23:32 发布

阅读量371

点赞数

分类专栏： python 文章标签： python 机器学习 sklearn

本文链接：https://blog.youkuaiyun.com/m0_50888396/article/details/127405978

版权

python 专栏收录该内容

10 篇文章

订阅专栏

本文介绍了如何使用Python进行JPEG图像文件的预处理，通过train_test_split函数实现数据集的6:2:2划分，包括train、val和test三部分，并将文件名写入相应文本文件。主要关注了文件操作和图像数据集划分的实用技巧。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

# coding: utf-8
from pathlib import Path
from sklearn.model_selection import train_test_split


def write_data(file_path: str, data: list):
    with open(file_path, 'w') as f:
        for i in data:
            f.write(i+'\n')


root = Path('/JPEGImages')
images = []

for image in root.glob('*.jpg'):
    images.append(image.stem)

# train:val:test = 6:2:2
image_trainval, image_test = train_test_split(images, test_size=0.2)
image_train, image_val = train_test_split(image_trainval, test_size=0.25)

write_data('/ImageSets/Main/all.txt', images)
write_data('/ImageSets/Main/trainval.txt', image_trainval)
write_data('/ImageSets/Main/train.txt', image_train)
write_data('/ImageSets/Main/val.txt', image_val)
write_data('/ImageSets/Main/test.txt', image_test)