Repository: openvinotoolkit/datumaro

Issue:
Image zip format
- Need to add a format that will allow to export/import zip archive of images.
- Add optional argument with compression method for export.


Discussion:

Resolving patch:
diff --git a/datumaro/plugins/image_zip_format.py b/datumaro/plugins/image_zip_format.py
new file mode 100644
--- /dev/null
+++ b/datumaro/plugins/image_zip_format.py
@@ -0,0 +1,114 @@
+# Copyright (C) 2021 Intel Corporation
+#
+# SPDX-License-Identifier: MIT
+
+from enum import Enum
+import logging as log
+import os
+import os.path as osp
+from zipfile import ZIP_BZIP2, ZIP_DEFLATED, ZIP_LZMA, ZIP_STORED, ZipFile
+
+from datumaro.components.converter import Converter
+from datumaro.components.extractor import DatasetItem, Importer, SourceExtractor
+
+from datumaro.util import parse_str_enum_value
+from datumaro.util.image import IMAGE_EXTENSIONS, ByteImage, encode_image
+
+class Compression(Enum):
+    ZIP_STORED = ZIP_STORED
+    ZIP_DEFLATED = ZIP_DEFLATED
+    ZIP_BZIP2 = ZIP_BZIP2
+    ZIP_LZMA = ZIP_LZMA
+
+class ImageZipPath:
+    DEFAULT_ARCHIVE_NAME = 'default.zip'
+    DEFAULT_COMPRESSION = Compression.ZIP_STORED
+
+class ImageZipExtractor(SourceExtractor):
+    def __init__(self, url, subset=None):
+        super().__init__(subset=subset)
+
+        assert url.endswith('.zip'), url
+
+        with ZipFile(url, 'r') as zf:
+            for path in zf.filelist:
+                item_id, extension = osp.splitext(path.filename)
+                if extension.lower() not in IMAGE_EXTENSIONS:
+                    continue
+                image = ByteImage(data=zf.read(path.filename))
+                self._items.append(DatasetItem(
+                    id=item_id, image=image, subset=self._subset
+                ))
+
+class ImageZipImporter(Importer):
+    @classmethod
+    def find_sources(cls, path):
+        return cls._find_sources_recursive(path, '.zip', 'image_zip')
+
+class ImageZipConverter(Converter):
+    DEFAULT_IMAGE_EXT = '.jpg'
+
+    @staticmethod
+    def _get_compression_method(s):
+        try:
+            return Compression[s.upper()]
+        except KeyError:
+            import argparse
+            raise argparse.ArgumentTypeError()
+
+    @classmethod
+    def build_cmdline_parser(cls, **kwargs):
+        parser = super().build_cmdline_parser(**kwargs)
+
+        parser.add_argument('--name', type=str,
+            default=ImageZipPath.DEFAULT_ARCHIVE_NAME,
+            help="Name of output zipfile (default: %(default)s)"
+        )
+
+        parser.add_argument('--compression', type=cls._get_compression_method,
+            default=ImageZipPath.DEFAULT_COMPRESSION.name,
+            help="Archive compression method.\nAvailable methods: {} "
+                "(default: %(default)s)" \
+                .format(', '.join(e.name for e in Compression))
+        )
+
+        return parser
+
+    def __init__(self, extractor, save_dir, name=None,
+            compression=None, **kwargs):
+        super().__init__(extractor, save_dir, **kwargs)
+
+        if name is None:
+            name = ImageZipPath.DEFAULT_ARCHIVE_NAME
+
+        compression = parse_str_enum_value(compression, Compression,
+            default=ImageZipPath.DEFAULT_COMPRESSION)
+
+        self._archive_name = name
+        self._compression = compression.value
+
+    def apply(self):
+        os.makedirs(self._save_dir, exist_ok=True)
+
+        archive_path = osp.join(self._save_dir, self._archive_name)
+
+        if osp.exists(archive_path):
+            raise FileExistsError('Zip file: %s, already exist,'
+                'specify archive name with --name extra argument' % archive_path)
+
+        with ZipFile(archive_path, 'w', self._compression) as zf:
+            for item in self._extractor:
+                if item.has_image:
+                    self._archive_image(zf, item)
+                else:
+                    log.debug("Item '%s' has no image info", item.id)
+
+    def _archive_image(self, zipfile, item):
+        image_name = self._make_image_filename(item)
+        if osp.isfile(item.image.path):
+            zipfile.write(item.image.path, arcname=image_name)
+        elif isinstance(item.image, ByteImage):
+            zipfile.writestr(image_name, item.image.get_bytes())
+        elif item.image.has_data:
+            zipfile.writestr(image_name,
+                encode_image(item.image.data, osp.splitext(image_name)[1]))


Test patch:
diff --git a/tests/assets/image_zip_dataset/1.zip b/tests/assets/image_zip_dataset/1.zip
new file mode 100644
Binary files /dev/null and b/tests/assets/image_zip_dataset/1.zip differ
diff --git a/tests/assets/image_zip_dataset/2.zip b/tests/assets/image_zip_dataset/2.zip
new file mode 100644
Binary files /dev/null and b/tests/assets/image_zip_dataset/2.zip differ
diff --git a/tests/cli/test_image_zip_format.py b/tests/cli/test_image_zip_format.py
new file mode 100644
--- /dev/null
+++ b/tests/cli/test_image_zip_format.py
@@ -0,0 +1,93 @@
+import os
+import os.path as osp
+from unittest import TestCase
+from zipfile import ZipFile
+
+import numpy as np
+import pytest
+
+from datumaro.cli.__main__ import main
+from datumaro.components.dataset import Dataset, DatasetItem
+from datumaro.util.test_utils import TestDir, compare_datasets
+
+from ..requirements import Requirements
+
+
+def run(test, *args, expected_code=0):
+    test.assertEqual(expected_code, main(args), str(args))
+
+def make_zip_archive(src_path, dst_path):
+    with ZipFile(dst_path, 'w') as archive:
+        for (dirpath, _, filenames) in os.walk(src_path):
+            for name in filenames:
+                path = osp.join(dirpath, name)
+                archive.write(path, osp.relpath(path, src_path))
+
+class ImageZipIntegrationScenarios(TestCase):
+    @pytest.mark.reqids(Requirements.DATUM_267)
+    def test_can_save_and_load(self):
+        source_dataset = Dataset.from_iterable([
+            DatasetItem(id='1', image=np.ones((5, 5, 3))),
+            DatasetItem(id='2', image=np.ones((2, 8, 3)))
+        ])
+
+        with TestDir() as test_dir:
+            source_dataset.export(test_dir, format='image_dir')
+            zip_path = osp.join(test_dir, 'images.zip')
+            make_zip_archive(test_dir, zip_path)
+
+            run(self, 'create', '-o', test_dir)
+            run(self, 'add', 'path', '-p', test_dir, '-f', 'image_zip', zip_path)
+
+            export_path = osp.join(test_dir, 'export.zip')
+            run(self, 'export', '-p', test_dir, '-f', 'image_zip',
+                '-o', test_dir, '--overwrite', '--',
+                '--name', osp.basename(export_path)
+            )
+
+            parsed_dataset = Dataset.import_from(export_path, format='image_zip')
+            compare_datasets(self, source_dataset, parsed_dataset)
+
+    @pytest.mark.reqids(Requirements.DATUM_267)
+    def test_can_export_zip_images_from_coco_dataset(self):
+        with TestDir() as test_dir:
+            coco_dir = osp.join(__file__[:__file__.rfind(osp.join('tests', ''))],
+                'tests', 'assets', 'coco_dataset')
+
+            run(self, 'create', '-o', test_dir)
+            run(self, 'add', 'path', '-p', test_dir, '-f', 'coco', coco_dir)
+
+            export_path = osp.join(test_dir, 'export.zip')
+            run(self, 'export', '-p', test_dir, '-f', 'image_zip',
+                '-o', test_dir, '--overwrite', '--',
+                '--name', osp.basename(export_path))
+
+            self.assertTrue(osp.isfile(export_path))
+            with ZipFile(export_path, 'r') as zf:
+                images = {f.filename for f in zf.filelist}
+                self.assertTrue(images == {'a.jpg', 'b.jpg'})
+
+    @pytest.mark.reqids(Requirements.DATUM_267)
+    def test_can_change_extension_for_images_in_zip(self):
+        source_dataset = Dataset.from_iterable([
+            DatasetItem(id='1', image=np.ones((5, 5, 3))),
+            DatasetItem(id='2', image=np.ones((2, 8, 3)))
+        ])
+
+        with TestDir() as test_dir:
+            source_dataset.export(test_dir, format='image_dir', image_ext='.jpg')
+            zip_path = osp.join(test_dir, 'images.zip')
+            make_zip_archive(test_dir, zip_path)
+
+            run(self, 'create', '-o', test_dir)
+            run(self, 'add', 'path', '-p', test_dir, '-f', 'image_zip', zip_path)
+
+            export_path = osp.join(test_dir, 'export.zip')
+            run(self, 'export', '-p', test_dir, '-f', 'image_zip',
+                '-o', test_dir, '--overwrite', '--',
+                '--name', osp.basename(export_path), '--image-ext', '.png')
+
+            self.assertTrue(osp.isfile(export_path))
+            with ZipFile(export_path, 'r') as zf:
+                images = {f.filename for f in zf.filelist}
+                self.assertTrue(images == {'1.png', '2.png'})
\ No newline at end of file
diff --git a/tests/requirements.py b/tests/requirements.py
--- a/tests/requirements.py
+++ b/tests/requirements.py
@@ -19,6 +19,7 @@ class Requirements:
     # GitHub issues (not bugs)
     # https://github.com/openvinotoolkit/datumaro/issues
     DATUM_244 = "Add Snyk integration"
+    DATUM_267 = "Add Image zip format"
 
     # GitHub issues (bugs)
     # https://github.com/openvinotoolkit/datumaro/issues
diff --git a/tests/test_image_zip_format.py b/tests/test_image_zip_format.py
new file mode 100644
--- /dev/null
+++ b/tests/test_image_zip_format.py
@@ -0,0 +1,109 @@
+import os.path as osp
+from unittest import TestCase
+
+import numpy as np
+import pytest
+
+from datumaro.components.extractor import DatasetItem
+from datumaro.components.project import Dataset
+from datumaro.plugins.image_zip_format import ImageZipPath, ImageZipConverter
+from datumaro.util.image import Image, save_image
+from datumaro.util.test_utils import TestDir, compare_datasets
+
+from .requirements import Requirements
+
+
+class ImageZipConverterTest(TestCase):
+    @pytest.mark.reqids(Requirements.DATUM_267)
+    def _test_can_save_and_load(self, source_dataset, test_dir,
+            **kwargs):
+        archive_path = osp.join(test_dir, kwargs.get('name',
+            ImageZipPath.DEFAULT_ARCHIVE_NAME))
+        ImageZipConverter.convert(source_dataset, test_dir, **kwargs)
+        parsed_dataset = Dataset.import_from(archive_path, 'image_zip')
+
+        compare_datasets(self, source_dataset, parsed_dataset)
+
+    @pytest.mark.reqids(Requirements.DATUM_267)
+    def test_can_save_and_load(self):
+        source_dataset = Dataset.from_iterable([
+            DatasetItem(id='1', image=np.ones((10, 6, 3))),
+            DatasetItem(id='2', image=np.ones((5, 4, 3))),
+        ])
+
+        with TestDir() as test_dir:
+            self._test_can_save_and_load(source_dataset, test_dir)
+
+    @pytest.mark.reqids(Requirements.DATUM_267)
+    def test_can_save_and_load_with_custom_archive_name(self):
+        source_dataset = Dataset.from_iterable([
+            DatasetItem(id='img_1', image=np.ones((10, 10, 3))),
+        ])
+
+        with TestDir() as test_dir:
+            self._test_can_save_and_load(source_dataset, test_dir,
+                name='my_archive.zip')
+
+    @pytest.mark.reqids(Requirements.DATUM_267)
+    def test_relative_paths(self):
+        source_dataset = Dataset.from_iterable([
+            DatasetItem(id='1', image=np.ones((10, 10, 3))),
+            DatasetItem(id='a/2', image=np.ones((4, 5, 3))),
+            DatasetItem(id='a/b/3', image=np.ones((20, 10, 3)))
+        ])
+
+        with TestDir() as test_dir:
+            self._test_can_save_and_load(source_dataset, test_dir)
+
+    @pytest.mark.reqids(Requirements.DATUM_267)
+    def test_can_save_and_load_custom_compresion_method(self):
+        source_dataset = Dataset.from_iterable([
+            DatasetItem(id='1', image=np.ones((5, 5, 3))),
+            DatasetItem(id='2', image=np.ones((4, 3, 3))),
+        ])
+
+        with TestDir() as test_dir:
+            self._test_can_save_and_load(source_dataset, test_dir,
+                compression='ZIP_DEFLATED')
+
+    @pytest.mark.reqids(Requirements.DATUM_267)
+    def test_can_save_and_load_with_arbitrary_extensions(self):
+        source_dataset = Dataset.from_iterable([
+            DatasetItem(id='subset/1',
+                image=Image(data=np.ones((10, 10, 3)), path='subset/1.png')),
+            DatasetItem(id='2',
+                image=Image(data=np.ones((4, 5, 3)), path='2.jpg')),
+        ])
+
+        with TestDir() as test_dir:
+            save_image(osp.join(test_dir, '2.jpg'),
+                source_dataset.get('2').image.data)
+            save_image(osp.join(test_dir, 'subset', '1.png'),
+                source_dataset.get('subset/1').image.data,
+                create_dir=True)
+
+            self._test_can_save_and_load(source_dataset, test_dir)
+
+DUMMY_DATASET_DIR = osp.join(osp.dirname(__file__), 'assets', 'image_zip_dataset')
+
+class ImageZipImporterTest(TestCase):
+    @pytest.mark.reqids(Requirements.DATUM_267)
+    def test_can_import(self):
+        source_dataset = Dataset.from_iterable([
+            DatasetItem(id='1', image=np.ones((10, 10, 3)))
+        ])
+
+        zip_path = osp.join(DUMMY_DATASET_DIR, '1.zip')
+        parsed_dataset = Dataset.import_from(zip_path, format='image_zip')
+        compare_datasets(self, source_dataset, parsed_dataset)
+
+
+    @pytest.mark.reqids(Requirements.DATUM_267)
+    def test_can_import_from_directory(self):
+        source_dataset = Dataset.from_iterable([
+            DatasetItem(id='1', image=np.ones((10, 10, 3))),
+            DatasetItem(id='2', image=np.ones((5, 10, 3)))
+        ])
+
+        parsed_dataset = Dataset.import_from(DUMMY_DATASET_DIR, format='image_zip')
+        compare_datasets(self, source_dataset, parsed_dataset)
\ No newline at end of file