You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
100 lines
3.1 KiB
Python
100 lines
3.1 KiB
Python
#!/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)
|
|
|
|
tracks = [(x.get('id'), x.get('label'))
|
|
for x in cvat_xml.findall('track')]
|
|
if tracks:
|
|
log.warn('Cannot parse interpolation tracks, ignoring {} tracks'.format(len(tracks)))
|
|
|
|
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()
|