从 csv 创建 pascol voc xml

Question

我被困在这里。 我正在从包含大量图像名称、类和边界框信息的 csv 文件生成注释 xml。

我想根据类分别制作xml文件。

例如：如果我想要 class == traffic_sign ，那么每个 xml（对于每个图像）将只包含 traffic_sign 和相应的信息。您可以查看以下示例：

如果我想再添加一个类（比如说 class== cross_road），那么 xml 将包含基于 csv 信息的 traffic_sign 和 cross_road。

我在这里举了一个例子：

我只想要traffic_sign和相应的信息，然后为图像1.jpg创建xml如下：

<annotation>
    <folder>./5/xmls</folder>
    <filename>1.jpg</filename>
    <path>./1.jpg</path>
    <source>
        <database>Unknown</database>
    </source>
    <size>
        <width>1920</width>
        <height>1080</height>
        <depth>3</depth>
    </size>
    <segmented>0</segmented>
    <object>
        <name>traffic_sign</name>
        <pose>Unspecified</pose>
        <truncated>0</truncated>
        <difficult>0</difficult>
        <bndbox>
            <xmin>886.89</xmin>
            <ymin>517.56</ymin>
            <xmax>931.46</xmax>
            <ymax>562.51</ymax>
        </bndbox>
    </object>
    <object>
        <name>traffic_sign</name>
        <pose>Unspecified</pose>
        <truncated>0</truncated>
        <difficult>0</difficult>
        <bndbox>
            <xmin>783.00</xmin>
            <ymin>432.58</ymin>
            <xmax>835.14</xmax>
            <ymax>485.09</ymax>
        </bndbox>
    </object>
</annotation>

我的csv格式是这样的：

filename           width    height  class            xmin     ymin   xmax   ymax
MP_KSC_021405.jpg   1920    1080    traffic_sign    1316.66 681.51  1358.14 717.92
MP_SEL_083698.jpg   1920    1080    traffic_sign    735.44  117.14  827.45  200.42
MP_SEL_083698.jpg   1920    1080    traffic_sign    733.06  201.22  825.07  289.26
MP_KSC_010168.jpg   1920    1080    traffic_sign    617.39  301.83  697.12  423.66

filename 包含所有图片，class 包含所有类。

我想要为每个图像单独的 xml。 每个 xml 文件都将包含与该图像相关的所有类和边界框。到目前为止，这是我在 Python 中所做的：

import pandas
import numpy
import os
# from tqdm import tqdm_notebook

from xml.etree.ElementTree import parse, Element, SubElement, ElementTree
import xml.etree.ElementTree as ET


#save_root1 = "./traffic_light"
save_root2 = "./test/xmls"


#if not os.path.exists(save_root1):
#    os.mkdir(save_root1)

if not os.path.exists(save_root2):
    os.mkdir(save_root2)

def write_xml(folder, filename, width, height, bbox_list):
    root = Element('annotation')
    SubElement(root, 'folder').text = folder
    SubElement(root, 'filename').text = filename
    SubElement(root, 'path').text = './images' +  filename
    source = SubElement(root, 'source')
    SubElement(source, 'database').text = 'Unknown'

    size = SubElement(root, 'size')
    SubElement(size, 'width').text = str(width)
    SubElement(size, 'height').text = str(height)
    SubElement(size, 'depth').text = '3'

    SubElement(root, 'segmented').text = '0'

    for i in bbox_list:
        obj = SubElement(root, 'object')
        SubElement(obj, 'name').text = i[0]
        SubElement(obj, 'pose').text = 'Unspecified'
        SubElement(obj, 'truncated').text = '0'
        SubElement(obj, 'difficult').text = '0'

        bbox = SubElement(obj, 'bndbox')
        SubElement(bbox, 'xmin').text = str(i[1])
        SubElement(bbox, 'ymin').text = str(i[2])
        SubElement(bbox, 'xmax').text = str(i[3])
        SubElement(bbox, 'ymax').text = str(i[4])

    indent(root)
    tree = ElementTree(root)
    tree.write('./'+folder + '/' + filename.split('.')[0] +'.xml')
    

seed_arr = []

file = open('./test/labels.csv', 'r', encoding='utf-8')
csv_reader = csv.reader(file)
#read csv data
for index, line in enumerate(csv_reader):
    # pass csv header == index[0]
    if index == 0:
        continue
    seed_arr.append(line)
file.close()

sign = 0
print(seed_arr[3])
for index, line in enumerate(seed_arr):
    label = line[3]
    if label == "traffic_sign":
        sign += 1    
        filename = line[0]
        width = line[1]
        height = line[2]
        class_name = line[3]
        xmin = line[4]
        ymin = line[5]
        xmax = line[6]
        ymax = line[7]
        sign+= 1

    if sign > 0:
        write_xml(save_root2, file_nm, width, height, bbox_list)

我希望在这里得到一些帮助。

非常感谢你们。

Answer 1

我猜您正在尝试合并具有相同文件名的所有条目，例如您的示例 CSV MP_SEL_083698.jpg 需要 XML 中的两个条目。

为此，您可以使用defaultdict() 构建一个字典，其中包含所有条目的列表，键为文件名。

然后可以将每个行列表传递给您的 write_xml() 函数以保存它们：

from collections import defaultdict
import os
import csv

from xml.etree.ElementTree import parse, Element, SubElement, ElementTree
import xml.etree.ElementTree as ET

save_root2 = "xmls"

if not os.path.exists(save_root2):
    os.mkdir(save_root2)


def write_xml(folder, filename, bbox_list):
    root = Element('annotation')
    SubElement(root, 'folder').text = folder
    SubElement(root, 'filename').text = filename
    SubElement(root, 'path').text = './images' +  filename
    source = SubElement(root, 'source')
    SubElement(source, 'database').text = 'Unknown'


    # Details from first entry
    e_filename, e_width, e_height, e_class_name, e_xmin, e_ymin, e_xmax, e_ymax = bbox_list[0]
    
    size = SubElement(root, 'size')
    SubElement(size, 'width').text = e_width
    SubElement(size, 'height').text = e_height
    SubElement(size, 'depth').text = '3'

    SubElement(root, 'segmented').text = '0'

    for entry in bbox_list:
        e_filename, e_width, e_height, e_class_name, e_xmin, e_ymin, e_xmax, e_ymax = entry
        
        obj = SubElement(root, 'object')
        SubElement(obj, 'name').text = e_class_name
        SubElement(obj, 'pose').text = 'Unspecified'
        SubElement(obj, 'truncated').text = '0'
        SubElement(obj, 'difficult').text = '0'

        bbox = SubElement(obj, 'bndbox')
        SubElement(bbox, 'xmin').text = e_xmin
        SubElement(bbox, 'ymin').text = e_ymin
        SubElement(bbox, 'xmax').text = e_xmax
        SubElement(bbox, 'ymax').text = e_ymax

    #indent(root)
    tree = ElementTree(root)
    
    xml_filename = os.path.join('.', folder, os.path.splitext(filename)[0] + '.xml')
    tree.write(xml_filename)
    

entries_by_filename = defaultdict(list)

with open('labels.csv', 'r', encoding='utf-8') as f_input_csv:
    csv_input = csv.reader(f_input_csv)
    header = next(csv_input)

    for row in csv_input:
        filename, width, height, class_name, xmin, ymin, xmax, ymax = row

        if class_name == "traffic_sign":
            entries_by_filename[filename].append(row)

for filename, entries in entries_by_filename.items():
    print(filename, len(entries))
    write_xml(save_root2, filename, entries)

您也可以使用next() 跳过标题行。