最新推送

yolo11_detect/cvattodetect.py

@@ -0,0 +1,100 @@
+import os
+import xml.etree.ElementTree as ET
+from pathlib import Path
+
+
+def cvat_to_yolo_detect(xml_path, output_dir, class_name_to_id=None):
+    """
+    将 CVAT 导出的 XML（目标检测模式）转换为 YOLO Detect 格式
+
+    Args:
+        xml_path (str): CVAT 导出的 XML 文件路径
+        output_dir (str): 输出 .txt 标注文件的目录
+        class_name_to_id (dict, optional): 类别名到 ID 的映射。
+            如果为 None，则自动从 XML 的 <labels> 中按顺序分配（0,1,2...）
+    """
+    output_dir = Path(output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    tree = ET.parse(xml_path)
+    root = tree.getroot()
+
+    # 自动提取类别顺序（如果未提供映射）
+    if class_name_to_id is None:
+        class_name_to_id = {}
+        labels_elem = root.find(".//labels")
+        if labels_elem is not None:
+            for idx, label in enumerate(labels_elem.findall("label")):
+                name = label.find("name").text
+                class_name_to_id[name] = idx
+        else:
+            print("⚠️ 未找到 <labels>，请手动提供 class_name_to_id")
+
+    print(f"类别映射: {class_name_to_id}")
+
+    # 遍历所有 <image>
+    for image in root.findall("image"):
+        img_name = image.get("name")
+        width = int(image.get("width"))
+        height = int(image.get("height"))
+
+        # 构建 .txt 文件名（去掉扩展名）
+        stem = Path(img_name).stem
+        txt_path = output_dir / f"{stem}.txt"
+
+        boxes = []
+        for box in image.findall("box"):
+            label = box.get("label")
+            if label not in class_name_to_id:
+                print(f"⚠️ 未知类别 '{label}'，跳过（图片: {img_name}）")
+                continue
+
+            class_id = class_name_to_id[label]
+            xtl = float(box.get("xtl"))
+            ytl = float(box.get("ytl"))
+            xbr = float(box.get("xbr"))
+            ybr = float(box.get("ybr"))
+
+            # 转为 YOLO 格式（归一化）
+            x_center = (xtl + xbr) / 2 / width
+            y_center = (ytl + ybr) / 2 / height
+            w = (xbr - xtl) / width
+            h = (ybr - ytl) / height
+
+            # 限制在 [0,1]（防止因标注误差越界）
+            x_center = max(0.0, min(1.0, x_center))
+            y_center = max(0.0, min(1.0, y_center))
+            w = max(0.0, min(1.0, w))
+            h = max(0.0, min(1.0, h))
+
+            boxes.append(f"{class_id} {x_center:.6f} {y_center:.6f} {w:.6f} {h:.6f}")
+
+        # 写入 .txt 文件（即使无框也创建空文件）
+        with open(txt_path, "w") as f:
+            f.write("\n".join(boxes))
+
+        print(f"✅ {img_name} → {len(boxes)} 个目标")
+
+    print(f"\n🎉 转换完成！YOLO 标注已保存至: {output_dir}")
+
+
+if __name__ == "__main__":
+    # ====== 配置区 ======
+    XML_PATH = "annotations.xml"  # 替换为你的 CVAT XML 路径
+    OUTPUT_LABELS_DIR = "labels"  # 输出的 YOLO .txt 目录
+
+    # 方式1：自动从 XML 提取类别（推荐）
+    #CLASS_MAP = None
+
+    # 方式2：手动指定（确保与训练时一致）
+    CLASS_MAP = {
+         "bag": 0,
+         "bag35": 1,
+    }
+
+    # ====== 执行转换 ======
+    cvat_to_yolo_detect(
+        xml_path=XML_PATH,
+        output_dir=OUTPUT_LABELS_DIR,
+        class_name_to_id=CLASS_MAP
+    )