diff --git a/CHANGELOG.md b/CHANGELOG.md index 4d5e55c25129..9781e527e86b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -23,6 +23,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Default `label-map` parameter value for VOC converter () - Randomness of random split transform () - `Transform.subsets()` method () +- Supported unknown image formats in TF Detection API converter () ### Security - diff --git a/datumaro/plugins/tf_detection_api_format/converter.py b/datumaro/plugins/tf_detection_api_format/converter.py index cdc1e5fe605e..382149dc3449 100644 --- a/datumaro/plugins/tf_detection_api_format/converter.py +++ b/datumaro/plugins/tf_detection_api_format/converter.py @@ -199,7 +199,7 @@ def _make_tf_example(self, item): def _save_image(self, item, path=None): src_ext = item.image.ext.lower() dst_ext = osp.splitext(osp.basename(path))[1].lower() - fmt = DetectionApiPath.IMAGE_EXT_FORMAT.get(dst_ext) + fmt = DetectionApiPath.IMAGE_EXT_FORMAT.get(dst_ext, '') if not fmt: log.warning("Item '%s': can't find format string for the '%s' " "image extension, the corresponding field will be empty." % \ diff --git a/datumaro/plugins/tf_detection_api_format/format.py b/datumaro/plugins/tf_detection_api_format/format.py index f4a879a692a2..e0999943f334 100644 --- a/datumaro/plugins/tf_detection_api_format/format.py +++ b/datumaro/plugins/tf_detection_api_format/format.py @@ -8,6 +8,6 @@ class DetectionApiPath: ANNOTATIONS_DIR = 'annotations' DEFAULT_IMAGE_EXT = '.jpg' - IMAGE_EXT_FORMAT = {'.jpg': 'jpeg', '.png': 'png'} + IMAGE_EXT_FORMAT = {'.jpg': 'jpeg', '.jpeg': 'jpeg', '.png': 'png'} LABELMAP_FILE = 'label_map.pbtxt' \ No newline at end of file diff --git a/tests/test_tfrecord_format.py b/tests/test_tfrecord_format.py index 4c3a78bbff3b..f9491c39effa 100644 --- a/tests/test_tfrecord_format.py +++ b/tests/test_tfrecord_format.py @@ -9,7 +9,7 @@ AnnotationType, Bbox, Mask, LabelCategories ) from datumaro.components.project import Project -from datumaro.util.image import Image +from datumaro.util.image import Image, ByteImage, encode_image from datumaro.util.test_utils import (TestDir, compare_datasets, test_save_and_load) from datumaro.util.tf_util import check_import @@ -135,6 +135,25 @@ def test_can_save_dataset_with_image_info(self): self._test_save_and_load(test_dataset, TfDetectionApiConverter.convert, test_dir) + def test_can_save_dataset_with_unknown_image_formats(self): + test_dataset = Dataset.from_iterable([ + DatasetItem(id=1, + image=ByteImage(data=encode_image(np.ones((5, 4, 3)), 'png'), + path='1/q.e'), + attributes={'source_id': ''} + ), + DatasetItem(id=2, + image=ByteImage(data=encode_image(np.ones((6, 4, 3)), 'png'), + ext='qwe'), + attributes={'source_id': ''} + ) + ], categories={ AnnotationType.label: LabelCategories(), }) + + with TestDir() as test_dir: + self._test_save_and_load(test_dataset, + partial(TfDetectionApiConverter.convert, save_images=True), + test_dir) + def test_labelmap_parsing(self): text = """ {