import os
import xml.etree.ElementTree as ET
from PIL import Image

voc_dir = "insulator"
annotations_dir = os.path.join(voc_dir, "Annotations")
images_dir = os.path.join(voc_dir, "images")
labels_dir = os.path.join(voc_dir, "labels")
os.makedirs(labels_dir, exist_ok=True)

def convert(size, box):
    dw, dh = 1.0 / size[0], 1.0 / size[1]
    x = (box[0] + box[1]) / 2.0
    y = (box[2] + box[3]) / 2.0
    w = box[1] - box[0]
    h = box[3] - box[2]
    return x * dw, y * dh, w * dw, h * dh

for xml_file in os.listdir(annotations_dir):
    if not xml_file.endswith(".xml"):
        continue

    image_id = os.path.splitext(xml_file)[0]
    img_path = os.path.join(images_dir, image_id + ".jpg")
    label_path = os.path.join(labels_dir, image_id + ".txt")

    if not os.path.exists(img_path):
        continue

    with Image.open(img_path) as img:
        w, h = img.size

    tree = ET.parse(os.path.join(annotations_dir, xml_file))
    root = tree.getroot()

    with open(label_path, "w") as f:
        for obj in root.findall("object"):
            cls = obj.find("name").text.strip()
            if cls != "insulator":
                continue  # 只保留 insulator 类别
            xmlbox = obj.find("bndbox")
            b = (
                float(xmlbox.find("xmin").text),
                float(xmlbox.find("xmax").text),
                float(xmlbox.find("ymin").text),
                float(xmlbox.find("ymax").text),
            )
            bb = convert((w, h), b)
            f.write(f"1 {' '.join([str(round(v, 6)) for v in bb])}\n")