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")
|