【数据转换】对mask图像数据转为xml标签数据

关键词： python 、mask 、 xml 、opencv 、IRSTD-1K、NUDT-SIRST

背景

Mask图像数据通常是通过某种分割算法或工具生成的，其中每个像素都被赋予了一个标签，表示其所属的目标或区域。最近我图省事情，在网上找了一些红外小目标的开源数据集，但是自己的任务是做目标检测，所以需要对mask图像数据转为xml格式数据。

xml与mask分析

XML格式的标注数据内容在深度学习、计算机视觉等领域中扮演着重要的角色。在图像标注领域，XML文件通常用于存储图像中各个目标或区域的标签信息。 XML格式的标注数据内容通常包含以下几个关键部分：

文件头部
图像信息
目标或区域列表
- 目标或区域属性
其他辅助信息

样例：

shell 复制代码

<annotation>
    <folder>aa</folder>
    <filename>XDU1.png</filename>
    <path>C:/Users/kiven/Desktop/aa/XDU1.png</path>
    <source>
       <database>Unknown</database>
    </source>
    <size>
       <width>512</width>
       <height>512</height>
       <depth>1</depth>
    <segmented>0</segmented>
    <object>
       <name>1</name>
       <pose>Unspecified</pose>
       <truncated>0</truncated>
       <difficult>0</difficult>
       <bndbox>
          <xmin>113</xmin>
          <ymin>251</ymin>
          <xmax>125</xmax>
          <ymax>265</ymax>
       </bndbox>
    </object>
    <object>
       <name>1</name>
       <pose>Unspecified</pose>
       <truncated>0</truncated>
       <difficult>0</difficult>
       <bndbox>
          <xmin>107</xmin>
          <ymin>290</ymin>
          <xmax>110</xmax>
          <ymax>296</ymax>
       </bndbox>
    </object>
</annotation>

对于IRSTD-1K和NUDT-SIRST数据集在这里查看后得到：mask图像为二值图，黑底白色表示目标区域。

转换

通过分析mask图像我们可以很清晰得到目标区域为白色区域（二值图），因此我们可以利用opencv中的 cv2.findContours函数得到白色目标，然后计算外接最小矩形即可得到目标区域坐标。然后根据xml信息进行逐步填充信息进行调试即可：

shell 复制代码

import os
import cv2

def put_xml(xml_name, img_name, img_dir, width, height, lab_objects):
    # write in xml file
    xml_file = open("C:/Users/kiven/Desktop/VOC2007/" + xml_name, 'w')
    xml_file.write('<annotation>\n')
    xml_file.write('    <folder>VOC2007</folder>\n')
    xml_file.write('    <filename>' + img_name + '</filename>\n')
    xml_file.write('    <path>' + img_dir + '</path>\n')
    xml_file.write('    <size>\n')
    xml_file.write('        <width>' + str(width) + '</width>\n')
    xml_file.write('        <height>' + str(height) + '</height>\n')
    xml_file.write('        <depth>3</depth>\n')
    xml_file.write('    </size>\n')

    for lab_object in lab_objects:

        xml_file.write('    <object>\n')
        xml_file.write('        <name>' + "0" + '</name>\n')
        xml_file.write('        <pose>Unspecified</pose>\n')
        xml_file.write('        <truncated>0</truncated>\n')
        xml_file.write('        <difficult>0</difficult>\n')
        xml_file.write('        <bndbox>\n')
        xml_file.write('            <xmin>' + str(lab_object[0]) + '</xmin>\n')
        xml_file.write('            <ymin>' + str(lab_object[1]) + '</ymin>\n')
        xml_file.write('            <xmax>' + str(lab_object[2]) + '</xmax>\n')
        xml_file.write('            <ymax>' + str(lab_object[3]) + '</ymax>\n')
        xml_file.write('        </bndbox>\n')
        xml_file.write('    </object>\n')

    xml_file.write('</annotation>')
    xml_file.close()


def make_xml(img_path, img_name):
    img_dir = img_path + "/" + img_name
    # 读取图像
    image = cv2.imread(img_dir, cv2.IMREAD_GRAYSCALE)
    height, width = image.shape[:2]
    # 查找所有的连通区域
    contours, _ = cv2.findContours(image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    # 遍历所有的连通区域
    boxx = []
    for contour in contours:
        # 计算每个连通区域的最小外接矩形
        x, y, w, h = cv2.boundingRect(contour)
        # 打印最小外接矩形的坐标和大小
        print(f"白色区域的最小外接矩形坐标: ({x}, {y})，宽度: {w}，高度: {h}")
        boxx.append([x, y, x + w, y + h])
    return img_dir, boxx, height, width


if __name__ == '__main__':
    img_path = "C:/Users/kiven/Desktop/VOC2007"
    for img_name in os.listdir(img_path):
        img_dir, boxx, height, width = make_xml(img_path, img_name)
        xml_name = img_name[:-4] + ".xml"
        put_xml(xml_name, img_name, img_dir, width, height, boxx)

结语

行！有了这个代码生成是xml标注信息，后续的实验数据又新增了不少