107 lines
3.1 KiB
Python
107 lines
3.1 KiB
Python
|
||
import os
|
||
import xml.etree.ElementTree as ET
|
||
from pathlib import Path
|
||
|
||
|
||
def cvat_to_yolo_seg(
|
||
xml_path,
|
||
output_dir,
|
||
class_name_to_id=None,
|
||
force_class_id=None
|
||
):
|
||
"""
|
||
将 CVAT 导出的 XML（polygon / segmentation）转换为 YOLO Segmentation 格式
|
||
|
||
Args:
|
||
xml_path (str): CVAT 导出的 XML 文件路径
|
||
output_dir (str): 输出 .txt 标注文件的目录
|
||
class_name_to_id (dict, optional):
|
||
类别名到 ID 的映射
|
||
force_class_id (dict, optional):
|
||
强制指定某些类别的 ID，例如 {"yemian": 0}
|
||
"""
|
||
|
||
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:
|
||
raise RuntimeError("❌ 未找到 <labels>，请手动提供 class_name_to_id")
|
||
|
||
print(f"原始类别映射: {class_name_to_id}")
|
||
|
||
# ----------------------------
|
||
# 强制修改类别 ID（新增功能）
|
||
# ----------------------------
|
||
if force_class_id:
|
||
for name, new_id in force_class_id.items():
|
||
if name in class_name_to_id:
|
||
old_id = class_name_to_id[name]
|
||
class_name_to_id[name] = new_id
|
||
print(f"强制修改类别映射: {name} {old_id} → {new_id}")
|
||
else:
|
||
print(f"类别 {name} 不存在，跳过")
|
||
|
||
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_path = output_dir / (Path(img_name).stem + ".txt")
|
||
lines = []
|
||
|
||
for polygon in image.findall("polygon"):
|
||
label = polygon.get("label")
|
||
if label not in class_name_to_id:
|
||
continue
|
||
|
||
class_id = class_name_to_id[label]
|
||
points_str = polygon.get("points")
|
||
|
||
points = []
|
||
for p in points_str.strip().split(";"):
|
||
x, y = p.split(",")
|
||
x = float(x) / width
|
||
y = float(y) / height
|
||
points.append(f"{x:.6f}")
|
||
points.append(f"{y:.6f}")
|
||
|
||
if len(points) < 6:
|
||
continue
|
||
|
||
lines.append(f"{class_id} " + " ".join(points))
|
||
|
||
with open(txt_path, "w", encoding="utf-8") as f:
|
||
f.write("\n".join(lines))
|
||
|
||
print("✅ CVAT segmentation → YOLO seg 转换完成")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
|
||
cvat_to_yolo_seg(
|
||
xml_path="annotations.xml",
|
||
output_dir="labels_seg",
|
||
force_class_id={
|
||
"yemian": 0
|
||
}
|
||
)
|