import os
import xml.etree.ElementTree as ET

def convert_voc_to_yolo(xml_path, output_dir):
    """将VOC格式XML转换为YOLO格式TXT"""
    tree = ET.parse(xml_path)
    root = tree.getroot()
    
    # 获取图像尺寸
    size = root.find('size')
    if size is None:
        return
        
    width_elem = size.find('width')
    height_elem = size.find('height')
    if width_elem is None or height_elem is None or width_elem.text is None or height_elem.text is None:
        return
        
    try:
        width = int(width_elem.text)
        height = int(height_elem.text)
    except (ValueError, TypeError):
        return
    
    # 类别映射
    class_map = {'猫': 0, '狗': 1}
    
    # 准备输出内容
    yolo_lines = []
    
    for obj in root.findall('object'):
        name_elem = obj.find('name')
        if name_elem is None or name_elem.text is None:
            continue
            
        class_name = name_elem.text
        class_id = class_map.get(class_name, -1)
        if class_id == -1:
            continue
            
        bndbox = obj.find('bndbox')
        if bndbox is None:
            continue
            
        try:
            xmin_elem = bndbox.find('xmin')
            xmin = int(xmin_elem.text) if xmin_elem is not None and xmin_elem.text is not None else 0
            
            ymin_elem = bndbox.find('ymin')
            ymin = int(ymin_elem.text) if ymin_elem is not None and ymin_elem.text is not None else 0
            
            xmax_elem = bndbox.find('xmax')
            xmax = int(xmax_elem.text) if xmax_elem is not None and xmax_elem.text is not None else 0
            
            ymax_elem = bndbox.find('ymax')
            ymax = int(ymax_elem.text) if ymax_elem is not None and ymax_elem.text is not None else 0
        except (ValueError, TypeError):
            continue
        
        # 转换为YOLO格式
        x_center = (xmin + xmax) / 2 / width
        y_center = (ymin + ymax) / 2 / height
        box_width = (xmax - xmin) / width
        box_height = (ymax - ymin) / height
        
        yolo_lines.append(f"{class_id} {x_center:.6f} {y_center:.6f} {box_width:.6f} {box_height:.6f}")
    
    # 写入TXT文件
    if yolo_lines:
        txt_filename = os.path.splitext(os.path.basename(xml_path))[0] + '.txt'
        txt_path = os.path.join(output_dir, txt_filename)
        
        with open(txt_path, 'w', encoding='utf-8') as f:
            f.write('\n'.join(yolo_lines))

def batch_convert(input_dir, output_dir):
    """批量转换目录中的所有XML文件"""
    os.makedirs(output_dir, exist_ok=True)
    
    for filename in os.listdir(input_dir):
        if filename.endswith('.xml'):
            xml_path = os.path.join(input_dir, filename)
            convert_voc_to_yolo(xml_path, output_dir)

if __name__ == "__main__":
    # 设置输入输出目录
    input_dir = r'F:\myprojects\啾啾救援识别系统\JJCarDetection\datasets\猫狗\train\labels'
    output_dir = input_dir  # 同一目录输出
    
    print(f"开始转换: {input_dir}")
    batch_convert(input_dir, output_dir)
    print(f"转换完成，结果保存在: {output_dir}")