split_validation.py

"""从测试集划分验证集"""
import os
import shutil
from pathlib import Path

# 基础路径
base_path = Path("test_data/bottle")

# 创建验证集目录结构
val_dirs = [
    "val/good",
    "val/broken_large", 
    "val/broken_small",
    "val/contamination"
]

for dir_path in val_dirs:
    (base_path / dir_path).mkdir(parents=True, exist_ok=True)

# 定义每个类别要移动的图片数量
splits = {
    "test/good": 5,           # 5张正常
    "test/broken_large": 5,   # 5张大破损
    "test/broken_small": 5,   # 5张小破损
    "test/contamination": 5   # 5张污染
}

# 移动文件
moved_count = 0
for src_dir, count in splits.items():
    src_path = base_path / src_dir
    val_dir = src_dir.replace("test/", "val/")
    val_path = base_path / val_dir
    
    # 获取前N张图片
    images = sorted(src_path.glob("*.png"))[:count]
    
    for img in images:
        dest = val_path / img.name
        shutil.move(str(img), str(dest))
        moved_count += 1
        print(f"Moved: {src_dir}/{img.name} -> {val_dir}/{img.name}")

print(f"\n✓ 共移动 {moved_count} 张图片到验证集")
print("\n数据集分布：")
train_count = len(list((base_path / "train/good").glob("*.png")))
val_count = len(list((base_path / "val").rglob("*.png")))
test_count = len(list((base_path / "test").rglob("*.png")))

print(f"  训练集: {train_count} 张正常图片")
print(f"  验证集: {val_count} 张")
print(f"  测试集: {test_count} 张")