Conversion of Ground Truth data for the USC dataset

usc_parser.py

import sys
import glob
import os
import argparse
import xml.etree.ElementTree

# dlib output format: <image file='*.jpg'><box top='' left='' width='' height=''/> </image>
# USC format: <Object><Rect x="142" y="22" width="28" height="73"/></Object>


class Box:
    top = 0
    left = 0
    width = 0
    height = 0


class ImageNode:
    filename = ''
    elements = None  # list cannot be initialized here!


ap = argparse.ArgumentParser()
ap.add_argument("-p", "--path", required=True,
                help="path to USC XML directory")
ap.add_argument("-o", "--output", required=True,
                help="path for the converted output file")
args = vars(ap.parse_args())


def writeXML(imageElements, outputFile):
    rootNode = xml.etree.ElementTree.Element('dataset')
    nameNode = xml.etree.ElementTree.SubElement(
        rootNode, 'name')
    nameNode.text = "Pedestrians"
    commentNode = xml.etree.ElementTree.SubElement(
        rootNode, 'comment')
    commentNode.text = "USC Pedestrian Dataset."
    imagesNode = xml.etree.ElementTree.SubElement(
        rootNode, 'images')

    for image in imageElements:
        imageNode = xml.etree.ElementTree.SubElement(
            imagesNode, 'image', file=image.filename)
        for box in image.elements:
            xml.etree.ElementTree.SubElement(
                imageNode, 'box', top=box.top, left=box.left, width=box.width, height=box.height)

    outputXMLTree = xml.etree.ElementTree.ElementTree(rootNode)
    outputXMLTree.write(
        outputFile, encoding="utf-8", xml_declaration=True)


def processFile(file):
    imagePath = ''
    # Validate file extension
    if file.endswith('.gt.xml'):
        imagePath = "%s.bmp" % file[:-7]
    elif file.endswith('.xml'):
        imagePath = "%s.bmp" % file[:-4]
    else:
        raise ValueError('Invalid file format - it must be a XML file')

    # Convert Elements
    node = ImageNode()
    node.filename = imagePath
    node.elements = []
    root = xml.etree.ElementTree.parse(file).getroot()
    for object in root.iter('Rect'):
        box = Box()
        box.top = object.attrib['x']
        box.left = object.attrib['y']
        box.width = object.attrib['width']
        box.height = object.attrib['height']
        node.elements.append(box)

    return node


if __name__ == '__main__':

    xml_path = args['path']

    print("Processing files in", xml_path)

    imageNodes = []
    for f in glob.glob(os.path.join(xml_path, "*.xml")):
        print("Parsing %s..." % f, end="")
        imageNodes.append(processFile(f))
        print("done")

    print("Generating XML...", end="")
    writeXML(imageNodes, args['output'])
    print("done")