XML标签格式转换为YOLO TXT格式

针对的是多边形（<polygon>）来描述对象的边界，而不是传统的矩形框（<bndbox>）

import xml.etree.ElementTree as ET
import os
from pathlib import Path# 解析VOC格式的XML文件，提取目标框的标签和坐标
def parse_voc_xml(xml_file):tree = ET.parse(xml_file)root = tree.getroot()# 获取图像的宽度和高度size = root.find('size')img_width = int(size.find('width').text)img_height = int(size.find('height').text)bboxes = []for obj in root.findall('object'):label = obj.find('name').text  # 获取目标的标签polygon = obj.find('polygon')  # 获取多边形的坐标# 获取多边形的所有顶点坐标x_coords = []y_coords = []for i in range(1, 5):  # 假设最多4个点（也可以扩展）x_coords.append(int(polygon.find(f'x{i}').text))y_coords.append(int(polygon.find(f'y{i}').text))# 计算边界框（最小矩形框）的坐标xmin = min(x_coords)xmax = max(x_coords)ymin = min(y_coords)ymax = max(y_coords)# 将 VOC 格式的坐标转换为 YOLO 格式x_center = (xmin + xmax) / 2y_center = (ymin + ymax) / 2width = xmax - xminheight = ymax - ymin# 归一化坐标x_center /= img_widthy_center /= img_heightwidth /= img_widthheight /= img_heightbboxes.append((label, x_center, y_center, width, height))return bboxes, img_width, img_height# 将标签保存为 YOLO 格式
def save_yolo_labels(image_filename, bboxes, output_dir):# YOLO 标签文件的路径txt_filename = os.path.join(output_dir, Path(image_filename).stem + '.txt')with open(txt_filename, 'w') as f:for label, x_center, y_center, width, height in bboxes:# YOLO 格式为：<class_id> <x_center> <y_center> <width> <height># 假设标签的类别是数字（如果是字符串，则需要将类别转换为数字映射）f.write(f"{label} {x_center} {y_center} {width} {height}\n")# 主函数：将 VOC XML 文件转换为 YOLO 格式
def convert_voc_to_yolo(voc_xml_dir, output_dir, class_mapping):# 如果输出目录不存在，则创建if not os.path.exists(output_dir):os.makedirs(output_dir)# 遍历所有 XML 文件for xml_file in Path(voc_xml_dir).glob('*.xml'):print(f"Processing {xml_file}")# 解析 VOC XML 文件bboxes, img_width, img_height = parse_voc_xml(xml_file)# 将标签转化为 YOLO 格式，且需要使用数字类别而不是字符串bboxes = [(class_mapping[label], x_center, y_center, width, height) for label, x_center, y_center, width, height in bboxes]# 保存 YOLO 格式的标签save_yolo_labels(xml_file.stem + '.jpg', bboxes, output_dir)  # 假设图像文件与XML文件同名，后缀为 .jpg# 主程序入口
if __name__ == "__main__":# 输入目录：包含VOC XML标注文件的目录voc_xml_dir = './original/labels'  # 替换为 VOC XML 文件目录路径# 输出目录：YOLO 格式标签文件的存储目录output_dir = './original/save'  # 替换为 YOLO 标签文件存储目录路径# 类别映射：将VOC中的标签名映射为数字类别class_mapping = {'car': 0,  'truck': 1,'bus': 2,'motorcycle': 3,'van': 4,'freight_car': 5,}# 转换 VOC 格式标签为 YOLO 格式convert_voc_to_yolo(voc_xml_dir, output_dir, class_mapping)
# 'car', 'truck', 'bus', 'motorcycle', 'van', 'freight_car'