从csv创建pascol voc xml

Question

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

我想在类的基础上单独制作xml文件。

例如:如果我想要class == traffic_sign，那么每个xml (对于每个图像)将只包含traffic_sign，相应的information.You可以参见下面的示例：

如果我想再添加一个类(比方说类== 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中需要两个条目。

要实现这一点，您可以使用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()跳过标题行。