142 lines
4.1 KiB
Python
142 lines
4.1 KiB
Python
import os
|
||
import xml.etree.ElementTree as ET
|
||
from pathlib import Path
|
||
|
||
|
||
def yolo_detect_to_cvat(
|
||
images_dir,
|
||
labels_dir,
|
||
output_xml,
|
||
class_id_to_name,
|
||
):
|
||
"""
|
||
将 YOLO Detect 格式转换为 CVAT XML（目标检测 box）
|
||
|
||
Args:
|
||
images_dir (str): 图片目录
|
||
labels_dir (str): YOLO txt 标注目录
|
||
output_xml (str): 输出 CVAT XML 路径
|
||
class_id_to_name (dict): {0: "xxx", 1: "yyy"}
|
||
"""
|
||
|
||
images_dir = Path(images_dir)
|
||
labels_dir = Path(labels_dir)
|
||
|
||
# ----------------------------
|
||
# 创建 XML 结构
|
||
# ----------------------------
|
||
annotations = ET.Element("annotations")
|
||
|
||
version = ET.SubElement(annotations, "version")
|
||
version.text = "1.1"
|
||
|
||
# labels
|
||
meta = ET.SubElement(annotations, "meta")
|
||
task = ET.SubElement(meta, "task")
|
||
labels_elem = ET.SubElement(task, "labels")
|
||
|
||
for class_id, name in class_id_to_name.items():
|
||
label = ET.SubElement(labels_elem, "label")
|
||
name_elem = ET.SubElement(label, "name")
|
||
name_elem.text = name
|
||
|
||
image_id = 0
|
||
|
||
# ----------------------------
|
||
# 遍历图片
|
||
# ----------------------------
|
||
for img_path in sorted(images_dir.iterdir()):
|
||
if img_path.suffix.lower() not in [".jpg", ".png", ".jpeg"]:
|
||
continue
|
||
|
||
import cv2
|
||
img = cv2.imread(str(img_path))
|
||
if img is None:
|
||
print(f"⚠️ 无法读取图片: {img_path}")
|
||
continue
|
||
|
||
height, width = img.shape[:2]
|
||
|
||
image_elem = ET.SubElement(
|
||
annotations,
|
||
"image",
|
||
{
|
||
"id": str(image_id),
|
||
"name": img_path.name,
|
||
"width": str(width),
|
||
"height": str(height),
|
||
}
|
||
)
|
||
|
||
label_txt = labels_dir / f"{img_path.stem}.txt"
|
||
|
||
# ----------------------------
|
||
# 有 / 无标签都要建 image
|
||
# ----------------------------
|
||
if label_txt.exists():
|
||
with open(label_txt, "r") as f:
|
||
for line in f:
|
||
line = line.strip()
|
||
if not line:
|
||
continue
|
||
|
||
class_id, xc, yc, w, h = map(float, line.split())
|
||
class_id = int(class_id)
|
||
|
||
label_name = class_id_to_name.get(class_id)
|
||
if label_name is None:
|
||
print(f"⚠️ 未知 class_id: {class_id}")
|
||
continue
|
||
|
||
# YOLO → CVAT 坐标
|
||
box_w = w * width
|
||
box_h = h * height
|
||
xtl = xc * width - box_w / 2
|
||
ytl = yc * height - box_h / 2
|
||
xbr = xtl + box_w
|
||
ybr = ytl + box_h
|
||
|
||
ET.SubElement(
|
||
image_elem,
|
||
"box",
|
||
{
|
||
"label": label_name,
|
||
"xtl": f"{xtl:.2f}",
|
||
"ytl": f"{ytl:.2f}",
|
||
"xbr": f"{xbr:.2f}",
|
||
"ybr": f"{ybr:.2f}",
|
||
"occluded": "0",
|
||
"source": "manual",
|
||
}
|
||
)
|
||
|
||
image_id += 1
|
||
print(f"✅ {img_path.name}")
|
||
|
||
# ----------------------------
|
||
# 写 XML
|
||
# ----------------------------
|
||
tree = ET.ElementTree(annotations)
|
||
tree.write(output_xml, encoding="utf-8", xml_declaration=True)
|
||
|
||
print(f"\n🎉 转换完成，CVAT XML 已生成：{output_xml}")
|
||
|
||
if __name__ == "__main__":
|
||
IMAGES_DIR = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/detect/1"
|
||
LABELS_DIR = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/detect/1"
|
||
OUTPUT_XML = "annotations.xml"
|
||
|
||
# ⚠️ 一定要和 YOLO 训练时一致
|
||
CLASS_ID_TO_NAME = {
|
||
#0: "hole",
|
||
#1: "crack"
|
||
0: "bag"
|
||
}
|
||
|
||
yolo_detect_to_cvat(
|
||
images_dir=IMAGES_DIR,
|
||
labels_dir=LABELS_DIR,
|
||
output_xml=OUTPUT_XML,
|
||
class_id_to_name=CLASS_ID_TO_NAME
|
||
)
|