# convert_cvat_to_yolo_keypoints.py
import xml.etree.ElementTree as ET
import os

# =================== 配置 ===================
xml_file = 'annotations.xml'           # 你的 CVAT XML 文件路径
images_dir = 'images'                  # 图像文件夹（用于读取宽高）
output_dir = 'labels_keypoints'        # 输出 YOLO 标签目录
os.makedirs(output_dir, exist_ok=True)

# 类别映射（根据你的 XML 中的 label name）
class_mapping = {
    'clamp1': 0,
    'clamp0': 1,
    'kongliao': 2,
    'duiliao': 3
}
# ============================================

def parse_points(points_str):
    """解析 CVAT 的 points 字符串，返回 [(x, y), ...]"""
    return [(float(p.split(',')[0]), float(p.split(',')[1])) for p in points_str.split(';')]

def normalize_bbox_and_kpts(image_w, image_h, bbox, keypoints):
    """归一化 bbox 和关键点"""
    cx, cy, w, h = bbox
    cx_n, cy_n, w_n, h_n = cx/image_w, cy/image_h, w/image_w, h/image_h
    kpts_n = []
    for x, y in keypoints:
        kpts_n.append(x / image_w)
        kpts_n.append(y / image_h)
        kpts_n.append(2)  # v=2: visible
    return (cx_n, cy_n, w_n, h_n), kpts_n

def points_to_bbox(points):
    """从点集生成最小外接矩形 (x_center, y_center, width, height)"""
    xs = [p[0] for p in points]
    ys = [p[1] for p in points]
    x_min, x_max = min(xs), max(xs)
    y_min, y_max = min(ys), max(ys)
    cx = (x_min + x_max) / 2
    cy = (y_min + y_max) / 2
    w = x_max - x_min
    h = y_max - y_min
    return cx, cy, w, h

# 解析 XML
tree = ET.parse(xml_file)
root = tree.getroot()

for image_elem in root.findall('image'):
    image_name = image_elem.get('name')
    image_w = int(image_elem.get('width'))
    image_h = int(image_elem.get('height'))
    label_file = os.path.splitext(image_name)[0] + '.txt'
    label_path = os.path.join(output_dir, label_file)

    with open(label_path, 'w') as f_out:
        # 查找 points 标注
        points_elem = image_elem.find("points[@label='clamp1']")
        if points_elem is None:
            print(f"⚠️ 警告：图像 {image_name} 缺少 clamp1 关键点标注")
            continue

        points_str = points_elem.get('points')
        keypoints = parse_points(points_str)

        if len(keypoints) != 4:
            print(f"⚠️ 警告：图像 {image_name} 的关键点数量不是 4，实际为 {len(keypoints)}")
            continue

        # 生成包围框
        bbox = points_to_bbox(keypoints)

        # 归一化
        (cx_n, cy_n, w_n, h_n), kpts_n = normalize_bbox_and_kpts(image_w, image_h, bbox, keypoints)

        # 写入 YOLO 格式
        line = [str(class_mapping['clamp1']), str(cx_n), str(cy_n), str(w_n), str(h_n)] + [str(k) for k in kpts_n]
        f_out.write(' '.join(line) + '\n')

print("✅ 所有标注已成功转换为 YOLO 关键点格式！")
print(f"📌 输出目录: {output_dir}")