84 lines
2.8 KiB
Python
84 lines
2.8 KiB
Python
import os
|
||
import shutil
|
||
import random
|
||
|
||
def split_train_val_test(train_dir, val_dir, test_dir, val_ratio=0.1, test_ratio=0.1):
|
||
"""
|
||
将 train 目录中的数据按比例划分为 train/val/test 三部分。
|
||
val 和 test 各自获得不同的一部分数据。
|
||
"""
|
||
# 检查目录
|
||
if not os.path.exists(train_dir):
|
||
raise FileNotFoundError(f"训练目录不存在: {train_dir}")
|
||
|
||
os.makedirs(val_dir, exist_ok=True)
|
||
os.makedirs(test_dir, exist_ok=True)
|
||
|
||
# 支持的图片格式
|
||
img_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp'}
|
||
|
||
# 获取所有图片文件
|
||
all_files = os.listdir(train_dir)
|
||
image_files = [f for f in all_files if os.path.splitext(f.lower())[1] in img_extensions]
|
||
|
||
# 配对图片和标签
|
||
pairs = []
|
||
for img in image_files:
|
||
name_no_ext = os.path.splitext(img)[0]
|
||
txt = name_no_ext + '.txt'
|
||
if txt in all_files:
|
||
pairs.append((img, txt))
|
||
else:
|
||
print(f"⚠️ 忽略 {img}：缺少对应标签文件")
|
||
|
||
if len(pairs) == 0:
|
||
print("❌ 未找到有效数据对")
|
||
return
|
||
|
||
total = len(pairs)
|
||
num_test = int(total * test_ratio)
|
||
num_val = int(total * val_ratio)
|
||
|
||
print(f"✅ 共找到 {total} 组有效数据")
|
||
print(f"✅ 将移动：val={num_val}, test={num_test}")
|
||
|
||
# 打乱并抽样
|
||
random.shuffle(pairs)
|
||
test_sample = pairs[:num_test]
|
||
val_sample = pairs[num_test:num_test + num_val]
|
||
# 剩下的留在 train
|
||
|
||
# 移动到 test
|
||
for img, txt in test_sample:
|
||
try:
|
||
shutil.move(os.path.join(train_dir, img), os.path.join(test_dir, img))
|
||
shutil.move(os.path.join(train_dir, txt), os.path.join(test_dir, txt))
|
||
except Exception as e:
|
||
print(f"❌ 移动到 test 失败: {img}, {e}")
|
||
|
||
# 移动到 val
|
||
for img, txt in val_sample:
|
||
try:
|
||
shutil.move(os.path.join(train_dir, img), os.path.join(val_dir, img))
|
||
shutil.move(os.path.join(train_dir, txt), os.path.join(val_dir, txt))
|
||
except Exception as e:
|
||
print(f"❌ 移动到 val 失败: {img}, {e}")
|
||
|
||
print(f"\n✅ 划分完成！")
|
||
print(f" train 保留: {total - num_val - num_test}")
|
||
print(f" val : {num_val}")
|
||
print(f" test : {num_test}")
|
||
print(f" 所有 val 和 test 数据已从 train 移出。")
|
||
|
||
|
||
# ========================
|
||
# 使用示例
|
||
# ========================
|
||
if __name__ == "__main__":
|
||
train_dir = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/obb3/train"
|
||
val_dir = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/obb3/val"
|
||
test_dir = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/obb3/test"
|
||
|
||
split_train_val_test(train_dir, val_dir, test_dir, val_ratio=0.1, test_ratio=0.1)
|
||
# ========================
|