import json
import os
import glob

def labelme_to_yolo_segmentation_batch(json_dir, output_dir, class_mapping, img_shape):
    """
    批量将 LabelMe JSON 文件转换为 YOLO 分割格式的 .txt 文件
    :param json_dir: 包含 JSON 文件的目录
    :param output_dir: 输出 .txt 文件的目录
    :param class_mapping: 类别映射字典，如 {"夹具": 0, "夹具1": 1}
    :param img_shape: 图像尺寸 (height, width)
    """
    # 确保输出目录存在
    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)]

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

    img_h, img_w = img_shape
    converted_count = 0

    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")

            with open(output_path, 'w', encoding='utf-8') as out_f:
                for shape in data['shapes']:
                    label = shape['label']
                    points = shape['points']

                    if label not in class_mapping:
                        print(f"⚠️ 跳过未知标签 '{label}' in {json_file}")
                        continue

                    class_id = class_mapping[label]
                    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}")

                    line = f"{class_id} {' '.join(normalized)}"
                    out_f.write(line + '\n')

            print(f"✅ 已转换: {os.path.basename(json_file)} -> {os.path.basename(output_path)}")
            converted_count += 1

        except Exception as e:
            print(f"❌ 转换失败 {json_file}: {e}")

    print(f"\n🎉 批量转换完成！共处理 {converted_count} 个文件。")
    print(f"📁 输出目录: {output_dir}")

# ================== 用户配置区 ==================
JSON_DIR = "/home/hx/yolo/ultralytics_yolo11-main/dataset"           # 输入：存放 LabelMe JSON 的文件夹
OUTPUT_DIR = "labels"                   # 输出：存放 YOLO .txt 的文件夹
CLASS_MAPPING = {
    "夹具": 0,
    "夹具1": 1
    # 可继续添加其他类别
}
IMG_SHAPE = (1440, 2506)                # 替换为你的图像实际尺寸 (height, width)

# ================== 执行转换 ==================
if __name__ == "__main__":
    labelme_to_yolo_segmentation_batch(
        json_dir=JSON_DIR,
        output_dir=OUTPUT_DIR,
        class_mapping=CLASS_MAPPING,
        img_shape=IMG_SHAPE
    )