diff --git a/cvat/utils/convert_to_voc.py b/cvat/utils/convert_to_voc.py index cd9cc4c4e27..8d846a60528 100644 --- a/cvat/utils/convert_to_voc.py +++ b/cvat/utils/convert_to_voc.py @@ -1,8 +1,25 @@ +""" +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. + +To use the script run: + +python convert_to_voc.py cvat.xml path_to_image_directory output_directory +""" import os import argparse import xml.etree.ElementTree from PIL import Image from pascal_voc_writer import Writer +import logging + +logger = logging.getLogger() +KNOWN_TAGS = {'box', 'image', 'attribute'} def process_cvat_xml(xml_file, img_dir, annotation_dir): @@ -18,6 +35,12 @@ def process_cvat_xml(xml_file, img_dir, annotation_dir): os.makedirs(annotation_dir) cvat_xml = xml.etree.ElementTree.parse(xml_file) + tracks = [(x.get('id'), x.get('label')) + for x in cvat_xml.findall('track')] + if tracks: + logger.warn('Cannot parse interpolation tracks, ignoring {} tracks' + .format(len(tracks))) + for img_tag in cvat_xml.findall('image'): filename = img_tag.get('name') @@ -27,6 +50,11 @@ def process_cvat_xml(xml_file, img_dir, annotation_dir): writer = Writer(filepath, width, height) + unknown_tags = {x.tag for x in img_tag.iter()}.difference(KNOWN_TAGS) + if unknown_tags: + logger.warn('Ignoring tags for image {}: {}' + .format(filepath, unknown_tags)) + for box in img_tag.findall('box'): label = box.get('label') xmin = float(box.get('xtl')) diff --git a/cvat/utils/tests/__init__.py b/cvat/utils/tests/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/cvat/utils/tests/test_process_cvat_xml.py b/cvat/utils/tests/test_process_cvat_xml.py new file mode 100644 index 00000000000..d1dc5c5c1eb --- /dev/null +++ b/cvat/utils/tests/test_process_cvat_xml.py @@ -0,0 +1,179 @@ +import tempfile +import shutil +import os +from unittest import TestCase, mock +from cvat.utils.convert_to_voc import process_cvat_xml + +XML_ANNOTATION_EXAMPLE = """ + + 1.0 + + + 1063 + My annotation task + 75 + annotation + 0 + + 2018-06-06 11:57:54.807162+03:00 + 2018-06-06 12:42:29.375251+03:00 + + + + + + 3086 + 0 + 74 + http://cvat.examle.com:8080/?id=3086 + + + + admin + + + + 2018-06-06 15:47:04.386866+03:00 + + + + false + a + + + + + true + a + + + + + false + b + + + + + false + c + + + true + a + + + +""" +XML_INTERPOLATION_EXAMPLE = """ + + 1.0 + + + 1062 + My interpolation task + 30084 + interpolation + 20 + + 2018-05-31 14:13:36.483219+03:00 + 2018-06-06 13:56:32.113705+03:00 + + + + + + 3085 + 0 + 30083 + http://cvat.example.com:8080/?id=3085 + + + + admin + + + + 2018-06-06 15:52:11.138470+03:00 + + + + 1 + + + 1 + + + 1 + + + + + 3 + + + 3 + + + 3 + + + +""" + + +class TestProcessCvatXml(TestCase): + def setUp(self): + self.test_dir = tempfile.mkdtemp() + + def tearDown(self): + shutil.rmtree(self.test_dir) + + @mock.patch('cvat.utils.convert_to_voc.logger') + @mock.patch('cvat.utils.convert_to_voc.Image') + def test_parse_annotation_xml(self, mock_image, mock_logger): + xml_filename = os.path.join(self.test_dir, 'annotations.xml') + with open(xml_filename, mode='x') as file: + file.write(XML_ANNOTATION_EXAMPLE) + + voc_dir = os.path.join(self.test_dir, 'voc_dir') + + width, height = 600, 400 + img = mock.MagicMock() + img.size = width, height + mock_image.open.return_value.__enter__.return_value = img + + images = ['C15_L1_0001', 'C15_L1_0002', 'C15_L1_0003', 'C15_L1_0040'] + expected_xmls = [os.path.join(voc_dir, x + '.xml') + for x in images] + expected_warn = "Ignoring tags for image img_dir/C15_L1_0040.jpg: " \ + "{'point'}" + process_cvat_xml(xml_filename, 'img_dir', voc_dir) + for exp in expected_xmls: + self.assertTrue(os.path.exists(exp)) + mock_logger.warn.assert_called_once_with(expected_warn) + + @mock.patch('cvat.utils.convert_to_voc.logger') + def test_parse_interpolation_xml(self, mock_logger): + xml_filename = os.path.join(self.test_dir, 'interpolations.xml') + with open(xml_filename, mode='x') as file: + file.write(XML_INTERPOLATION_EXAMPLE) + + voc_dir = os.path.join(self.test_dir, 'voc_dir') + expected_warn = 'Cannot parse interpolation tracks, ignoring 2 tracks' + + process_cvat_xml(xml_filename, 'img_dir', voc_dir) + + self.assertTrue(os.path.exists(voc_dir)) + self.assertTrue(len(os.listdir(voc_dir)) == 0) + mock_logger.warn.assert_called_once_with(expected_warn)