如何更改代码以读取和创建多个xml文件?

时间:2018-10-15 13:38:07

标签: python

我有这段代码可用来将保存在名为“ images”的文件夹中的图像从.png更改为.xml,并提供有关它们的附加信息。当我运行此代码时,我只得到图像000001的.xml文件,我理解该文件,因为我有代码选择该特定图像。我不确定如何一次在文件中选择多个图像。我具有从000000到000355的图像。任何建议都非常好!真的不想手动运行该代码355次!

import os
import cv2
from lxml import etree
import xml.etree.cElementTree as ET


def write_xml(folder, img, objects, tl, br, savedir):
    if not os.path.isdir(savedir):
        os.mkdir(savedir)

    image = cv2.imread(img.path)
    height, width, depth = image.shape

    annotation = ET.Element('annotation')
    ET.SubElement(annotation, 'folder').text = folder
    ET.SubElement(annotation, 'filename').text = img.name
    ET.SubElement(annotation, 'segmented').text = '0'
    size = ET.SubElement(annotation, 'size')
    ET.SubElement(size, 'width').text = str(width)
    ET.SubElement(size, 'height').text = str(height)
    ET.SubElement(size, 'depth').text = str(depth)
    for obj, topl, botr in zip(objects, tl, br):
        ob = ET.SubElement(annotation, 'object')
        ET.SubElement(ob, 'name').text = obj
        ET.SubElement(ob, 'pose').text = 'Unspecified'
        ET.SubElement(ob, 'truncated').text = '0'
        ET.SubElement(ob, 'difficult').text = '0'
        bbox = ET.SubElement(ob, 'bndbox')
        ET.SubElement(bbox, 'xmin').text = str(topl[0])
        ET.SubElement(bbox, 'ymin').text = str(topl[1])
        ET.SubElement(bbox, 'xmax').text = str(botr[0])
        ET.SubElement(bbox, 'ymax').text = str(botr[1])

    xml_str = ET.tostring(annotation)
    root = etree.fromstring(xml_str)
    xml_str = etree.tostring(root, pretty_print=True)
    save_path = os.path.join(savedir, img.name.replace('png', 'xml'))
    with open(save_path, 'wb') as temp_xml:
        temp_xml.write(xml_str)


if __name__ == '__main__':
    """
    for testing
    """

    folder = 'images'
    img = [im for im in os.scandir('images') if '000001' in im.name][0]
    objects = ['auv']
    tl = [(10, 10)]
    br = [(100, 100)]
    savedir = 'annotations'
    write_xml(folder, img, objects, tl, br, savedir)

1 个答案:

答案 0 :(得分:1)

基本思想是使循环遍历每个图像文件,然后对每个图像文件执行相同的操作:

for img in os.scandir('images'):
    objects = ['auv']
    tl = [(10, 10)]
    br = [(100, 100)]
    savedir = 'annotations'
    write_xml(folder, img, objects, tl, br, savedir)

(您可能需要更改图像列表的表达式,因为它现在可能包含您不想处理的内容。)