109 lines
3.8 KiB
Python
109 lines
3.8 KiB
Python
|
import os
|
|||
|
|
import shutil
|
|||
|
|
import random
|
|||
|
|
from pathlib import Path
|
|||
|
|
|
|||
|
|
|
|||
|
|
def split_train_to_val(train_dir, val_dir, ratio=0.1, seed=42):
|
|||
|
|
"""
|
|||
|
|
从 train_dir 随机抽取 ratio 比例的 **带标签图像** 到 val_dir。
|
|||
|
|
自动判断是分类结构（有子文件夹）还是平铺结构（无子文件夹）。
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
train_dir (str): 训练集路径
|
|||
|
|
val_dir (str): 验证集路径（会自动创建）
|
|||
|
|
ratio (float): 抽取比例，如 0.1 表示 10%
|
|||
|
|
seed (int): 随机种子，保证可复现
|
|||
|
|
"""
|
|||
|
|
train_path = Path(train_dir)
|
|||
|
|
val_path = Path(val_dir)
|
|||
|
|
|
|||
|
|
if not train_path.exists():
|
|||
|
|
raise FileNotFoundError(f"训练目录不存在: {train_path}")
|
|||
|
|
|
|||
|
|
# 设置随机种子
|
|||
|
|
random.seed(seed)
|
|||
|
|
|
|||
|
|
# 获取所有一级子项
|
|||
|
|
items = [p for p in train_path.iterdir()]
|
|||
|
|
|
|||
|
|
# 判断是否为分类结构：所有子项都是目录且非空
|
|||
|
|
is_classification = all(p.is_dir() for p in items) and len(items) > 0
|
|||
|
|
|
|||
|
|
# 定义图像扩展名
|
|||
|
|
IMG_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.gif'}
|
|||
|
|
|
|||
|
|
if is_classification:
|
|||
|
|
print("📁 检测到分类结构（含类别子文件夹）")
|
|||
|
|
for class_dir in items:
|
|||
|
|
class_name = class_dir.name
|
|||
|
|
src_class_dir = train_path / class_name
|
|||
|
|
dst_class_dir = val_path / class_name
|
|||
|
|
dst_class_dir.mkdir(parents=True, exist_ok=True)
|
|||
|
|
|
|||
|
|
# 只找图像文件
|
|||
|
|
image_files = [
|
|||
|
|
f for f in src_class_dir.iterdir()
|
|||
|
|
if f.is_file() and f.suffix.lower() in IMG_EXTENSIONS
|
|||
|
|
]
|
|||
|
|
if not image_files:
|
|||
|
|
print(f" ⚠️ 类别 '{class_name}' 中无图像文件，跳过")
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
num_val = max(1, int(len(image_files) * ratio))
|
|||
|
|
val_images = random.sample(image_files, num_val)
|
|||
|
|
|
|||
|
|
for img in val_images:
|
|||
|
|
# 移动图像
|
|||
|
|
shutil.move(str(img), str(dst_class_dir / img.name))
|
|||
|
|
# 移动同名 .txt 标签
|
|||
|
|
txt_file = img.with_suffix('.txt')
|
|||
|
|
if txt_file.exists():
|
|||
|
|
shutil.move(str(txt_file), str(dst_class_dir / txt_file.name))
|
|||
|
|
|
|||
|
|
print(f" ✅ 类别 '{class_name}': {len(val_images)} 张图像已移至 val")
|
|||
|
|
|
|||
|
|
else:
|
|||
|
|
print("📄 检测到平铺结构（YOLO格式：图像 + 同名 .txt 标签）")
|
|||
|
|
val_path.mkdir(parents=True, exist_ok=True)
|
|||
|
|
|
|||
|
|
# 只收集图像文件（作为采样单元）
|
|||
|
|
image_files = [
|
|||
|
|
f for f in train_path.iterdir()
|
|||
|
|
if f.is_file() and f.suffix.lower() in IMG_EXTENSIONS
|
|||
|
|
]
|
|||
|
|
|
|||
|
|
if not image_files:
|
|||
|
|
print("⚠️ 训练目录中未找到任何图像文件（支持格式: jpg, png 等）")
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
# 随机抽取图像
|
|||
|
|
num_val = max(1, int(len(image_files) * ratio))
|
|||
|
|
val_images = random.sample(image_files, num_val)
|
|||
|
|
|
|||
|
|
# 移动选中的图像及其标签
|
|||
|
|
for img in val_images:
|
|||
|
|
# 移动图像
|
|||
|
|
shutil.move(str(img), str(val_path / img.name))
|
|||
|
|
# 移动同名 .txt
|
|||
|
|
txt_file = img.with_suffix('.txt')
|
|||
|
|
if txt_file.exists():
|
|||
|
|
shutil.move(str(txt_file), str(val_path / txt_file.name))
|
|||
|
|
|
|||
|
|
print(f"✅ 平铺结构: 已移动 {len(val_images)} 张图像及其标签到 {val_path}")
|
|||
|
|
|
|||
|
|
print(f"\n🎉 分割完成！验证集已保存至: {val_path}")
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ======================
|
|||
|
|
# 使用示例
|
|||
|
|
# ======================
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
TRAIN_DIR = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/point2/train"
|
|||
|
|
VAL_DIR = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/point2/val1"
|
|||
|
|
split_train_to_val(
|
|||
|
|
train_dir=TRAIN_DIR,
|
|||
|
|
val_dir=VAL_DIR,
|
|||
|
|
ratio=0.1,
|
|||
|
|
seed=42
|
|||
|
|
)
|