159 lines
4.7 KiB
Python
159 lines
4.7 KiB
Python
|
# yolo_obb_to_cvat_full.py
|
|||
|
|
import os
|
|||
|
|
import math
|
|||
|
|
import xml.etree.ElementTree as ET
|
|||
|
|
from pathlib import Path
|
|||
|
|
import cv2
|
|||
|
|
|
|||
|
|
|
|||
|
|
def recover_rotated_box(points):
|
|||
|
|
"""
|
|||
|
|
YOLO OBB 四点 → 中心点、宽高、旋转角度
|
|||
|
|
点顺序为：p1 p2 p3 p4，逆时针/顺时针均可接受
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
cx = sum(p[0] for p in points) / 4
|
|||
|
|
cy = sum(p[1] for p in points) / 4
|
|||
|
|
|
|||
|
|
# 点1→点2 作为宽方向
|
|||
|
|
vx = points[1][0] - points[0][0]
|
|||
|
|
vy = points[1][1] - points[0][1]
|
|||
|
|
|
|||
|
|
w = math.sqrt(vx*vx + vy*vy)
|
|||
|
|
|
|||
|
|
# 点2→点3 作为高方向
|
|||
|
|
vx2 = points[2][0] - points[1][0]
|
|||
|
|
vy2 = points[2][1] - points[1][1]
|
|||
|
|
|
|||
|
|
h = math.sqrt(vx2*vx2 + vy2*vy2)
|
|||
|
|
|
|||
|
|
# 角度 = 边1 与水平线夹角
|
|||
|
|
angle_rad = math.atan2(vy, vx)
|
|||
|
|
angle_deg = math.degrees(angle_rad)
|
|||
|
|
|
|||
|
|
# CVAT 角度必须是 0~360
|
|||
|
|
if angle_deg < 0:
|
|||
|
|
angle_deg += 360
|
|||
|
|
|
|||
|
|
return cx, cy, w, h, angle_deg
|
|||
|
|
|
|||
|
|
|
|||
|
|
def yolo_obb_to_cvat_xml(label_dir, image_dir, class_id_to_name, output_xml):
|
|||
|
|
|
|||
|
|
label_dir = Path(label_dir)
|
|||
|
|
image_dir = Path(image_dir)
|
|||
|
|
|
|||
|
|
# ======== 构建基本 XML 结构（与你发的一模一样） ========
|
|||
|
|
root = ET.Element("annotations")
|
|||
|
|
|
|||
|
|
ET.SubElement(root, "version").text = "1.1"
|
|||
|
|
|
|||
|
|
# meta/task
|
|||
|
|
meta = ET.SubElement(root, "meta")
|
|||
|
|
task = ET.SubElement(meta, "task")
|
|||
|
|
|
|||
|
|
ET.SubElement(task, "id").text = "1"
|
|||
|
|
ET.SubElement(task, "name").text = "yolo_obb_import"
|
|||
|
|
ET.SubElement(task, "size").text = str(len(list(label_dir.glob("*.txt"))))
|
|||
|
|
ET.SubElement(task, "mode").text = "annotation"
|
|||
|
|
ET.SubElement(task, "overlap").text = "0"
|
|||
|
|
ET.SubElement(task, "bugtracker").text = ""
|
|||
|
|
ET.SubElement(task, "created").text = ""
|
|||
|
|
ET.SubElement(task, "updated").text = ""
|
|||
|
|
ET.SubElement(task, "subset").text = "default"
|
|||
|
|
ET.SubElement(task, "start_frame").text = "0"
|
|||
|
|
ET.SubElement(task, "stop_frame").text = str(len(list(label_dir.glob("*.txt"))) - 1)
|
|||
|
|
ET.SubElement(task, "frame_filter").text = ""
|
|||
|
|
|
|||
|
|
# labels
|
|||
|
|
labels = ET.SubElement(task, "labels")
|
|||
|
|
for name in class_id_to_name.values():
|
|||
|
|
lab = ET.SubElement(labels, "label")
|
|||
|
|
ET.SubElement(lab, "name").text = name
|
|||
|
|
ET.SubElement(lab, "color").text = "#ffffff"
|
|||
|
|
ET.SubElement(lab, "type").text = "any"
|
|||
|
|
ET.SubElement(lab, "attributes")
|
|||
|
|
|
|||
|
|
ET.SubElement(meta, "dumped").text = ""
|
|||
|
|
|
|||
|
|
# ======== 处理每张图片 ========
|
|||
|
|
for idx, txt_file in enumerate(sorted(label_dir.glob("*.txt"))):
|
|||
|
|
stem = txt_file.stem
|
|||
|
|
img_file = None
|
|||
|
|
|
|||
|
|
# 自动匹配图片
|
|||
|
|
for ext in ["jpg", "png", "jpeg", "bmp", "webp"]:
|
|||
|
|
p = image_dir / f"{stem}.{ext}"
|
|||
|
|
if p.exists():
|
|||
|
|
img_file = p
|
|||
|
|
break
|
|||
|
|
|
|||
|
|
if img_file is None:
|
|||
|
|
print(f"⚠ 找不到图片: {stem}")
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
img = cv2.imread(str(img_file))
|
|||
|
|
H, W = img.shape[:2]
|
|||
|
|
|
|||
|
|
# image tag
|
|||
|
|
image_elem = ET.SubElement(root, "image", {
|
|||
|
|
"id": str(idx),
|
|||
|
|
"name": img_file.name,
|
|||
|
|
"width": str(W),
|
|||
|
|
"height": str(H)
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
# 读取 txt
|
|||
|
|
with open(txt_file, "r") as f:
|
|||
|
|
for line in f.readlines():
|
|||
|
|
parts = line.strip().split()
|
|||
|
|
cls_id = int(parts[0])
|
|||
|
|
coords = list(map(float, parts[1:]))
|
|||
|
|
|
|||
|
|
pts = []
|
|||
|
|
for i in range(0, 8, 2):
|
|||
|
|
x = coords[i] * W
|
|||
|
|
y = coords[i+1] * H
|
|||
|
|
pts.append((x, y))
|
|||
|
|
|
|||
|
|
# 四点 → 框参数
|
|||
|
|
cx, cy, bw, bh, ang = recover_rotated_box(pts)
|
|||
|
|
|
|||
|
|
xtl = cx - bw / 2
|
|||
|
|
ytl = cy - bh / 2
|
|||
|
|
xbr = cx + bw / 2
|
|||
|
|
ybr = cy + bh / 2
|
|||
|
|
|
|||
|
|
# box 标签（完全跟你的一样）
|
|||
|
|
ET.SubElement(image_elem, "box", {
|
|||
|
|
"label": class_id_to_name[cls_id],
|
|||
|
|
"source": "manual",
|
|||
|
|
"occluded": "0",
|
|||
|
|
"xtl": f"{xtl:.2f}",
|
|||
|
|
"ytl": f"{ytl:.2f}",
|
|||
|
|
"xbr": f"{xbr:.2f}",
|
|||
|
|
"ybr": f"{ybr:.2f}",
|
|||
|
|
"rotation": f"{ang:.2f}",
|
|||
|
|
"z_order": "0"
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
print(f"✔ 处理 {img_file.name}")
|
|||
|
|
|
|||
|
|
# 保存 XML
|
|||
|
|
tree = ET.ElementTree(root)
|
|||
|
|
tree.write(output_xml, encoding="utf-8", xml_declaration=True)
|
|||
|
|
print(f"\n🎉 已生成 XML：{output_xml}")
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
CLASS_MAP = {
|
|||
|
|
0: "clamp1",
|
|||
|
|
1: "clamp0",
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
yolo_obb_to_cvat_xml(
|
|||
|
|
label_dir="/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/1/222",
|
|||
|
|
image_dir="/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/1/111",
|
|||
|
|
class_id_to_name=CLASS_MAP,
|
|||
|
|
output_xml="converted_annotations.xml"
|
|||
|
|
)
|