xml_to_yolo_converter.py

import os
import xml.etree.ElementTree as ET

# Configurações
classes = ['With Helmet', 'Without Helmet']  # nomes exatamente iguais aos do XML
xml_dir = './dataset/labels/xml/'      # Pasta com os arquivos .xml
yolo_dir = './dataset/labels/yolo/'   # Pasta de saída para arquivos .txt

os.makedirs(yolo_dir, exist_ok=True)

for xml_file in os.listdir(xml_dir):
    if not xml_file.endswith('.xml'):
        continue

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

    size = root.find('size')
    w = int(size.find('width').text)
    h = int(size.find('height').text)

    yolo_lines = []

    for obj in root.findall('object'):
        cls_name = obj.find('name').text
        if cls_name not in classes:
            continue
        cls_id = classes.index(cls_name)

        bndbox = obj.find('bndbox')
        xmin = float(bndbox.find('xmin').text)
        ymin = float(bndbox.find('ymin').text)
        xmax = float(bndbox.find('xmax').text)
        ymax = float(bndbox.find('ymax').text)

        # Conversão para YOLO
        x_center = ((xmin + xmax) / 2) / w
        y_center = ((ymin + ymax) / 2) / h
        width = (xmax - xmin) / w
        height = (ymax - ymin) / h

        yolo_lines.append(f"{cls_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}")

    txt_file = xml_file.replace('.xml', '.txt')
    with open(os.path.join(yolo_dir, txt_file), 'w') as f:
        f.write('\n'.join(yolo_lines))

print("✅ Conversão concluída!")