46 lines
2.5 KiB
Python
46 lines
2.5 KiB
Python
import os
|
|
import random
|
|
import shutil
|
|
|
|
def split_dataset(images_dir, labels_dir, train_ratio=0.8):
|
|
# 创建 train 和 val 文件夹
|
|
train_images_dir = os.path.join('/home/admin-root/haotian/python哈汽锻8安全帽识别/安全帽头数据集20250428/2200/train', 'images')
|
|
val_images_dir = os.path.join('/home/admin-root/haotian/python哈汽锻8安全帽识别/安全帽头数据集20250428/2200/val', 'images')
|
|
train_labels_dir = os.path.join('/home/admin-root/haotian/python哈汽锻8安全帽识别/安全帽头数据集20250428/2200/train', 'labels')
|
|
val_labels_dir = os.path.join('/home/admin-root/haotian/python哈汽锻8安全帽识别/安全帽头数据集20250428/2200/val', 'labels')
|
|
|
|
os.makedirs(train_images_dir, exist_ok=True)
|
|
os.makedirs(val_images_dir, exist_ok=True)
|
|
os.makedirs(train_labels_dir, exist_ok=True)
|
|
os.makedirs(val_labels_dir, exist_ok=True)
|
|
|
|
# 获取所有图像文件的列表
|
|
image_files = [f for f in os.listdir(images_dir) if os.path.isfile(os.path.join(images_dir, f))]
|
|
random.shuffle(image_files)
|
|
|
|
# 计算划分点
|
|
split_index = int(len(image_files) * train_ratio)
|
|
|
|
# 划分文件
|
|
train_files = image_files[:split_index]
|
|
val_files = image_files[split_index:]
|
|
|
|
# 移动图像文件到对应的文件夹
|
|
for file in train_files:
|
|
shutil.move(os.path.join(images_dir, file), os.path.join(train_images_dir, file))
|
|
label_file = file.replace('.jpg', '.txt') # 假设标签文件与图像文件同名，但扩展名为 .txt
|
|
shutil.move(os.path.join(labels_dir, label_file), os.path.join(train_labels_dir, label_file))
|
|
|
|
for file in val_files:
|
|
shutil.move(os.path.join(images_dir, file), os.path.join(val_images_dir, file))
|
|
label_file = file.replace('.jpg', '.txt') # 假设标签文件与图像文件同名，但扩展名为 .txt
|
|
shutil.move(os.path.join(labels_dir, label_file), os.path.join(val_labels_dir, label_file))
|
|
|
|
print(f"数据集已成功划分为训练集和验证集，训练集占比 {train_ratio * 100}%")
|
|
|
|
# 使用示例
|
|
images_dir = '/home/admin-root/haotian/python哈汽锻8安全帽识别/安全帽头数据集20250428/2200/images/images' # 替换为你的 images 文件夹路径
|
|
labels_dir = '/home/admin-root/haotian/python哈汽锻8安全帽识别/安全帽头数据集20250428/2200/images/labels' # 替换为你的 labels 文件夹路径
|
|
train_ratio = 0.8 # 训练集占比，默认为 80%
|
|
|
|
split_dataset(images_dir, labels_dir, train_ratio) |