VOC格式数据集标签转化为YOLO格式标签

VOC格式数据集标签转化为YOLO格式标签

import os
import xml.etree.ElementTree as ET

def convert_voc_to_yolo(voc_xml_path, yolo_txt_path, class_id_map):
    tree = ET.parse(voc_xml_path)
    root = tree.getroot()

    size = root.find('size')
    img_width = float(size.find('width').text)
    img_height = float(size.find('height').text)

    with open(yolo_txt_path, 'w') as yolo_file:
        for obj in root.findall('object'):
            class_name = obj.find('name').text
            if class_name not in class_id_map:
                continue

            class_id = class_id_map[class_name]

            bbox = obj.find('bndbox')
            x_min = float(bbox.find('xmin').text)
            y_min = float(bbox.find('ymin').text)
            x_max = float(bbox.find('xmax').text)
            y_max = float(bbox.find('ymax').text)

            x_center = (x_min + x_max) / (2.0 * img_width)
            y_center = (y_min + y_max) / (2.0 * img_height)
            width = (x_max - x_min) / img_width
            height = (y_max - y_min) / img_height

            yolo_line = f"{class_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}\n"
            yolo_file.write(yolo_line)

def batch_convert_voc_to_yolo(input_folder, output_folder, class_id_map):
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    for filename in os.listdir(input_folder):
        if filename.endswith('.xml'):
            voc_xml_path = os.path.join(input_folder, filename)
            yolo_txt_path = os.path.join(output_folder, os.path.splitext(filename)[0] + '.txt')
            convert_voc_to_yolo(voc_xml_path, yolo_txt_path, class_id_map)

if __name__ == "__main__":
    input_folder = r"F:\HSR_Date\zc\anotations_dj"
    output_folder = r"F:\HSR_Date\zc\labels_dj"
    # class_id_map = {"Brace sleeve-f": 0 , "Brace sleeve-b": 0}  # Update with your class names and IDs
    class_id_map = {"DJLS": 0}  # Update with your class names and IDs

    batch_convert_voc_to_yolo(input_folder, output_folder, class_id_map)