import os
import json
import shutil
from tqdm import tqdm

# ======================
# ⚙️ 可配置路径
# ======================
# 原始数据目录（Labelme）
INPUT_DIR = r"\\192.168.50.222\share\zyh\data\坡口切割\10-11"

# 输出 YOLO 数据集目录（标准格式）
OUTPUT_DIR =r"\\192.168.50.222\share\zyh\data\坡口切割\yolo"

# ======================
# 🚀 核心函数
# ======================

def convert_labelme_to_yolo(json_path, output_txt, image_w, image_h):
    """将 Labelme JSON 转为 YOLO txt"""
    with open(json_path, "r", encoding="utf-8") as f:
        data = json.load(f)

    lines = []
    for shape in data.get("shapes", []):
        if shape["label"] != "line":
            continue
        points = shape["points"]
        if len(points) != 2:
            continue

        x1, y1 = points[0]
        x2, y2 = points[1]

        # 归一化
        x1 /= image_w
        y1 /= image_h
        x2 /= image_w
        y2 /= image_h

        lines.append(f"0 {x1:.6f} {y1:.6f} {x2:.6f} {y2:.6f}")

    os.makedirs(os.path.dirname(output_txt), exist_ok=True)
    with open(output_txt, "w", encoding="utf-8") as f:
        f.write("\n".join(lines))


def process_split(split):
    """处理 train / val 子集"""
    print(f"\n📂 处理中：{split} 集")

    img_src_dir = os.path.join(INPUT_DIR, "images", split)
    json_src_dir = os.path.join(INPUT_DIR, "labels", split)

    img_dst_dir = os.path.join(OUTPUT_DIR, split, "images")
    label_dst_dir = os.path.join(OUTPUT_DIR, split, "labels")

    os.makedirs(img_dst_dir, exist_ok=True)
    os.makedirs(label_dst_dir, exist_ok=True)

    json_files = [f for f in os.listdir(json_src_dir) if f.endswith(".json")]

    for json_file in tqdm(json_files):
        json_path = os.path.join(json_src_dir, json_file)

        with open(json_path, "r", encoding="utf-8") as f:
            data = json.load(f)

        image_name = data["imagePath"]
        image_w = data["imageWidth"]
        image_h = data["imageHeight"]

        # 图片路径
        src_img_path = os.path.join(img_src_dir, image_name)
        dst_img_path = os.path.join(img_dst_dir, image_name)

        if not os.path.exists(src_img_path):
            print(f"⚠️ 缺失图片：{src_img_path}")
            continue

        # 拷贝图片
        shutil.copy2(src_img_path, dst_img_path)

        # 生成 YOLO txt
        txt_name = os.path.splitext(image_name)[0] + ".txt"
        txt_path = os.path.join(label_dst_dir, txt_name)

        convert_labelme_to_yolo(json_path, txt_path, image_w, image_h)


def generate_data_yaml():
    """生成 data.yaml 文件"""
    yaml_path = os.path.join(OUTPUT_DIR, "data.yaml")
    yaml_content = f"""path: {OUTPUT_DIR}
train: train/images
val: val/images

names:
  0: line
"""
    with open(yaml_path, "w", encoding="utf-8") as f:
        f.write(yaml_content)
    print(f"\n✅ 已生成 data.yaml: {yaml_path}")


if __name__ == "__main__":
    print("🚀 开始转换 Labelme → YOLO 数据集 (标准格式)")

    process_split("train")
    process_split("val")

    generate_data_yaml()

    print("\n🎉 完成！YOLO 数据集已生成于：", OUTPUT_DIR)