123 lines
3.4 KiB
Python
123 lines
3.4 KiB
Python
|
# convert_yolo_seg_to_cvat.py
|
|||
|
|
import os
|
|||
|
|
import cv2
|
|||
|
|
import xml.etree.ElementTree as ET
|
|||
|
|
|
|||
|
|
# ================== 配置 ==================
|
|||
|
|
labels_dir = "./labels/labels" # YOLO segmentation txt 目录
|
|||
|
|
images_dir = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/6111c/1"
|
|||
|
|
output_xml = "2annotations_cvat_seg.xml"
|
|||
|
|
|
|||
|
|
# 类别映射（id -> name）
|
|||
|
|
class_mapping_reverse = {
|
|||
|
|
0: "yemian",
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# ================== 工具函数 ==================
|
|||
|
|
def unnormalize_polygon(img_w, img_h, coords):
|
|||
|
|
"""
|
|||
|
|
YOLO seg 归一化坐标 -> 像素坐标
|
|||
|
|
coords: [x1, y1, x2, y2, ...]
|
|||
|
|
"""
|
|||
|
|
pts = []
|
|||
|
|
for i in range(0, len(coords), 2):
|
|||
|
|
x = float(coords[i]) * img_w
|
|||
|
|
y = float(coords[i + 1]) * img_h
|
|||
|
|
pts.append((x, y))
|
|||
|
|
return pts
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ================== 构建 XML ==================
|
|||
|
|
root = ET.Element("annotations")
|
|||
|
|
ET.SubElement(root, "version").text = "1.1"
|
|||
|
|
|
|||
|
|
meta = ET.SubElement(root, "meta")
|
|||
|
|
task = ET.SubElement(meta, "task")
|
|||
|
|
ET.SubElement(task, "name").text = "converted_from_yolo_seg"
|
|||
|
|
|
|||
|
|
txt_files = sorted([f for f in os.listdir(labels_dir) if f.endswith(".txt")])
|
|||
|
|
ET.SubElement(task, "size").text = str(len(txt_files))
|
|||
|
|
|
|||
|
|
# labels
|
|||
|
|
labels_elem = ET.SubElement(task, "labels")
|
|||
|
|
for name in class_mapping_reverse.values():
|
|||
|
|
lab = ET.SubElement(labels_elem, "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 = ""
|
|||
|
|
|
|||
|
|
# ================== 处理每个 txt ==================
|
|||
|
|
for idx, txt_file in enumerate(txt_files):
|
|||
|
|
base = os.path.splitext(txt_file)[0]
|
|||
|
|
|
|||
|
|
# 查找图片
|
|||
|
|
img_path = None
|
|||
|
|
for ext in [".jpg", ".jpeg", ".png", ".bmp"]:
|
|||
|
|
p = os.path.join(images_dir, base + ext)
|
|||
|
|
if os.path.exists(p):
|
|||
|
|
img_path = p
|
|||
|
|
break
|
|||
|
|
|
|||
|
|
if img_path is None:
|
|||
|
|
print(f"❗缺少图片：{base}.*（跳过）")
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
img = cv2.imread(img_path)
|
|||
|
|
if img is None:
|
|||
|
|
print(f"❗无法读取图片：{img_path}")
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
h, w = img.shape[:2]
|
|||
|
|
|
|||
|
|
image_elem = ET.SubElement(
|
|||
|
|
root, "image",
|
|||
|
|
id=str(idx),
|
|||
|
|
name=os.path.basename(img_path),
|
|||
|
|
width=str(w),
|
|||
|
|
height=str(h)
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
yolo_path = os.path.join(labels_dir, txt_file)
|
|||
|
|
with open(yolo_path, "r") as f:
|
|||
|
|
lines = [l.strip() for l in f.readlines() if l.strip()]
|
|||
|
|
|
|||
|
|
if not lines:
|
|||
|
|
print(f"⚠ 空标签：{txt_file}")
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
# 支持一张图多个实例
|
|||
|
|
for line in lines:
|
|||
|
|
parts = line.split()
|
|||
|
|
cls_id = int(parts[0])
|
|||
|
|
coords = parts[1:]
|
|||
|
|
|
|||
|
|
if len(coords) < 6 or len(coords) % 2 != 0:
|
|||
|
|
print(f"⚠ 非法 segmentation 点数：{txt_file}")
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
label_name = class_mapping_reverse.get(cls_id, str(cls_id))
|
|||
|
|
|
|||
|
|
pts = unnormalize_polygon(w, h, coords)
|
|||
|
|
pts_str = ";".join([f"{x:.2f},{y:.2f}" for x, y in pts])
|
|||
|
|
|
|||
|
|
ET.SubElement(
|
|||
|
|
image_elem,
|
|||
|
|
"polygon",
|
|||
|
|
label=label_name,
|
|||
|
|
source="manual",
|
|||
|
|
occluded="0",
|
|||
|
|
points=pts_str,
|
|||
|
|
z_order="0"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# ================== 写出 XML ==================
|
|||
|
|
tree = ET.ElementTree(root)
|
|||
|
|
tree.write(output_xml, encoding="utf-8", xml_declaration=True)
|
|||
|
|
|
|||
|
|
print("\n🎉 YOLO Segmentation → CVAT 转换完成！")
|
|||
|
|
print(f"📄 输出文件：{output_xml}")
|
|||
|
|
print(f"📊 共处理 {len(txt_files)} 个 txt")
|