Fix unknown formats in tf detection api export (cvat-ai#40)

* Support jpeg (along with jpg) * Fix unknown formats
TOsmanov · Oct 13, 2020 · a00a4ec · a00a4ec
1 parent 245c770
commit a00a4ec
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 3 deletions.
diff --git 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 (<https://github.com/openvinotoolkit/datumaro/pull/34>)
 - Randomness of random split transform (<https://github.com/openvinotoolkit/datumaro/pull/38>)
 - `Transform.subsets()` method (<https://github.com/openvinotoolkit/datumaro/pull/38>)
+- Supported unknown image formats in TF Detection API converter (<https://github.com/openvinotoolkit/datumaro/pull/40>)
 
 ### Security
 -

diff --git 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
@@ -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'
diff --git 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 = """
             {