import json
import os
import glob

def labelme_to_yolo_segmentation_batch(json_dir, output_dir, target_label="夹具1", class_id=1, img_shape=None):
    """
    批量将 LabelMe JSON 文件转换为 YOLO 分割格式的 .txt 文件
    仅转换指定标签（如 "夹具1"），忽略其他所有标签，且不生成空文件。

    :param json_dir: 包含 LabelMe JSON 文件的目录
    :param output_dir: 输出 .txt 文件的目录
    :param target_label: 要转换的目标标签名称，如 "夹具1"
    :param class_id: 对应的 YOLO 类别 ID（通常从 0 开始）
    :param img_shape: 图像尺寸 (height, width)，如 (1440, 2506)
    """
    if img_shape is None:
        raise ValueError("必须提供 img_shape 参数，例如 (1440, 2506)")

    # 确保输出目录存在
    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  # 标记是否有目标标签

            # 打开输出文件准备写入
            with open(output_path, 'w', encoding='utf-8') as out_f:
                for shape in data.get('shapes', []):
                    label = shape['label']
                    points = shape['points']

                    # 只处理目标标签
                    if label != target_label:
                        continue

                    # 归一化坐标
                    normalized = []
                    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))
                        normalized.append(f"{nx:.6f}")
                        normalized.append(f"{ny:.6f}")

                    # 写入 YOLO 行
                    line = f"{class_id} {' '.join(normalized)}"
                    out_f.write(line + '\n')
                    has_valid_shapes = True

            # 如果没有找到目标标签，删除空文件
            if not has_valid_shapes:
                os.remove(output_path)
                print(f"🟡 跳过: {os.path.basename(json_file)} -> 未包含 '{target_label}'，不生成 .txt")
                skipped_count += 1
            else:
                print(f"✅ 已转换: {os.path.basename(json_file)} -> {os.path.basename(output_path)}")
                converted_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}")
    print("="*50)


# ================== 用户配置区 ==================
JSON_DIR = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/seg/class3"  # LabelMe JSON 文件夹路径
OUTPUT_DIR = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/seg/label"                         # 输出 YOLO 标注文件夹
TARGET_LABEL = "液面"                         # 只转换这个标签
#TARGET_LABEL = "夹具1"                         # 只转换这个标签i
CLASS_ID = 0                                  # YOLO 中该类的 ID
#IMG_SHAPE = (1440, 2560)                      # 图像实际尺寸 (高度, 宽度)
IMG_SHAPE = (640, 640)                      # 图像实际尺寸 (高度, 宽度)
# ================== 执行转换 ==================
if __name__ == "__main__":
    print(f"🚀 开始转换 LabelMe → YOLO (仅 '{TARGET_LABEL}')")
    labelme_to_yolo_segmentation_batch(
        json_dir=JSON_DIR,
        output_dir=OUTPUT_DIR,
        target_label=TARGET_LABEL,
        class_id=CLASS_ID,
        img_shape=IMG_SHAPE
    )