import json
import os
import glob

def labelme_to_yolo_keypoints_batch(json_dir, output_dir, target_label="夹具1", class_id=0, img_shape=None, num_keypoints=None):
    """
    批量将 LabelMe JSON 文件中的关键点标注转换为 YOLO Pose (Keypoints) 格式的 .txt 文件
    仅转换指定标签（如 "夹具1"），忽略其他标签，且不生成空文件。

    :param json_dir: 包含 LabelMe JSON 文件的目录
    :param output_dir: 输出 .txt 文件的目录
    :param target_label: 要转换的目标标签名称
    :param class_id: 对应的 YOLO 类别 ID
    :param img_shape: 图像尺寸 (height, width)，如 (1440, 2506)
    :param num_keypoints: 预期关键点数量（可选，用于校验）
    """
    if img_shape is None:
        raise ValueError("必须提供 img_shape 参数，例如 (1440, 2506)")

    if num_keypoints is not None and num_keypoints <= 0:
        raise ValueError("num_keypoints 必须为正整数")

    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)

    # 获取所有 .json 文件（排除 _mask.json 等）
    json_files = glob.glob(os.path.join(json_dir, "*.json"))
    json_files = [f for f in json_files if os.path.isfile(f) and not f.endswith("_mask.json")]

    if not json_files:
        print(f"❌ 在 {json_dir} 中未找到任何 JSON 文件")
        return

    img_h, img_w = img_shape
    converted_count = 0
    skipped_count = 0

    print(f"🔍 开始转换关键点，仅处理标签: '{target_label}' (class_id={class_id})")

    for json_file in json_files:
        try:
            with open(json_file, 'r', encoding='utf-8') as f:
                data = json.load(f)

            base_name = os.path.splitext(os.path.basename(json_file))[0]
            output_path = os.path.join(output_dir, f"{base_name}.txt")

            has_valid_shapes = False
            keypoints_found = []

            with open(output_path, 'w', encoding='utf-8') as out_f:
                for shape in data.get('shapes', []):
                    label = shape['label']
                    points = shape['points']  # 每个 shape 是一个关键点 [x, y]

                    if label != target_label:
                        continue

                    # 收集所有关键点坐标
                    for x, y in points:
                        nx = max(0.0, min(1.0, x / img_w))
                        ny = max(0.0, min(1.0, y / img_h))
                        keypoints_found.extend([f"{nx:.6f}", f"{ny:.6f}", "2"])  # v=2: visible

                    has_valid_shapes = True

                # 写入 YOLO 行（支持多个 shape？通常一个标签一个实例）
                if has_valid_shapes:
                    if num_keypoints is not None and len(keypoints_found) // 3 != num_keypoints:
                        print(f"⚠️  {os.path.basename(json_file)}: 关键点数量不匹配，期望 {num_keypoints}，实际 {len(keypoints_found)//3}")
                    line = f"{class_id} {' '.join(keypoints_found)}"
                    out_f.write(line + '\n')
                    print(f"✅ 已转换: {os.path.basename(json_file)} -> {os.path.basename(output_path)}")
                    converted_count += 1
                else:
                    # 删除空文件
                    os.remove(output_path)

            if not has_valid_shapes:
                print(f"🟡 跳过: {os.path.basename(json_file)} -> 未包含 '{target_label}'，不生成 .txt")
                skipped_count += 1

        except Exception as e:
            print(f"❌ 转换失败 {json_file}: {e}")
            if os.path.exists(output_path):
                os.remove(output_path)

    print("\n" + "="*50)
    print(f"🎉 批量关键点转换完成！")
    print(f"📊 成功转换: {converted_count}")
    print(f"📊 跳过文件: {skipped_count}")
    print(f"📁 输出目录: {output_dir}")
    if num_keypoints:
        print(f"📍 关键点数量: {num_keypoints}")
    print("="*50)


# ================== 用户配置区 ==================
JSON_DIR = "/home/hx/yolo/yolo11_point/folder_end"   # LabelMe JSON 文件夹路径
OUTPUT_DIR = "labels_keypoints"                    # 输出 YOLO 关键点标签目录
TARGET_LABEL = "point"                              # 要提取的关键点标签
CLASS_ID = 0                                       # YOLO 中该类的 ID
IMG_SHAPE = (1440, 2506)                           # (高度, 宽度)
NUM_KEYPOINTS = 4                                  # 可选：指定关键点数量用于校验

# ================== 执行转换 ==================
if __name__ == "__main__":
    print(f"🚀 开始转换 LabelMe 关键点 → YOLO Pose 格式 (仅 '{TARGET_LABEL}')")
    labelme_to_yolo_keypoints_batch(
        json_dir=JSON_DIR,
        output_dir=OUTPUT_DIR,
        target_label=TARGET_LABEL,
        class_id=CLASS_ID,
        img_shape=IMG_SHAPE,
        num_keypoints=NUM_KEYPOINTS
    )