James
James

Reputation: 127

How to change code to read and create multiple xml files?

I have this code that I am using to change the images I have saved in a folder called 'images' from .png to .xml with the additional information about them. When I run this code I only get the .xml file for image 000001 which I understand because I am having the code select that specific image. I am unsure how though to select multiple images in my file at a single time. I have images named from 000000 to 000355. Any advice would be great! really do not want to manually run the code 355 times!

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)

Upvotes: 1

Views: 389

Answers (1)

Scott Hunter
Scott Hunter

Reputation: 49803

The basic idea is to make loop going through each of your image files, and do what you did for a single image before for each:

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

(You might need to change the expression for your list of images, as it might now include things you don't want to process.)

Upvotes: 1

Related Questions