create pascol voc xml from csv

Question

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

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

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

如果我想再添加一个类(比如class==cross_load(，那么xml将包括traffic_sign和基于csv信息的cross_load。

我在这里举了一个例子：

我只想要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

我猜您正试图将所有条目组合到同一个文件名中，例如您的示例CSVMP_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()跳过标题行。