"""
Author: Fanxing Shang
Date: 2023-1-10
   //////////////////// __voc-yolo脚本__  ->  'voc_yolo.py'//////////////////////

    __用途__  ->  将voc格式的.xml数据集转换陈yolo格式的.txt数据集
"""

import os
import xml.etree.ElementTree as ET


def convert_voc_to_yolo(voc_dir, yolo_dir, classes_file):
    # 创建YOLO标注文件目录
    if not os.path.exists(yolo_dir):
        os.makedirs(yolo_dir)

    # 读取类别文件
    with open(classes_file) as f:
        classes = f.read().strip().split()

    # 处理每一个VOC XML标注文件
    for xml_file in os.listdir(voc_dir):
        if xml_file.endswith('.xml'):
            xml_path = os.path.join(voc_dir, xml_file)
            tree = ET.parse(xml_path)
            root = tree.getroot()

            # 获取图片尺寸
            size = root.find('size')
            width = int(size.find('width').text)
            height = int(size.find('height').text)

            yolo_annotations = []

            # 获取每一个目标对象
            for obj in root.findall('object'):
                cls = obj.find('name').text
                if cls not in classes:
                    continue
                cls_id = classes.index(cls)

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

                # 计算YOLO格式的坐标
                x_center = (xmin + xmax) / 2.0 / width
                y_center = (ymin + ymax) / 2.0 / height
                w = (xmax - xmin) / width
                h = (ymax - ymin) / height

                # 格式化为六位小数
                yolo_annotations.append(f"{cls_id} {x_center:.6f} {y_center:.6f} {w:.6f} {h:.6f}")

            # 保存YOLO标注文件
            yolo_file = os.path.join(yolo_dir, xml_file.replace('.xml', '.txt'))
            with open(yolo_file, 'w') as f:
                f.write('\n'.join(yolo_annotations))


# 设置路径
voc_dir = 'd:\\Desktop\\奴隶矿场\\9\\labelimg'  # VOC XML标注文件目录
yolo_dir = 'd:\\Desktop\\奴隶矿场\\9\\labels'  # YOLO标注文件目录
classes_file = 'd:\\Desktop\\奴隶矿场\\9\\classes.txt'  # 类别文件，每行一个类别名称

# 转换
convert_voc_to_yolo(voc_dir, yolo_dir, classes_file)