如何拆分预取的图像数据集(从datasetbuilder下载)

import os
import shutil
from sklearn.model_selection import train_test_split
import numpy as np

# 假设你的数据集路径如下
dataset_path = 'path_to_your_dataset'
images_dir = os.path.join(dataset_path, 'images')
labels_file = os.path.join(dataset_path, 'labels.csv')

# 读取标签文件（假设是CSV格式）
labels = np.genfromtxt(labels_file, delimiter=',', dtype=str)

# 分离特征和标签
image_files = labels[:, 0]  # 假设第一列是图像文件名
image_labels = labels[:, 1]  # 假设第二列是标签

# 首先将数据集分为训练+验证 和 测试集
train_val_images, test_images, train_val_labels, test_labels = train_test_split(
    image_files, image_labels, test_size=0.15, random_state=42)

# 再将训练+验证集分为训练集和验证集
train_images, val_images, train_labels, val_labels = train_test_split(
    train_val_images, train_val_labels, test_size=0.1765, random_state=42)  # 0.15 / 0.85 ≈ 0.1765

# 创建保存路径
os.makedirs(os.path.join(dataset_path, 'train'), exist_ok=True)
os.makedirs(os.path.join(dataset_path, 'val'), exist_ok=True)
os.makedirs(os.path.join(dataset_path, 'test'), exist_ok=True)

# 移动文件到对应的文件夹
for img, label in zip(train_images, train_labels):
    shutil.move(os.path.join(images_dir, img), os.path.join(dataset_path, 'train', img))

for img, label in zip(val_images, val_labels):
    shutil.move(os.path.join(images_dir, img), os.path.join(dataset_path, 'val', img))

for img, label in zip(test_images, test_labels):
    shutil.move(os.path.join(images_dir, img), os.path.join(dataset_path, 'test', img))