问题背景

有了训练集和测试集,但是没有验证集,只能从训练集中划分 20% 作为验证集,剩下的 80% 作为训练集。

在这里插入图片描述

在这里插入图片描述

我的划分前的训练集中包含 png 格式图像和 json 格式的标签,需要新建训练集和验证集文件夹,并分别将图像和标签从源文件夹移动到这两个新建的目标文件夹。

代码实现

安装依赖:

pip install shutil
pip install scikit-learn

这里提供满足我的需求的代码,供参考。

import os
import shutil
from sklearn.model_selection import train_test_split

# 创建文件夹
def mkdir(path):
    folder = os.path.exists(path)
    if not folder:
        os.makedirs(path)
        print(f'-- new folder "{path}" --')
    else:
        print(f'-- the folder "{path}" is already here --')

dataset_path = "/home/xxx/HardDisk/Datasets/weld/weldSeam615/train_val"
train_set_save_path = "/home/xxx/HardDisk/Datasets/weld/weldRGB615/train"
val_set_save_path = "/home/xxx/HardDisk/Datasets/weld/weldRGB615/val"
mkdir(train_set_save_path)
mkdir(val_set_save_path)

file_pathes = os.listdir(dataset_path)
# 获取文件夹下所有 png 格式的图像的名称(不包含后缀名)
img_names = []
for file_path in file_pathes:
    if os.path.splitext(file_path)[1] == ".png":
        file_name = os.path.splitext(file_path)[0]
        img_names.append(file_name)

# 划分训练集和验证集
train_set, val_set = train_test_split(img_names, test_size=0.2, random_state=42)
print(f"train_set size: {len(train_set)}, val_set size: {len(val_set)}")

# 训练集处理:将图像和标签文件移动到目标文件夹
for file_name in train_set:
    img_src_path = os.path.join(dataset_path, file_name+".png")
    img_dst_path = os.path.join(train_set_save_path, file_name+".png")
    shutil.copyfile(img_src_path, img_dst_path)

    json_src_path = os.path.join(dataset_path, file_name+".json")
    json_dst_path = os.path.join(train_set_save_path, file_name+".json")
    shutil.copyfile(json_src_path, json_dst_path)

# 验证集处理:将图像和标签文件移动到目标文件夹
for file_name in val_set:
    img_src_path = os.path.join(dataset_path, file_name+".png")
    img_dst_path = os.path.join(val_set_save_path, file_name+".png")
    shutil.copyfile(img_src_path, img_dst_path)

    json_src_path = os.path.join(dataset_path, file_name+".json")
    json_dst_path = os.path.join(val_set_save_path, file_name+".json")
    shutil.copyfile(json_src_path, json_dst_path)

备注

图像来源网络,侵删。

Logo

魔乐社区(Modelers.cn) 是一个中立、公益的人工智能社区,提供人工智能工具、模型、数据的托管、展示与应用协同服务,为人工智能开发及爱好者搭建开放的学习交流平台。社区通过理事会方式运作,由全产业链共同建设、共同运营、共同享有,推动国产AI生态繁荣发展。

更多推荐