Pascal VOC annotation to Apple createML annotations conversion tool

convert_voc2createml.py

import argparse
import json
import logging
import pathlib
from xml.etree import ElementTree

log = logging.getLogger(__name__)


def _parse_arguments() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description='Pascal VOC to Apple createML annotations converter')

    parser.add_argument('--verbose', '-v', action='store_true',
                        help='Verbose mode')
    parser.add_argument('--source', '-s', default='./', type=str, help='Source path')
    parser.add_argument('--destination', '-d', default='./annotations.json', type=str, help='Destination file')

    return parser.parse_args()


def _parse_annotation(box):
    label = box.find('name').text
    bndbox = box.find('bndbox')
    xmin = float(bndbox.find('xmin').text)
    ymin = float(bndbox.find('ymin').text)
    xmax = float(bndbox.find('xmax').text)
    ymax = float(bndbox.find('ymax').text)

    return {
        'label': label,
        'coordinates': {
            # Bounding box's origin is the center in createML
            # It is the top-left corner in VOC
            'x': (xmin + xmax) / 2.0,
            'y': (ymin + ymax) / 2.0,
            'width': xmax - xmin,
            'height': ymax - ymin,
        },
    }


def _process_single_image(xml_file: pathlib.Path):
    log.info("Processing: %s", xml_file)
    tree = ElementTree.parse(xml_file)
    root = tree.getroot()

    image_filename = root.find('filename').text
    boxes = root.iterfind('object')

    annotations = [_parse_annotation(box) for box in boxes]

    return {
        'imagefilename': image_filename,
        'annotation': annotations,
    }


def main():
    # Parsing script input parameters
    args = _parse_arguments()
    level = logging.DEBUG if args.verbose else logging.INFO
    logging.basicConfig(level=level)

    source_path = pathlib.Path(args.source)
    destination_file = pathlib.Path(args.destination)

    # Build content of the annotations.json
    images_content = [_process_single_image(xml_file) for xml_file in source_path.glob('*.xml')]

    # Save output annotation file
    with open(destination_file, 'w') as file:
        json.dump(images_content, file, indent=4)
    
    log.info("Saved %s annotated images in %s", len(images_content), destination_file)


if __name__ == "__main__":
    main()