"""用训练好的模型对未标注图片做推理，输出labelme JSON预标注"""
import json
import argparse
from pathlib import Path
from ultralytics import YOLO

IMG_W, IMG_H = 3200, 1800


def yolo_box_to_labelme_shape(cx_n, cy_n, w_n, h_n):
    cx = cx_n * IMG_W
    cy = cy_n * IMG_H
    w = w_n * IMG_W
    h = h_n * IMG_H
    x1, y1 = cx - w / 2, cy - h / 2
    x2, y2 = cx + w / 2, cy + h / 2
    return {
        "label": "goat",
        "points": [[round(x1, 4), round(y1, 4)], [round(x2, 4), round(y2, 4)]],
        "group_id": None,
        "description": "",
        "shape_type": "rectangle",
        "flags": {},
        "mask": None,
    }


def make_labelme_json(img_path: Path, shapes: list) -> dict:
    return {
        "version": "5.5.0",
        "flags": {},
        "shapes": shapes,
        "imagePath": img_path.name,
        "imageData": None,
        "imageHeight": IMG_H,
        "imageWidth": IMG_W,
    }


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("model_path", help="best.pt路径")
    parser.add_argument("unlabeled_root", help="Unlabeled_images根目录")
    parser.add_argument("out_root", help="Auto_labeled输出根目录")
    parser.add_argument("--conf", type=float, default=0.25)
    parser.add_argument("--iou", type=float, default=0.45)
    parser.add_argument("--imgsz", type=int, default=1280)
    args = parser.parse_args()

    model = YOLO(args.model_path)
    unlabeled_root = Path(args.unlabeled_root)
    out_root = Path(args.out_root)

    img_paths = sorted(unlabeled_root.rglob("*.jpg"))
    print(f"Found {len(img_paths)} unlabeled images")

    for img_path in img_paths:
        rel = img_path.parent.relative_to(unlabeled_root)
        period = rel.parts[-1]
        cam = rel.parts[0]
        out_dir = out_root / cam / f"{period}Label"
        out_dir.mkdir(parents=True, exist_ok=True)

        results = model.predict(
            str(img_path),
            conf=args.conf,
            iou=args.iou,
            imgsz=args.imgsz,
            verbose=False,
        )

        shapes = []
        for box in results[0].boxes:
            cx_n, cy_n, w_n, h_n = box.xywhn[0].tolist()
            shapes.append(yolo_box_to_labelme_shape(cx_n, cy_n, w_n, h_n))

        out_json = out_dir / (img_path.stem + ".json")
        out_json.write_text(
            json.dumps(make_labelme_json(img_path, shapes), ensure_ascii=False, indent=2),
            encoding="utf-8",
        )

    print(f"Done. Output: {out_root}")


if __name__ == "__main__":
    main()