旋转目标检测数据格式转换：RoLabelImg 与 DOTA 格式

最新推荐文章于 2025-02-21 21:53:22 发布

Lunar*

最新推荐文章于 2025-02-21 21:53:22 发布

阅读量666

点赞数 9

分类专栏：目标检测文章标签：目标检测人工智能

本文链接：https://blog.csdn.net/qq_45141261/article/details/144528094

版权

目标检测专栏收录该内容

11 篇文章

订阅专栏

引言

在旋转目标检测任务中，数据格式的选择与处理对于模型的训练至关重要。本文将介绍两种常见的旋转目标检测数据格式：RoLabelImg 格式 和 DOTA 格式，并详细说明二者之间的转换方法及其实现代码。

数据格式简介

1. RoLabelImg 格式

RoLabelImg 是一种基于 XML 的标注格式，其结构类似于标准的 Pascal VOC 数据格式，但针对旋转框增加了 robndbox 标签，存储了目标的中心点坐标 (cx, cy)、宽 (w)、高 (h)、以及旋转角度 (angle)。
数据示例（XML 文件）：

<annotation>
    <folder>Unknown</folder>
    <filename>example.jpg</filename>
    <size>
        <width>1024</width>
        <height>768</height>
        <depth>3</depth>
    </size>
    <object>
        <name>car</name>
        <difficult>0</difficult>
        <type>robndbox</type>
        <robndbox>
            <cx>500.0</cx>
            <cy>300.0</cy>
            <w>100.0</w>
            <h>50.0</h>
            <angle>0.7854</angle> <!-- 弧度制 -->
        </robndbox>
    </object>
</annotation>

2. DOTA 格式

DOTA 是一种基于文本的标注格式，通常用于遥感图像目标检测。它基于四点坐标表示旋转框，每一行表示一个目标的标注信息。
数据示例（TXT 文件）：

x1 y1 x2 y2 x3 y3 x4 y4 class difficult
300.0 250.0 350.0 250.0 350.0 300.0 300.0 300.0 car 0

x1, y1, x2, y2, x3, y3, x4, y4：旋转框四个顶点的坐标。
class：目标类别。
difficult：是否为困难样本。

格式转换原理

1. RoLabelImg → DOTA：

根据 robndbox 中的 (cx, cy, w, h, angle)，计算出旋转框的四个顶点 (x1, y1, x2, y2, x3, y3, x4, y4)。
通过数学公式，将旋转中心点、宽高和旋转角度转化为每个顶点的坐标。

2. DOTA → RoLabelImg：

根据四个顶点 (x1, y1, …, x4, y4)，计算旋转框的中心点 (cx, cy)、宽高 (w, h) 和旋转角度 (angle)。
将这些参数写入 XML 文件。

代码实现

1. RoLabelImg → DOTA 格式

import os
import math
import xml.etree.ElementTree as ET
from typing import List, Tuple
from tqdm import tqdm


def rotate_point(xc: float, yc: float, xp: float, yp: float, theta: float) -> Tuple[float, float]:
    """旋转坐标点

    Args:
        xc: 中心点 x 坐标
        yc: 中心点 y 坐标
        xp: 顶点 x 坐标
        yp: 顶点 y 坐标
        theta: 旋转角度（弧度制）

    Returns:
        旋转后的点坐标 (x', y')
    """
    xoff = xp - xc
    yoff = yp - yc
    cos_theta = math.cos(theta)
    sin_theta = math.sin(theta)
    p_resx = cos_theta * xoff + sin_theta * yoff
    p_resy = -sin_theta * xoff + cos_theta * yoff
    return float(format(xc + p_resx, '.1f')), float(format(yc + p_resy, '.1f'))


def find_top_left_point(points: List[Tuple[float, float]]) -> List[Tuple[float, float]]:
    """找到左上角点并重新排序顶点

    Args:
        points: 四个顶点的坐标列表 [(x1, y1), (x2, y2), ...]

    Returns:
        按顺时针排列的顶点列表
    """
    points_dict = {point[1]: point[0] for point in points}
    sorted_keys = sorted(points_dict.keys())  # 按 y 坐标排序
    temp = [points_dict[sorted_keys[0]], points_dict[sorted_keys[1]]]  # 找到最小的两个 y 坐标对应的 x 值
    minx = min(temp)
    miny = sorted_keys[0] if minx == temp[0] else sorted_keys[1]
    top_left = [minx, miny]  # 左上角点

    # 找到左上角点的索引位置并重新排列
    top_left_index = points.index(tuple(top_left))
    return points[top_left_index:] + points[:top_left_index]


def rolabelimg2dota(xml_path: str, txt_path: str) -> None:
    """将 RoLabelImg 格式 XML 转换为 DOTA 格式 TXT

    Args:
        xml_path: RoLabelImg 格式 XML 文件所在的文件夹路径
        txt_path: 转换后的 DOTA 格式 TXT 文件保存路径
    """
    os.makedirs(txt_path, exist_ok=True)

    file_list = [f for f in os.listdir(xml_path) if f.lower().endswith('.xml')]
    for xml_name in tqdm(file_list, desc="Converting RoLabelImg to DOTA"):
        txt_name = xml_name[:-4] + '.txt'
        txt_file = os.path.join(txt_path, txt_name)
        tree = ET.parse(os.path.join(xml_path, xml_name))
        root = tree.getroot()

        with open(txt_file, "w+", encoding='UTF-8') as out_file:
            for obj in root.findall('object'):
                name = obj.find('name').text
                difficult = obj.find('difficult').text

                robndbox = obj.find('robndbox')
                cx, cy, w, h, angle = [float(robndbox.find(tag).text) for tag in ['cx', 'cy', 'w', 'h', 'angle']]

                # 计算旋转框四个顶点
                points = [
                    rotate_point(cx, cy, cx + w / 2 * sign_x, cy + h / 2 * sign_y, -angle)
                    for sign_x, sign_y in [(-1, -1), (1, -1), (1, 1), (-1, 1)]
                ]

                # 找到左上角点并重新排序顶点
                points = find_top_left_point(points)

                # 写入 DOTA 格式
                data = " ".join(map(str, [*sum(points, ()), name, difficult])) + "\n"
                out_file.write(data)

2. DOTA → RoLabelImg 格式

import os
import math
from tqdm import tqdm
from typing import Dict, List
from xml.etree.ElementTree import Element, SubElement, tostring
from xml.dom.minidom import parseString
from PIL import Image


def create_xml(file_name: str, width: int, height: int, difficult: int, objects: List[Dict]) -> bytes:
    """创建 RoLabelImg 格式的 XML 文件

    Args:
        file_name: 图像文件名
        width: 图像宽度
        height: 图像高度
        difficult: 是否为困难样本（0 或 1）
        objects: 标注对象列表，每个对象包含 class, cx, cy, w, h, angle

    Returns:
        XML 文件内容的 bytes 数据
    """
    annotation = Element("annotation")
    folder = SubElement(annotation, "folder")
    folder.text = "Unknown"
    filename = SubElement(annotation, "filename")
    filename.text = file_name

    size = SubElement(annotation, "size")
    SubElement(size, "width").text = str(width)
    SubElement(size, "height").text = str(height)
    SubElement(size, "depth").text = "3"

    for obj in objects:
        object_node = SubElement(annotation, "object")
        SubElement(object_node, "name").text = obj['class']
        SubElement(object_node, "difficult").text = str(difficult)
        robndbox = SubElement(object_node, "robndbox")
        for tag, value in obj.items():
            if tag != 'class':
                SubElement(robndbox, tag).text = str(value)

    xml_str = parseString(tostring(annotation)).toprettyxml(indent='\t', encoding='utf-8')
    return xml_str


def dota_to_rolabelimg(input_folder: str, output_folder: str, image_folder: str) -> None:
    """将 DOTA 格式 TXT 转换为 RoLabelImg 格式 XML

    Args:
        input_folder: DOTA 格式 TXT 文件所在的文件夹路径
        output_folder: 转换后的 RoLabelImg XML 文件保存路径
        image_folder: 图像文件夹路径，用于获取图像宽高信息
    """
    os.makedirs(output_folder, exist_ok=True)

    for file_name in tqdm(os.listdir(input_folder), desc="Converting DOTA to RoLabelImg"):
        if file_name.endswith(".txt"):
            input_file_path = os.path.join(input_folder, file_name)
            output_file_name = file_name.replace(".txt", ".xml")
            output_path = os.path.join(output_folder, output_file_name)
            image_path = os.path.join(image_folder, file_name.replace(".txt", ".jpg"))

            # 获取图像尺寸
            img = Image.open(image_path)
            image_width, image_height = img.size

            objects = []
            with open(input_file_path, "r") as input_file:
                for line in input_file:
                    data = line.strip().split(" ")
                    x1, y1, x2, y2, x3, y3, x4, y4, obj_class, difficult = data
                    x1, y1, x2, y2, x3, y3, x4, y4 = map(float, [x1, y1, x2, y2, x3, y3, x4, y4])

                    # 计算旋转框参数
                    cx = (x1 + x3) / 2
                    cy = (y1 + y3) / 2
                    w = math.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)
                    h = math.sqrt((x2 - x3) ** 2 + (y2 - y3) ** 2)
                    angle = math.atan2(y2 - y1, x2 - x1)

                    objects.append({"class": obj_class, "cx": cx, "cy": cy, "w": w, "h": h, "angle": angle})

            # 创建 XML 文件内容
            xml_str = create_xml(file_name, image_width, image_height, int(difficult), objects)

            # 写入 XML 文件
            with open(output_path, "wb") as output_file:
                output_file.write(xml_str)