diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md index 787a65e9..c649d0b4 100644 --- a/CONTRIBUTORS.md +++ b/CONTRIBUTORS.md @@ -18,5 +18,8 @@ # Contributors - **[Victor Salimonov](https://github.com/VikTorSalimonov)** - * Documentation, screencasts + +- **[Sebastián Yonekura](https://github.com/syonekura)** + * [convert_to_voc.py](cvat/utils/convert_to_voc.py) - an utility for + converting CVAT XML to PASCAL VOC data annotation format. \ No newline at end of file diff --git a/cvat/requirements/base.txt b/cvat/requirements/base.txt index 8f609236..bd216aa4 100644 --- a/cvat/requirements/base.txt +++ b/cvat/requirements/base.txt @@ -21,4 +21,4 @@ rq==0.10.0 scipy==1.0.1 sqlparse==0.2.4 django-sendfile==0.3.11 -dj-pagination==2.3.2 +dj-pagination==2.3.2 \ No newline at end of file diff --git a/cvat/utils/__init__.py b/cvat/utils/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/cvat/utils/convert_to_voc.py b/cvat/utils/convert_to_voc.py new file mode 100644 index 00000000..8d846a60 --- /dev/null +++ b/cvat/utils/convert_to_voc.py @@ -0,0 +1,80 @@ +""" +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): + """ + Transforms a single XML in CVAT format to multiple PASCAL VOC format + XMls. + + :param xml_file: CVAT format XML + :param img_dir: image directory of the dataset + :param annotation_dir: directory of annotations with PASCAL VOC format + :return: + """ + 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') + + filepath = os.path.join(img_dir, filename) + with Image.open(filepath) as img: + width, height = img.size + + 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')) + ymin = float(box.get('ytl')) + xmax = float(box.get('xbr')) + ymax = float(box.get('ybr')) + + writer.addObject(label, xmin, ymin, xmax, ymax) + + fname = os.path.splitext(filename)[0] + '.xml' + writer.save(os.path.join(annotation_dir, fname)) + + +parser = argparse.ArgumentParser(description='Transforms CVAT XML to Pascal ' + 'VOC format') +parser.add_argument('cvat_xml', type=argparse.FileType(), help='CVAT XML file') +parser.add_argument('img_dir', help='Image directory of the dataset') +parser.add_argument('annotation_dir', help='Output directory of ' + 'XML annotations') + +if __name__ == '__main__': + args = vars(parser.parse_args()) + process_cvat_xml(args['cvat_xml'], args['img_dir'], args['annotation_dir']) diff --git a/cvat/utils/requirements.txt b/cvat/utils/requirements.txt new file mode 100644 index 00000000..5528df43 --- /dev/null +++ b/cvat/utils/requirements.txt @@ -0,0 +1 @@ +pascal-voc-writer==0.1.4 \ No newline at end of file diff --git a/cvat/utils/tests/__init__.py b/cvat/utils/tests/__init__.py new file mode 100644 index 00000000..e69de29b 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 00000000..d1dc5c5c --- /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)