import os
import shutil

# ========== 用户配置 ==========
SOURCE_DIR = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/obb/train"                     # 要检查的目录（.表示当前目录）
TARGET_DIR = "selected_files"        # 目标子文件夹
# =================================

os.makedirs(TARGET_DIR, exist_ok=True)

for filename in os.listdir(SOURCE_DIR):
    if filename.endswith(".txt"):
        txt_path = os.path.join(SOURCE_DIR, filename)

        # 读取行数
        with open(txt_path, "r", encoding="utf-8") as f:
            lines = f.readlines()

        # 判断是否为 3 行
        if len(lines) == 3:
            base = os.path.splitext(filename)[0]

            # 找到同名的所有文件
            for f in os.listdir(SOURCE_DIR):
                if f.startswith(base + "."):
                    src_path = os.path.join(SOURCE_DIR, f)
                    dst_path = os.path.join(TARGET_DIR, f)

                    print(f"Moving: {src_path} → {dst_path}")
                    shutil.move(src_path, dst_path)

print("处理完成！")