coco2voc.py

import xml.etree.cElementTree as ET
import pathlib
current_path = pathlib.Path().absolute()


import os.path
import json
def convert_coco_to_voc(coco_path):
    """
        Get all images in coco json file and download them into a folder matching the
        VOC format.
    :param coco_path: Path to coco file
    :return:
    """
    # Let's get the file from the given path
    file = open(coco_path)
    # Now transform it into a python dict.
    coco_data = json.load(file)
    count = 0
    new_images = []
    print(coco_data['images'])
    train_num = int(0.8* len(coco_data['images']))
    for ann in coco_data['annotations']:
        root = ET.Element("annotation")
        folder = ET.SubElement(root, "folder").text = 'train' if count <= train_num else 'test'
        filename = ET.SubElement(root, "filename").text = '{}.jpg'.format(ann['image_id'])
        path = ET.SubElement(root, "path").text = '{}/{}/{}.jpg'.format(current_path, 'JPEGImages', ann['image_id'])
        source = ET.SubElement(root, "source")
        database = ET.SubElement(source, "source").text = 'Unknown'
        size = ET.SubElement(root, "size")
        width = ET.SubElement(size, "width").text = '640'
        height = ET.SubElement(size, "height").text = '360'
        depth = ET.SubElement(size, "depth").text = '3'
        segmented = ET.SubElement(root, "segmented").text = '0'

        object_test = ET.SubElement(root, "object")
        ET.SubElement(object_test, "name").text = "goomba"
        ET.SubElement(object_test, "pose").text = "Unspecified"
        ET.SubElement(object_test, "truncated").text = "0"
        ET.SubElement(object_test, "difficult").text = "0"
        bndbox = ET.SubElement(object_test, "bndbox")
        ET.SubElement(bndbox, "xmin").text = str(ann['bbox'][0])
        ET.SubElement(bndbox, "ymin").text = str(ann['bbox'][1])
        ET.SubElement(bndbox, "xmax").text =  str(ann['bbox'][0] + ann['bbox'][2])
        ET.SubElement(bndbox, "ymax").text = str(ann['bbox'][1] + ann['bbox'][3])

        tree = ET.ElementTree(root)
        tree.write("Annotations/{}.xml".format(ann['image_id']))
        count += 1