import os
import json
import shutil
import random
from PIL import Image
import argparse

def convert_single_json_to_yolo(json_path, class_mapping):
    """将单个 {image}.json 转换为 YOLO 格式，并返回 (image_path, yolo_lines)"""
    with open(json_path, "r", encoding="utf-8") as f:
        data = json.load(f)

    # 假设 JSON 结构: { "image": "xxx.jpg", "annotations": [...] }
    if isinstance(data, list):
        item = data[0]  # 如果是列表，取第一个
    else:
        item = data     # 如果是字典

    image_name = item["image"]
    image_path = os.path.join(os.path.dirname(json_path), image_name)

    if not os.path.exists(image_path):
        print(f"⚠️ 图像不存在: {image_path}")
        return None, None

    try:
        with Image.open(image_path) as img:
            img_w, img_h = img.size
    except Exception as e:
        print(f"❌ 无法打开图像 {image_path}: {e}")
        return None, None

    yolo_lines = []
    for ann in item.get("annotations", []):
        label = ann["label"]
        if label not in class_mapping:
            print(f"⚠️ 未知类别 '{label}'，跳过")
            continue

        class_id = class_mapping[label]
        coords = ann["coordinates"]
        x_min = coords["x"]
        y_min = coords["y"]
        width = coords["width"]
        height = coords["height"]

        x_center = x_min + width / 2
        y_center = y_min + height / 2

        x_norm = x_center / img_w
        y_norm = y_center / img_h
        w_norm = width / img_w
        h_norm = height / img_h

        yolo_lines.append(f"{class_id} {x_norm:.6f} {y_norm:.6f} {w_norm:.6f} {h_norm:.6f}")

    return image_path, yolo_lines

def main():
    parser = argparse.ArgumentParser(description="从同目录的 image.jpg + image.json 构建 YOLO 数据集")
    parser.add_argument("--data-dir", required=True, help="包含图像和JSON的原始目录")
    parser.add_argument("--output-dir", default="yolo_dataset", help="输出 YOLO 数据集目录")
    parser.add_argument("--classes", required=True, help='类别映射，如 "person:0"')
    parser.add_argument("--train", type=float, default=0.7, help="训练集比例 (default: 0.7)")
    parser.add_argument("--val", type=float, default=0.2, help="验证集比例 (default: 0.2)")
    parser.add_argument("--test", type=float, default=0.1, help="测试集比例 (default: 0.1)")
    parser.add_argument("--dataset-root", default=None, help="data.yaml 中的 path 字段（默认为 ./，表示当前目录）")

    args = parser.parse_args()

    # 验证比例
    total = args.train + args.val + args.test
    if abs(total - 1.0) > 1e-5:
        raise ValueError(f"比例之和必须为1，当前: train={args.train}, val={args.val}, test={args.test}")

    class_mapping = {}
    for pair in args.classes.split(","):
        name, idx = pair.split(":")
        class_mapping[name.strip()] = int(idx.strip())
    print(f"类别映射: {class_mapping}")

    # 找到所有 .json 文件
    json_files = [f for f in os.listdir(args.data_dir) if f.lower().endswith(".json")]
    print(f"🔍 找到 {len(json_files)} 个 JSON 文件")

    valid_pairs = []
    for json_file in json_files:
        json_path = os.path.join(args.data_dir, json_file)
        image_path, yolo_lines = convert_single_json_to_yolo(json_path, class_mapping)
        if image_path is not None and yolo_lines:
            valid_pairs.append((image_path, yolo_lines, json_path))
        elif image_path is not None:
            print(f"ℹ️  {json_file} 无有效标注，跳过")

    print(f"✅ 共 {len(valid_pairs)} 个有效样本")

    # 随机打乱并划分
    random.seed(42)  # 可复现
    random.shuffle(valid_pairs)

    n = len(valid_pairs)
    n_train = int(n * args.train)
    n_val = int(n * args.val)
    n_test = n - n_train - n_val

    splits = [
        ("train", valid_pairs[:n_train]),
        ("val", valid_pairs[n_train:n_train + n_val]),
        ("test", valid_pairs[n_train + n_val:])
    ]

    # 创建输出目录
    output_images = os.path.join(args.output_dir, "images")
    output_labels = os.path.join(args.output_dir, "labels")
    for split in ["train", "val", "test"]:
        os.makedirs(os.path.join(output_images, split), exist_ok=True)
        os.makedirs(os.path.join(output_labels, split), exist_ok=True)

    # 复制图像 + 生成标签
    for split_name, items in splits:
        print(f"\n📦 正在处理 {split_name} ({len(items)} 张图)...")
        for image_path, yolo_lines, json_path in items:
            # 复制图像
            image_name = os.path.basename(image_path)
            dst_image = os.path.join(output_images, split_name, image_name)
            shutil.copy2(image_path, dst_image)

            # 写入标签
            txt_name = os.path.splitext(image_name)[0] + ".txt"
            dst_txt = os.path.join(output_labels, split_name, txt_name)
            with open(dst_txt, "w") as f:
                f.write("\n".join(yolo_lines))

    # 生成符合 Ultralytics YOLO 格式的 data.yaml
    data_yaml = os.path.join(args.output_dir, "data.yaml")

    # 构建 names: {0: 'person', 1: 'car', ...}
    reversed_map = {v: k for k, v in class_mapping.items()}
    max_id = max(reversed_map.keys()) if reversed_map else -1
    names_lines = []
    for i in range(max_id + 1):
        name = reversed_map.get(i, f"class{i}")
        names_lines.append(f"  {i}: {name}")

    if args.dataset_root:
        yaml_path = args.dataset_root
    else:
        yaml_path = "./"
        
    with open(data_yaml, "w", encoding="utf-8") as f:
        f.write(f"""# Ultralytics YOLO 🚀
path: {yaml_path}
train: images/train
val: images/val
test: images/test

names:
{chr(10).join(names_lines)}
""")

    print(f"\n🎉 完成！YOLO 数据集保存在: {os.path.abspath(args.output_dir)}")
    print(f"📊 训练集: {n_train}, 验证集: {n_val}, 测试集: {n_test}")

if __name__ == "__main__":
    main()

运行命令

python json_to_txt.py \
--data-dir /Users/tiger/Desktop/images \
--output-dir /Users/tiger/Desktop/result \
--classes "person:0" \
--train 0.7 \
--val 0.2 \
--test 0.1 \
--dataset-root parson