import os
import shutil
from sklearn.model_selection import train_test_split
# 定义原始数据集路径
original_images_path = "D:\\PCB_defect\\datasets\\Joint_PCB_New\\images"
original_labels_path = "D:\\PCB_defect\\datasets\\Joint_PCB_New\\labels"
# 定义新数据集路径
new_dataset_path = "D:\\PCB_defect\\datasets\\Joint_PCB_A"
train_path = os.path.join(new_dataset_path, "images", "train")
test_path = os.path.join(new_dataset_path, "images", "test")
val_path = os.path.join(new_dataset_path, "images", "val")
train_labels_path = os.path.join(new_dataset_path, "labels", "train")
test_labels_path = os.path.join(new_dataset_path, "labels", "test")
val_labels_path = os.path.join(new_dataset_path, "labels", "val")
# 创建新的文件夹结构
os.makedirs(train_path, exist_ok=True)
os.makedirs(test_path, exist_ok=True)
os.makedirs(val_path, exist_ok=True)
os.makedirs(train_labels_path, exist_ok=True)
os.makedirs(test_labels_path, exist_ok=True)
os.makedirs(val_labels_path, exist_ok=True)
# 获取所有图像文件名
image_files = [f for f in os.listdir(original_images_path) if f.endswith(".bmp")]
# 划分数据集
train_files, test_files = train_test_split(image_files, test_size=0.2, random_state=42)
test_files, val_files = train_test_split(test_files, test_size=0.5, random_state=42)
# 复制图像和标签文件到新的文件夹
for file in train_files:
shutil.copy(os.path.join(original_images_path, file), os.path.join(train_path, file))
shutil.copy(os.path.join(original_labels_path, file.replace(".bmp", ".txt")), os.path.join(train_labels_path, file.replace(".bmp", ".txt")))
for file in test_files:
shutil.copy(os.path.join(original_images_path, file), os.path.join(test_path, file))
shutil.copy(os.path.join(original_labels_path, file.replace(".bmp", ".txt")), os.path.join(test_labels_path, file.replace(".bmp", ".txt")))
for file in val_files:
shutil.copy(os.path.join(original_images_path, file), os.path.join(val_path, file))
shutil.copy(os.path.join(original_labels_path, file.replace(".bmp", ".txt")), os.path.join(val_labels_path, file.replace(".bmp", ".txt")))
"数据集已成功划分为训练集、测试集和验证集。"
将数据集划分为训练集、测试集和验证集
最新推荐文章于 2024-09-09 22:19:37 发布