158 lines
4.7 KiB
Python
158 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 = [1080,1920]
|
||
|
||
# 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"已生成 XML：{output_xml}")
|
||
|
||
if __name__ == "__main__":
|
||
CLASS_MAP = {
|
||
0: "clamp",
|
||
}
|
||
|
||
yolo_obb_to_cvat_xml(
|
||
label_dir="/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/1/labels/labels",
|
||
image_dir="/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/1/zjdata17",
|
||
class_id_to_name=CLASS_MAP,
|
||
output_xml="annotations.xml"
|
||
)
|