127 lines
3.9 KiB
Python
127 lines
3.9 KiB
Python
# cvat_xml_to_yolo_obb.py
|
|
# 仅在有标注时生成 .txt 文件
|
|
|
|
import xml.etree.ElementTree as ET
|
|
import numpy as np
|
|
from pathlib import Path
|
|
|
|
def rotate_box(x_center, y_center, w, h, angle_deg):
|
|
"""
|
|
将旋转框转为 4 个角点坐标（未归一化）
|
|
"""
|
|
angle_rad = np.radians(angle_deg)
|
|
|
|
# 四个角点相对于中心的偏移
|
|
corners = np.array([
|
|
[-w/2, -h/2],
|
|
[ w/2, -h/2],
|
|
[ w/2, h/2],
|
|
[-w/2, h/2]
|
|
])
|
|
|
|
# 旋转矩阵（顺时针）
|
|
cos_a, sin_a = np.cos(angle_rad), np.sin(angle_rad)
|
|
rotation_matrix = np.array([[cos_a, -sin_a],
|
|
[sin_a, cos_a]])
|
|
|
|
# 旋转并平移
|
|
rotated_corners = np.dot(corners, rotation_matrix.T) + np.array([x_center, y_center])
|
|
return rotated_corners # 返回 (4, 2) 数组
|
|
|
|
def cvat_xml_to_yolo_obb(cvat_xml_path, output_dir, class_name_to_id=None):
|
|
"""
|
|
将 CVAT annotations.xml 转为 YOLO-OBB 格式
|
|
- 仅在有有效标注时创建 .txt 文件
|
|
"""
|
|
if class_name_to_id is None:
|
|
class_name_to_id = {"clamp": 0} # ✅ 请根据你的实际类别修改
|
|
|
|
tree = ET.parse(cvat_xml_path)
|
|
root = tree.getroot()
|
|
|
|
# 创建 labels 输出目录
|
|
labels_dir = Path(output_dir) / "labels"
|
|
labels_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
# 统计信息
|
|
processed_images = 0
|
|
saved_files = 0
|
|
skipped_classes = 0
|
|
|
|
for image_elem in root.findall("image"):
|
|
image_name = image_elem.get("name")
|
|
img_w = float(image_elem.get("width"))
|
|
img_h = float(image_elem.get("height"))
|
|
|
|
label_file = (labels_dir / Path(image_name).stem).with_suffix(".txt")
|
|
|
|
boxes = image_elem.findall("box")
|
|
valid_annotations = []
|
|
|
|
for box in boxes:
|
|
label = box.get("label")
|
|
if label not in class_name_to_id:
|
|
print(f"⚠️ 跳过未知类别: {label} (图片: {image_name})")
|
|
skipped_classes += 1
|
|
continue
|
|
|
|
class_id = class_name_to_id[label]
|
|
xtl = float(box.get("xtl"))
|
|
ytl = float(box.get("ytl"))
|
|
xbr = float(box.get("xbr"))
|
|
ybr = float(box.get("ybr"))
|
|
|
|
# 计算中心点和宽高
|
|
x_center = (xtl + xbr) / 2
|
|
y_center = (ytl + ybr) / 2
|
|
w = xbr - xtl
|
|
h = ybr - ytl
|
|
|
|
# 获取旋转角度
|
|
angle = float(box.get("rotation", 0.0))
|
|
|
|
# 计算 4 个角点
|
|
corners = rotate_box(x_center, y_center, w, h, angle)
|
|
|
|
# 归一化到 [0,1]
|
|
corners[:, 0] /= img_w
|
|
corners[:, 1] /= img_h
|
|
|
|
# 展平并生成行
|
|
points = corners.flatten()
|
|
line = str(class_id) + " " + " ".join(f"{coord:.6f}" for coord in points)
|
|
valid_annotations.append(line)
|
|
|
|
# ✅ 只有存在有效标注时才写入文件
|
|
if valid_annotations:
|
|
with open(label_file, 'w') as f:
|
|
f.write("\n".join(valid_annotations) + "\n")
|
|
saved_files += 1
|
|
print(f"✅ 已生成: {label_file}")
|
|
# else: # 无标注，不创建文件
|
|
# print(f"🟡 无标注，跳过: {image_name}")
|
|
|
|
processed_images += 1
|
|
|
|
print(f"\n🎉 转换完成！")
|
|
print(f"📊 处理图片: {processed_images}")
|
|
print(f"✅ 生成标签: {saved_files}")
|
|
if skipped_classes:
|
|
print(f"⚠️ 跳过类别: {skipped_classes} 个标注")
|
|
|
|
# ==================== 使用示例 ====================
|
|
|
|
if __name__ == "__main__":
|
|
# ✅ 请修改以下路径和类别
|
|
CVAT_XML_PATH = "annotations.xml" # 你的 annotations.xml 文件
|
|
OUTPUT_DIR = "yolo_obb_dataset1" # 输出目录
|
|
CLASS_MAPPING = {
|
|
"clamp": 0,
|
|
}
|
|
|
|
# 执行转换
|
|
cvat_xml_to_yolo_obb(
|
|
cvat_xml_path=CVAT_XML_PATH,
|
|
output_dir=OUTPUT_DIR,
|
|
class_name_to_id=CLASS_MAPPING
|
|
) |