Quick hack of cvat converter.py to handle interpolation. (Breaks support for non-interpolated annotation, however.)

converter.py

#!/usr/bin/env python
#
# SPDX-License-Identifier: MIT
"""
Given a CVAT XML and a directory with the image dataset, this script reads the
CVAT XML and writes the annotations in PASCAL VOC format into a given
directory.

This implementation only supports bounding boxes in CVAT annotation format, and
warns if it encounter any tracks or annotations that are not bounding boxes,
ignoring them in both cases.
"""

import os
import argparse
import glog as log
from lxml import etree
from pascal_voc_writer import Writer


def parse_args():
    """Parse arguments of command line"""
    parser = argparse.ArgumentParser(
        description='Convert CVAT XML annotations to PASCAL VOC format'
    )

    parser.add_argument(
        '--cvat-xml', metavar='FILE', required=True,
        help='input file with CVAT annotation in xml format'
    )

    parser.add_argument(
        '--image-dir', metavar='DIRECTORY', required=True,
        help='directory which contains original images'
    )

    parser.add_argument(
        '--output-dir', metavar='DIRECTORY', required=True,
        help='directory for output annotations in PASCAL VOC format'
    )

    return parser.parse_args()


def process_cvat_xml(xml_file, image_dir, output_dir):
    """
    Transforms a single XML in CVAT format to multiple PASCAL VOC format
    XMls.

    :param xml_file: CVAT format XML
    :param image_dir: image directory of the dataset
    :param output_dir: directory of annotations with PASCAL VOC format
    :return:
    """
    KNOWN_TAGS = {'box', 'image', 'attribute'}
    os.makedirs(output_dir, exist_ok=True)
    cvat_xml = etree.parse(xml_file)

    basename = os.path.splitext( os.path.basename( xml_file ) )[0]

    tracks= cvat_xml.findall( './/track' )
    #print( "Tracks: "

    # Build up a list of each bounding box per frame

    frames = {}

    for track in tracks:
        trackid = int(track.get("id"))
        label = track.get("label")
        boxes = track.findall( './box' )
        for box in boxes:
            frameid  = int(box.get('frame'))
            outside  = int(box.get('outside'))
            occluded = int(box.get('occluded'))
            keyframe = int(box.get('keyframe'))
            xtl      = float(box.get('xtl'))
            ytl      = float(box.get('ytl'))
            xbr      = float(box.get('xbr'))
            ybr      = float(box.get('ybr'))
            
            frame = frames.get( frameid, {} )
            
            if outside == 0:
                frame[ trackid ] = { 'xtl': xtl, 'ytl': ytl, 'xbr': xbr, 'ybr': ybr, 'label': label }
                
            frames[ frameid ] = frame

    # Spit out a list of each object for each frame
    for frameid in sorted(frames.keys()):
        print( frameid )

        image_name = "%s_%08d.jpg" % (basename, frameid)
        width = 1270
        height = 780
        image_path = os.path.join(image_dir, image_name)
        if not os.path.exists(image_path):
            log.warn('{} image cannot be found. Is `{}` image directory correct?'.
                format(image_path, image_dir))
        writer = Writer(image_path, width, height)


        frame = frames[frameid]

        objids = sorted(frame.keys())

        for objid in objids:

            obj = frame[objid]
            print( "    %d: (%f,%f) (%f,%f) '%s'" % (objid, obj['xtl'],obj['ytl'],obj['xbr'],obj['ybr'],obj['label']) )


        for objid in objids:

            box = frame[objid]

            label = box.get('label')
            xmin = float(box.get('xtl'))
            ymin = float(box.get('ytl'))
            xmax = float(box.get('xbr'))
            ymax = float(box.get('ybr'))

            writer.addObject(label, xmin, ymin, xmax, ymax)

        anno_name = os.path.basename(os.path.splitext(image_name)[0] + '.xml')
        anno_dir = os.path.dirname(os.path.join(output_dir, image_name))
        os.makedirs(anno_dir, exist_ok=True)
        writer.save(os.path.join(anno_dir, anno_name))


#    for img_tag in cvat_xml.findall('image'):
#        image_name = img_tag.get('name')
#        width = img_tag.get('width')
#        height = img_tag.get('height')
#        image_path = os.path.join(image_dir, image_name)
#        if not os.path.exists(image_path):
#            log.warn('{} image cannot be found. Is `{}` image directory correct?'.
#                format(image_path, image_dir))
#        writer = Writer(image_path, width, height)
#
#        unknown_tags = {x.tag for x in img_tag.iter()}.difference(KNOWN_TAGS)
#        if unknown_tags:
#            log.warn('Ignoring tags for image {}: {}'.format(image_path, unknown_tags))
#
#        for box in img_tag.findall('box'):
#            label = box.get('label')
#            xmin = float(box.get('xtl'))
#            ymin = float(box.get('ytl'))
#            xmax = float(box.get('xbr'))
#            ymax = float(box.get('ybr'))
#
#            writer.addObject(label, xmin, ymin, xmax, ymax)
#
#        anno_name = os.path.basename(os.path.splitext(image_name)[0] + '.xml')
#        anno_dir = os.path.dirname(os.path.join(output_dir, image_name))
#        os.makedirs(anno_dir, exist_ok=True)
#        writer.save(os.path.join(anno_dir, anno_name))


def main():
    args = parse_args()
    process_cvat_xml(args.cvat_xml, args.image_dir, args.output_dir)


if __name__ == "__main__":
    main()