三维重构终版

yolov8/voc_to_yolo.py

@@ -0,0 +1,119 @@
+import xml.etree.ElementTree as ET
+import os
+import glob
+
+
+def voc_to_yolo(xml_file_path, output_dir, class_mapping):
+    """
+    将单个 PASCAL VOC anntation (.xml) 文件转换为 YOLO (.txt) 格式。
+
+    Args:
+        xml_file_path (str): 输入的 .xml 文件路径。
+        output_dir (str): 输出 .txt 文件的目标文件夹。
+        class_mapping (dict): 类别名称到类别ID的映射字典。
+    """
+    try:
+        # 解析 XML 文件
+        tree = ET.parse(xml_file_path)
+        root = tree.getroot()
+
+        # 获取图像尺寸
+        size = root.find('size')
+        if size is None:
+            print(f"警告: 在 {xml_file_path} 中未找到 <size> 标签，跳过此文件。")
+            return
+
+        img_width = int(size.find('width').text)
+        img_height = int(size.find('height').text)
+
+        # 准备用于写入的YOLO标注列表
+        yolo_annotations = []
+
+        # 遍历所有 object
+        for obj in root.findall('object'):
+            # 获取类别名称
+            class_name = obj.find('name').text
+            if class_name not in class_mapping:
+                print(f"警告: 类别 '{class_name}' 不在预定义的 class_mapping 中，跳过此物体。")
+                continue
+
+            class_id = class_mapping[class_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)
+
+            # --- 核心转换公式 ---
+            x_center = (xmin + xmax) / 2.0 / img_width
+            y_center = (ymin + ymax) / 2.0 / img_height
+            width = (xmax - xmin) / img_width
+            height = (ymax - ymin) / img_height
+
+            # 将结果添加到列表
+            yolo_annotations.append(f"{class_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}")
+
+        # 如果文件中有有效的物体，则写入 .txt 文件
+        if yolo_annotations:
+            # 构建输出文件名
+            base_filename = os.path.basename(xml_file_path)
+            txt_filename = os.path.splitext(base_filename)[0] + '.txt'
+            output_path = os.path.join(output_dir, txt_filename)
+
+            # 写入文件
+            with open(output_path, 'w') as f:
+                f.write('\n'.join(yolo_annotations))
+
+            # print(f"成功转换: {xml_file_path} -> {output_path}")
+
+    except Exception as e:
+        print(f"处理文件 {xml_file_path} 时发生错误: {e}")
+
+
+def main():
+    # --- 用户需要配置的参数 ---
+
+    # 1. 定义你的类别和对应的ID (从0开始)
+    # 根据你的截图，你只有一个类别 "Space weld workpiece"
+    # 请确保这里的名称与你XML文件中的<name>标签完全一致！
+    CLASS_MAPPING = {
+        'Space weld workpiece': 0,
+        # 如果有其他类别，在这里添加，例如:
+        # 'weld_seam': 1,
+    }
+
+    # 2. 定义输入和输出文件夹
+    # 输入文件夹: 存放 .xml 文件的目录
+    input_xml_dir = '../label/up_xml'
+
+    # 输出文件夹: 存放转换后的 .txt 文件的目录
+    output_txt_dir = 'data_up'
+
+    # --- 脚本执行部分 ---
+
+    # 自动创建输出文件夹（如果不存在）
+    if not os.path.exists(output_txt_dir):
+        os.makedirs(output_txt_dir)
+        print(f"已创建输出文件夹: {output_txt_dir}")
+
+    # 查找所有 .xml 文件
+    xml_files = glob.glob(os.path.join(input_xml_dir, '*.xml'))
+
+    if not xml_files:
+        print(f"错误: 在目录 '{input_xml_dir}' 中没有找到任何 .xml 文件。请检查路径。")
+        return
+
+    print(f"找到 {len(xml_files)} 个 .xml 文件。开始转换...")
+
+    # 遍历并转换每个文件
+    for xml_file in xml_files:
+        voc_to_yolo(xml_file, output_txt_dir, CLASS_MAPPING)
+
+    print("\n转换完成！")
+    print(f"所有 YOLO 格式的标签文件已保存在: {output_txt_dir}")
+
+
+if __name__ == '__main__':
+    main()