DOTA数据与XML文件之间的相互转化

在目标检测任务中，当我们利用roLabelImg项目标注好了数据之后，会发现数据的保存格式是XML格式，并不能直接用于目标检测的算法中。因此，我们需要将XML格式的文件转换为适用于目标检测的DOTA格式（txt文件），并能够实现从DOTA格式回转到XML格式。下面我将介绍如何实现这两种格式的转换。

1. XML到DOTA的转换

此部分代码主要将标注工具生成的XML文件转换为DOTA格式的文本文件。每个XML文件通常包含一张图像的多个对象标注，包括对象的类别、位置和难度等级。

输入:

• xml_dir: 包含XML文件的目录路径。
• xml_name: XML文件的名称。
• img_dir: 原始图像文件的目录路径。
• savedImg_dir: 转换后图像保存的目录路径。

输出:

• txt_path: 生成的DOTA格式文本文件的保存路径。

# -*- coding: utf8 -*-

"""
This code is used to convert from xml files to dota files
"""


import os
import xml.etree.ElementTree as ET
import math
import cv2 as cv
import argparse
import numpy as np

def voc_to_dota(xml_dir, xml_name, img_dir, savedImg_dir):
    txt_name = xml_name[:-4] + '.txt'  # txt文件名字：去掉xml 加上.txt
    txt_path = xml_dir + '/txt_label_v1'  # txt文件目录：在xml目录下创建的txtl_label文件夹
    if not os.path.exists(txt_path):
        os.makedirs(txt_path)
    txt_file = os.path.join(txt_path, txt_name)  # txt完整的含名文件路径

    xml_file = os.path.join(xml_dir, xml_name)
    tree = ET.parse(os.path.join(xml_file))  # 解析xml文件 然后转换为DOTA格式文件
    root = tree.getroot()
    # img_name = xml_name[:-4] + '.jpg'
    # img_path = os.path.join(img_dir, img_name)
    # img = cv.imdecode(np.fromfile(img_path, dtype=np.uint8), cv.IMREAD_COLOR)
    with open(txt_file, "w+", encoding='UTF-8') as out_file:
        # out_file.write('imagesource:null' + '\n' + 'gsd:null' + '\n')
        for obj in root.findall('object'):
            name = obj.find('name').text
            difficult = obj.find('difficult').text
            try:
                robndbox = obj.find('robndbox')
                cx = float(robndbox.find('cx').text)
                cy = float(robndbox.find('cy').text)
                w = float(robndbox.find('w').text)
                h = float(robndbox.find('h').text)
                angle = float(robndbox.find('angle').text)
                # 其他处理逻辑
            except AttributeError:
                print(f"文件 {xml_name} 中的 {obj} 缺少 'robndbox' 标签")
                # 可以选择跳过这个文件或做其他处理
                continue  # 继续处理下一个object
            p0x, p0y = rotatePoint(cx, cy, cx - w / 2, cy - h / 2, -angle)
            p1x, p1y = rotatePoint(cx, cy, cx + w / 2, cy - h / 2, -angle)
            p2x, p2y = rotatePoint(cx, cy, cx + w / 2, cy + h / 2, -angle)
            p3x, p3y = rotatePoint(cx, cy, cx - w / 2, cy + h / 2, -angle)

            # 找最左上角的点
            dict = {p0y: p0x, p1y: p1x, p2y: p2x, p3y: p3x}
            list = find_topLeftPopint(dict)
            # print((list))
            if list[0] == p0x:
                list_xy = [p0x, p0y, p1x, p1y, p2x, p2y, p3x, p3y]
            elif list[0] == p1x:
                list_xy = [p1x, p1y, p2x, p2y, p3x, p3y, p0x, p0y]
            elif list[0] == p2x:
                list_xy = [p2x, p2y, p3x, p3y, p0x, p0y, p1x, p1y]
            else:
                list_xy = [p3x, p3y, p0x, p0y, p1x, p1y, p2x, p2y]

            # 在原图上画矩形 看是否转换正确

            # cv.line(img, (int(list_xy[0]), int(list_xy[1])), (int(list_xy[2]), int(list_xy[3])), color=(255, 0, 0),
            #         thickness=3)
            # cv.line(img, (int(list_xy[2]), int(list_xy[3])), (int(list_xy[4]), int(list_xy[5])), color=(0, 255, 0),
            #         thickness=3)
            # cv.line(img, (int(list_xy[4]), int(list_xy[5])), (int(list_xy[6]), int(list_xy[7])), color=(0, 0, 255),
            #         thickness=2)
            # cv.line(img, (int(list_xy[6]), int(list_xy[7])), (int(list_xy[0]), int(list_xy[1])), color=(255, 255, 0),
            #         thickness=2)

            data = str(list_xy[0]) + " " + str(list_xy[1]) + " " + str(list_xy[2]) + " " + str(list_xy[3]) + " " + \
                   str(list_xy[4]) + " " + str(list_xy[5]) + " " + str(list_xy[6]) + " " + str(list_xy[7]) + " "
            data = data + name + " " + difficult + "\n"
            out_file.write(data)
        if not os.path.exists(savedImg_dir):
            os.makedirs(savedImg_dir)
        print(xml_name)
        out_img = os.path.join(savedImg_dir, xml_name[:-4] + '.jpg')
        print(out_img)
        # cv.imshow("dam", img)
        # cv.waitKey(0)
        # cv.imwrite(out_img, img)
        # is_saved = cv.imwrite(out_img, img)
        
        # cv.imencode('.jpg', img)[1].tofile(out_img)#英文或中文路径均适用
        # 打印返回值以确认图像是否成功保存
        # print("Image saved:", is_saved)

        # 根据返回值进行条件判断
        # if not is_saved:
        #     print(f"Failed to save image at {out_img}. Check path and image data.")


def find_topLeftPopint(dict):
    dict_keys = sorted(dict.keys())  # y值
    temp = [dict[dict_keys[0]], dict[dict_keys[1]]]
    minx = min(temp)
    if minx == temp[0]:
        miny = dict_keys[0]
    else:
        miny = dict_keys[1]
    return [minx, miny]


# 转换成四点坐标
def rotatePoint(xc, yc, xp, yp, theta):
    xoff = xp - xc
    yoff = yp - yc
    cosTheta = math.cos(theta)
    sinTheta = math.sin(theta)
    pResx = cosTheta * xoff + sinTheta * yoff
    pResy = - sinTheta * xoff + cosTheta * yoff
    # pRes = (xc + pResx, yc + pResy)
    # 保留一位小数点
    return float(format(xc + pResx, '.1f')), float(format(yc + pResy, '.1f'))
    # return xc + pResx, yc + pResy


def parse_args():
    parser = argparse.ArgumentParser(description='数据格式转换')
    parser.add_argument('--xml-dir', default=r'D:\project\代码\train_v4\data_v2\test_contact\txt2xml', help='original xml file dictionary')
    parser.add_argument('--img-dir', default=r'D:\project\代码\data_xml_jpg_v3\YANGTZE\YANGTZE_label_2015', help='original image dictionary')
    parser.add_argument('--outputImg-dir', default=r'D:\project\代码\data_xml_jpg_v3\YANGTZE\YANGTZE_label_2015',
                        help='saved image dictionary after dealing ')

    args = parser.parse_args()
    return args


if __name__ == '__main__':
    args = parse_args()
    xml_path = args.xml_dir
    xmlFile_list = os.listdir(xml_path)
    print(xmlFile_list)
    for i in range(0, len(xmlFile_list)):
        if ('.xml' in xmlFile_list[i]) or ('.XML' in xmlFile_list[i]):
            voc_to_dota(xml_path, xmlFile_list[i], args.img_dir, args.outputImg_dir)
            print('----------------------------------------{}{}----------------------------------------'
                  .format(xmlFile_list[i], ' has Done!'))
        else:
            print(xmlFile_list[i] + ' is not xml file')

2. DOTA到XML的转换

此部分代码主要用于将DOTA格式的文本文件转换回XML格式。通常这种转换用于验证或兼容其他使用XML标注的系统。

输入:

• txt_dir: 包含DOTA文本文件的目录。
• img_dir: 相应的图像文件目录。
• xml_dir: 输出的XML文件存放目录。

输出:

• XML文件，每个文件包含一个图像的标注信息。

import os
import shutil
import xml.etree.ElementTree as ET
import math
from xml.dom import minidom

# 文件夹路径
img_dir = r"D:\修复数据\研究生\project\代码\train_v4\data_v2\test\images" # 图像路径
txt_dir = r'D:\修复数据\研究生\project\代码\train_v4\data_v2\test\labels_cleaned' # dota数据路径
xml_dir = r"D:\修复数据\研究生\project\代码\train_v4\data_v2\test\labels_cleaned\txt2xml" # 输出的xml路径

# 确保XML目录存在
if not os.path.exists(xml_dir):
    os.makedirs(xml_dir)

    
def calculate_angle(x1, y1, x2, y2):
    """
    计算两点形成的直线与水平线的夹角（弧度）。
    """
    angle = math.atan2(float(y2) - float(y1), float(x2) - float(x1))
    # 确保角度为正
    angle = angle % (2 * math.pi)
    return angle

def distance(x1, y1, x2, y2):
    return math.sqrt((float(x2) - float(x1))**2 + (float(y2) - float(y1))**2)
    

def convert_txt_to_xml(txt_file, img_file, xml_file_path):
    tree = ET.ElementTree()
    root = ET.Element("annotation")
    tree._setroot(root)

    folder = ET.SubElement(root, "folder")
    folder.text = "2020_16_12800_large_patch_0.2" # 这里的路径可以随意指定不会影响标注
    
    filename = ET.SubElement(root, "filename")
    filename.text = os.path.basename(img_file)
    
    path = ET.SubElement(root, "path")
    path.text = img_file.replace('/', '\')
    
    source = ET.SubElement(root, "source")
    database = ET.SubElement(source, "database")
    database.text = "Unknown"
    
    size = ET.SubElement(root, "size")
    width = ET.SubElement(size, "width")
    width.text = "12800"
    height = ET.SubElement(size, "height")
    height.text = "12800"
    depth = ET.SubElement(size, "depth")
    depth.text = "3"
    
    segmented = ET.SubElement(root, "segmented")
    segmented.text = "0"
    
    with open(txt_file, 'r') as f:
        # next(f)
        # next(f)
        for line in f:
            parts = line.strip().split(' ')
            obj = ET.SubElement(root, "object")
            
            type_element = ET.SubElement(obj, "type")
            type_element.text = "robndbox"  # 添加类型信息
            
            name = ET.SubElement(obj, "name")
            name.text = parts[8]  # 类名
            
            pose = ET.SubElement(obj, "pose")
            pose.text = "Unspecified"  # 增加姿态信息
            
            truncated = ET.SubElement(obj, "truncated")
            truncated.text = "0"  # 增加截断信息
            
            difficult = ET.SubElement(obj, "difficult")
            difficult.text = "0"  # 增加难度信息

            robndbox = ET.SubElement(obj, "robndbox")
            cx = ET.SubElement(robndbox, "cx")
            cy = ET.SubElement(robndbox, "cy")
            w = ET.SubElement(robndbox, "w")
            h = ET.SubElement(robndbox, "h")
            angle = ET.SubElement(robndbox, "angle")

            # 计算中心点坐标、宽度、高度和旋转角度
            cx.text = str((float(parts[0]) + float(parts[4])) / 2)
            cy.text = str((float(parts[1]) + float(parts[5])) / 2)


            # 计算所有可能的边长
            distances = [
                distance(parts[0], parts[1], parts[2], parts[3]), # p0到p1的距离
                distance(parts[2], parts[3], parts[4], parts[5]), # p1到p2的距离
                distance(parts[4], parts[5], parts[6], parts[7]), # p2到p3的距离
                distance(parts[6], parts[7], parts[0], parts[1]), # p3到p0的距离
            ]

            # 排序边长，最长的是宽度，第二长的是高度
            distances_sorted = sorted(distances, reverse=True)
            w.text = str(distances_sorted[0])
            h.text = str(distances_sorted[2])
            print(distances_sorted[0])
            print(distances_sorted[2])
            # 排序找到最长边对应的点
            dis_dic = [
                (distance(parts[0], parts[1], parts[2], parts[3]), (parts[0], parts[1], parts[2], parts[3])),  # p0到p1的距离
                (distance(parts[2], parts[3], parts[4], parts[5]), (parts[2], parts[3], parts[4], parts[5])),  # p1到p2的距离
                (distance(parts[4], parts[5], parts[6], parts[7]), (parts[4], parts[5], parts[6], parts[7])),  # p2到p3的距离
                (distance(parts[6], parts[7], parts[0], parts[1]), (parts[6], parts[7], parts[0], parts[1])),  # p3到p0的距离
            ]
            longest_edge = max(dis_dic, key=lambda x: x[0])

            # 使用最长边上的点计算角度
            angle.text = str(calculate_angle(longest_edge[1][0], longest_edge[1][1], longest_edge[1][2], longest_edge[1][3]))
            
    rough_string = ET.tostring(root, 'utf-8')
    reparsed = minidom.parseString(rough_string)
    with open(xml_file_path, 'w', encoding='utf-8') as xml_file:
        # 使用toprettyxml方法来美化并添加换行和缩进
        xml_file.write(reparsed.toprettyxml(indent="  "))

for txt_filename in os.listdir(txt_dir):
    if txt_filename.endswith(".txt"):
        base_name = txt_filename[:-4]
        img_file = os.path.join(img_dir, base_name + ".jpg")  # 假设图片文件存在
        xml_file_path = os.path.join(xml_dir, base_name + ".xml")
        txt_file = os.path.join(txt_dir, txt_filename)
        
        convert_txt_to_xml(txt_file, img_file, xml_file_path)

通过以上两个代码，我们就可以实现dota格式的txt标注文件与xml标注文件之间的丝滑切换，希望对大家有帮助！