利用sklearn工具中的train_test_split快速划分实例分割数据集

目录

1.导入相关的包

2. 文件夹路径

3.获取所有文件的路径列表

4. 划分数据集为训练集和测试集最终保存测试集

5.将进一步上面得到的训练集划分为训练集和验证集，保存训练集和验证集

6.将以上的结果保存。

---

1.导入相关的包

import os
from sklearn.model_selection import train_test_split

2. 文件夹路径

# 文件夹路径
original_images = "......\JPEGImages"
annotated_images = "......\SegmentationClass"

3.获取所有文件的路径列表

original_files = [os.path.join(original_images, file) for file in os.listdir(original_images)]
annotated_files = [os.path.join(annotated_images, file) for file in os.listdir(annotated_images)]

4. 划分数据集为训练集和测试集最终保存测试集

train_original, test_original, train_annotated, test_annotated = train_test_split(
    original_files, annotated_files, test_size=0.2, random_state=42)

5.将进一步上面得到的训练集划分为训练集和验证集，保存训练集和验证集

train_original, val_original, train_annotated, val_annotated = train_test_split(
    train_original, train_annotated, test_size=0.1, random_state=42)

6.将以上的结果保存。