# yolo_obb_to_cvat_full.py
import os
import math
import xml.etree.ElementTree as ET
from pathlib import Path
import cv2


def recover_rotated_box(points):
    """
    YOLO OBB 四点 → 中心点、宽高、旋转角度
    点顺序为：p1 p2 p3 p4，逆时针/顺时针均可接受
    """

    cx = sum(p[0] for p in points) / 4
    cy = sum(p[1] for p in points) / 4

    # 点1→点2 作为宽方向
    vx = points[1][0] - points[0][0]
    vy = points[1][1] - points[0][1]

    w = math.sqrt(vx*vx + vy*vy)

    # 点2→点3 作为高方向
    vx2 = points[2][0] - points[1][0]
    vy2 = points[2][1] - points[1][1]

    h = math.sqrt(vx2*vx2 + vy2*vy2)

    # 角度 = 边1 与水平线夹角
    angle_rad = math.atan2(vy, vx)
    angle_deg = math.degrees(angle_rad)

    # CVAT 角度必须是 0~360
    if angle_deg < 0:
        angle_deg += 360

    return cx, cy, w, h, angle_deg


def yolo_obb_to_cvat_xml(label_dir, image_dir, class_id_to_name, output_xml):

    label_dir = Path(label_dir)
    image_dir = Path(image_dir)

    # ======== 构建基本 XML 结构（与你发的一模一样） ========
    root = ET.Element("annotations")

    ET.SubElement(root, "version").text = "1.1"

    # meta/task
    meta = ET.SubElement(root, "meta")
    task = ET.SubElement(meta, "task")

    ET.SubElement(task, "id").text = "1"
    ET.SubElement(task, "name").text = "yolo_obb_import"
    ET.SubElement(task, "size").text = str(len(list(label_dir.glob("*.txt"))))
    ET.SubElement(task, "mode").text = "annotation"
    ET.SubElement(task, "overlap").text = "0"
    ET.SubElement(task, "bugtracker").text = ""
    ET.SubElement(task, "created").text = ""
    ET.SubElement(task, "updated").text = ""
    ET.SubElement(task, "subset").text = "default"
    ET.SubElement(task, "start_frame").text = "0"
    ET.SubElement(task, "stop_frame").text = str(len(list(label_dir.glob("*.txt"))) - 1)
    ET.SubElement(task, "frame_filter").text = ""

    # labels
    labels = ET.SubElement(task, "labels")
    for name in class_id_to_name.values():
        lab = ET.SubElement(labels, "label")
        ET.SubElement(lab, "name").text = name
        ET.SubElement(lab, "color").text = "#ffffff"
        ET.SubElement(lab, "type").text = "any"
        ET.SubElement(lab, "attributes")

    ET.SubElement(meta, "dumped").text = ""

    # ======== 处理每张图片 ========
    for idx, txt_file in enumerate(sorted(label_dir.glob("*.txt"))):
        stem = txt_file.stem
        img_file = None

        # 自动匹配图片
        for ext in ["jpg", "png", "jpeg", "bmp", "webp"]:
            p = image_dir / f"{stem}.{ext}"
            if p.exists():
                img_file = p
                break

        if img_file is None:
            print(f"⚠ 找不到图片: {stem}")
            continue

        img = cv2.imread(str(img_file))
        H, W = [1080,1920]

        # image tag
        image_elem = ET.SubElement(root, "image", {
            "id": str(idx),
            "name": img_file.name,
            "width": str(W),
            "height": str(H)
        })

        # 读取 txt
        with open(txt_file, "r") as f:
            for line in f.readlines():
                parts = line.strip().split()
                cls_id = int(parts[0])
                coords = list(map(float, parts[1:]))

                pts = []
                for i in range(0, 8, 2):
                    x = coords[i] * W
                    y = coords[i+1] * H
                    pts.append((x, y))

                # 四点 → 框参数
                cx, cy, bw, bh, ang = recover_rotated_box(pts)

                xtl = cx - bw / 2
                ytl = cy - bh / 2
                xbr = cx + bw / 2
                ybr = cy + bh / 2

                # box 标签（完全跟你的一样）
                ET.SubElement(image_elem, "box", {
                    "label": class_id_to_name[cls_id],
                    "source": "manual",
                    "occluded": "0",
                    "xtl": f"{xtl:.2f}",
                    "ytl": f"{ytl:.2f}",
                    "xbr": f"{xbr:.2f}",
                    "ybr": f"{ybr:.2f}",
                    "rotation": f"{ang:.2f}",
                    "z_order": "0"
                })

        print(f"✔ 处理 {img_file.name}")

    # 保存 XML
    tree = ET.ElementTree(root)
    tree.write(output_xml, encoding="utf-8", xml_declaration=True)
    print(f"已生成 XML：{output_xml}")

if __name__ == "__main__":
    CLASS_MAP = {
        0: "clamp",
    }

    yolo_obb_to_cvat_xml(
        label_dir="/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/1/labels/labels",
        image_dir="/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/1/zjdata17",
        class_id_to_name=CLASS_MAP,
        output_xml="annotations.xml"
    )