# convert_to_detection.py
import os
import shutil


def convert_labels(input_dir, output_dir=None):
    """
    将包含关键点的标签文件转换为仅保留目标检测 bbox 的标签文件
    输入格式：class cx cy w h x1 y1 v1 x2 y2 v2 x3 y3 v3 x4 y4 v4
    输出格式：class cx cy w h  （仅保留前5个）

    :param input_dir: 原始标签目录路径，如 'dataset/labels/train'
    :param output_dir: 输出目录路径，如果为 None，则覆盖原文件
    """
    if output_dir is None:
        output_dir = input_dir  # 覆盖原文件
        print(f"⚠️  警告：将覆盖原文件目录 {input_dir}")
    else:
        os.makedirs(output_dir, exist_ok=True)
        print(f"✅ 输出目录已创建: {output_dir}")

    # 确认输入目录存在
    if not os.path.exists(input_dir):
        raise FileNotFoundError(f"输入目录不存在: {input_dir}")

    # 统计处理文件数
    count = 0

    for filename in os.listdir(input_dir):
        if filename.endswith('.txt'):
            input_path = os.path.join(input_dir, filename)
            output_path = os.path.join(output_dir, filename)

            with open(input_path, 'r') as f:
                lines = f.readlines()

            new_lines = []
            for line in lines:
                parts = line.strip().split()
                if not parts:
                    continue
                # 只保留前5个数据：class, cx, cy, w, h
                det_line = ' '.join(parts[:5])
                new_lines.append(det_line)

            # 写入新文件
            with open(output_path, 'w') as f:
                f.write('\n'.join(new_lines))

            count += 1

    print(f"✅ 转换完成！共处理 {count} 个文件。")
    print(f"📁 输出目录: {output_dir}")


# ======================
# 使用示例
# ======================

if __name__ == "__main__":
    # ✅ 修改以下路径为你的实际路径
    #INPUT_LABELS_DIR = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/point2/train"  # 原始标签目录
    #OUTPUT_LABELS_DIR = None # 转换后的目录

    # 如果你想覆盖原文件，可以设置 OUTPUT_LABELS_DIR = None
    # OUTPUT_LABELS_DIR = None

    #convert_labels(INPUT_LABELS_DIR, OUTPUT_LABELS_DIR)

    # 如果你还有 val 目录，也处理一下
    INPUT_VAL_DIR = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/point2/val"
    OUTPUT_VAL_DIR = None

    convert_labels(INPUT_VAL_DIR, OUTPUT_VAL_DIR)

    print("🎉 所有标签已转换完毕！现在可以训练 detect 模型了。")