# convert_yolo_seg_to_cvat.py
import os
import cv2
import xml.etree.ElementTree as ET

# ================== 配置 ==================
labels_dir = "./labels/labels"   # YOLO segmentation txt 目录
images_dir = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/6111c/1"
output_xml = "2annotations_cvat_seg.xml"

# 类别映射（id -> name）
class_mapping_reverse = {
    0: "yemian",
}

# ================== 工具函数 ==================
def unnormalize_polygon(img_w, img_h, coords):
    """
    YOLO seg 归一化坐标 -> 像素坐标
    coords: [x1, y1, x2, y2, ...]
    """
    pts = []
    for i in range(0, len(coords), 2):
        x = float(coords[i]) * img_w
        y = float(coords[i + 1]) * img_h
        pts.append((x, y))
    return pts


# ================== 构建 XML ==================
root = ET.Element("annotations")
ET.SubElement(root, "version").text = "1.1"

meta = ET.SubElement(root, "meta")
task = ET.SubElement(meta, "task")
ET.SubElement(task, "name").text = "converted_from_yolo_seg"

txt_files = sorted([f for f in os.listdir(labels_dir) if f.endswith(".txt")])
ET.SubElement(task, "size").text = str(len(txt_files))

# labels
labels_elem = ET.SubElement(task, "labels")
for name in class_mapping_reverse.values():
    lab = ET.SubElement(labels_elem, "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 = ""

# ================== 处理每个 txt ==================
for idx, txt_file in enumerate(txt_files):
    base = os.path.splitext(txt_file)[0]

    # 查找图片
    img_path = None
    for ext in [".jpg", ".jpeg", ".png", ".bmp"]:
        p = os.path.join(images_dir, base + ext)
        if os.path.exists(p):
            img_path = p
            break

    if img_path is None:
        print(f"❗缺少图片：{base}.*（跳过）")
        continue

    img = cv2.imread(img_path)
    if img is None:
        print(f"❗无法读取图片：{img_path}")
        continue

    h, w = img.shape[:2]

    image_elem = ET.SubElement(
        root, "image",
        id=str(idx),
        name=os.path.basename(img_path),
        width=str(w),
        height=str(h)
    )

    yolo_path = os.path.join(labels_dir, txt_file)
    with open(yolo_path, "r") as f:
        lines = [l.strip() for l in f.readlines() if l.strip()]

    if not lines:
        print(f"⚠ 空标签：{txt_file}")
        continue

    # 支持一张图多个实例
    for line in lines:
        parts = line.split()
        cls_id = int(parts[0])
        coords = parts[1:]

        if len(coords) < 6 or len(coords) % 2 != 0:
            print(f"⚠ 非法 segmentation 点数：{txt_file}")
            continue

        label_name = class_mapping_reverse.get(cls_id, str(cls_id))

        pts = unnormalize_polygon(w, h, coords)
        pts_str = ";".join([f"{x:.2f},{y:.2f}" for x, y in pts])

        ET.SubElement(
            image_elem,
            "polygon",
            label=label_name,
            source="manual",
            occluded="0",
            points=pts_str,
            z_order="0"
        )

# ================== 写出 XML ==================
tree = ET.ElementTree(root)
tree.write(output_xml, encoding="utf-8", xml_declaration=True)

print("\n🎉 YOLO Segmentation → CVAT 转换完成！")
print(f"📄 输出文件：{output_xml}")
print(f"📊 共处理 {len(txt_files)} 个 txt")