Develop DatumaroBinaryFormat to export/import the dataset header & DatasetItem (#828)

- Support DatasetItem by Datumaro binary format but annotations are not yet.
- Quick fix for CI failure because of dill

Signed-off-by: Kim, Vinnam <[email protected]>

Author: vinnamkim

.github/workflows/pr_check.yml

@@ -31,7 +31,7 @@ jobs:
           python-version: ${{ matrix.python-version }}
       - name: Installing dependencies
         run: |
-          pip install -e '.[default,tf,tfds]' pytest pytest-cov
+          pip install -e '.[default,tf,tfds-dev]' pytest pytest-cov
       - name: Unit testing
         run: |
           pytest -v tests/unit/ --cov

datumaro/components/media.py

@@ -8,6 +8,7 @@
 import os.path as osp
 import shutil
 import weakref
+from enum import IntEnum
 from typing import Callable, Iterable, Iterator, List, Optional, Tuple, Union
 
 import cv2
@@ -19,7 +20,22 @@
 BboxIntCoords = Tuple[int, int, int, int]  # (x, y, w, h)
 
 
+class MediaType(IntEnum):
+    NO_MEDIA = 0
+    UNKNOWN = 1
+    IMAGE = 2
+    BYTE_IMAGE = 3
+    VIDEO_FRAME = 4
+    VIDEO = 5
+    POINT_CLOUD = 6
+    MULTIFRAME_IMAGE = 7
+    ROI_IMAGE = 8
+    MOSAIC_IMAGE = 9
+
+
 class MediaElement:
+    MEDIA_TYPE = MediaType.UNKNOWN
+
     def __init__(self, path: str) -> None:
         assert path, "Path can't be empty"
         self._path = path
@@ -42,6 +58,8 @@ def __eq__(self, other: object) -> bool:
 
 
 class Image(MediaElement):
+    MEDIA_TYPE = MediaType.IMAGE
+
     def __init__(
         self,
         data: Union[np.ndarray, Callable[[str], np.ndarray], None] = None,
@@ -169,6 +187,8 @@ def save(self, path):
 
 
 class ByteImage(Image):
+    MEDIA_TYPE = MediaType.BYTE_IMAGE
+
     _FORMAT_MAGICS = (
         (b"\x89PNG\r\n\x1a\n", ".png"),
         (b"\xff\xd8\xff", ".jpg"),
@@ -235,6 +255,8 @@ def save(self, path):
 
 
 class VideoFrame(Image):
+    MEDIA_TYPE = MediaType.VIDEO_FRAME
+
     def __init__(self, video: Video, index: int):
         self._video = video
         self._index = index
@@ -333,6 +355,8 @@ def _navigate_to(self, idx: int) -> VideoFrame:
 
 
 class Video(MediaElement, Iterable[VideoFrame]):
+    MEDIA_TYPE = MediaType.VIDEO
+
     """
     Provides random access to the video frames.
     """
@@ -501,13 +525,17 @@ def __hash__(self):
 
 
 class PointCloud(MediaElement):
+    MEDIA_TYPE = MediaType.POINT_CLOUD
+
     def __init__(self, path: str, extra_images: Optional[List[Image]] = None):
         self._path = path
 
         self.extra_images: List[Image] = extra_images or []
 
 
 class MultiframeImage(MediaElement):
+    MEDIA_TYPE = MediaType.MULTIFRAME_IMAGE
+
     def __init__(
         self,
         images: Optional[Iterable[Union[str, Image, np.ndarray, Callable[[str], np.ndarray]]]],
@@ -538,6 +566,8 @@ def data(self) -> List[Image]:
 
 
 class RoIImage(Image):
+    MEDIA_TYPE = MediaType.ROI_IMAGE
+
     def __init__(
         self,
         data: Union[np.ndarray, Callable[[str], np.ndarray], None] = None,
@@ -580,6 +610,8 @@ def save(self, path):
 
 
 class MosaicImage(Image):
+    MEDIA_TYPE = MediaType.MOSAIC_IMAGE
+
     def __init__(self, imgs: List[ImageWithRoI], size: Tuple[int, int]) -> None:
         def _get_mosaic_img(_) -> np.ndarray:
             h, w = self.size

datumaro/plugins/data_formats/datumaro/base.py

@@ -51,7 +51,10 @@ def __init__(self, path):
         self._related_images_dir = related_images_dir
 
         super().__init__(subset=osp.splitext(osp.basename(path))[0])
+        self._load_impl(path)
 
+    def _load_impl(self, path: str) -> None:
+        """Actual implementation of loading Datumaro format."""
         parsed_anns = parse_json_file(path)
         self._infos = self._load_infos(parsed_anns)
         self._categories = self._load_categories(parsed_anns)

datumaro/plugins/data_formats/datumaro/exporter.py

@@ -342,9 +342,14 @@ def _convert_points_categories(self, obj):
 
 class DatumaroExporter(Exporter):
     DEFAULT_IMAGE_EXT = DatumaroPath.IMAGE_EXT
-    WRITER_CLS = _SubsetWriter
     PATH_CLS = DatumaroPath
 
+    def create_writer(self, subset: str) -> _SubsetWriter:
+        return _SubsetWriter(
+            context=self,
+            ann_file=osp.join(self._annotations_dir, subset + self.PATH_CLS.ANNOTATION_EXT),
+        )
+
     def apply(self):
         os.makedirs(self._save_dir, exist_ok=True)
 
@@ -359,13 +364,8 @@ def apply(self):
         self._pcd_dir = osp.join(self._save_dir, self.PATH_CLS.PCD_DIR)
         self._related_images_dir = osp.join(self._save_dir, self.PATH_CLS.RELATED_IMAGES_DIR)
 
-        writers = {
-            subset: self.WRITER_CLS(
-                context=self,
-                ann_file=osp.join(self._annotations_dir, subset + self.PATH_CLS.ANNOTATION_EXT),
-            )
-            for subset in self._extractor.subsets()
-        }
+        writers = {subset: self.create_writer(subset) for subset in self._extractor.subsets()}
+
         for writer in writers.values():
             writer.add_infos(self._extractor.infos())
             writer.add_categories(self._extractor.categories())

datumaro/plugins/data_formats/datumaro_binary/base.py

@@ -2,8 +2,64 @@
 #
 # SPDX-License-Identifier: MIT
 
+import struct
+from io import BufferedWriter
+from typing import Optional
+
+from datumaro.components.errors import DatasetImportError
+from datumaro.plugins.data_formats.datumaro_binary.format import DatumaroBinaryPath
+from datumaro.plugins.data_formats.datumaro_binary.mapper import DictMapper
+
 from ..datumaro.base import DatumaroBase
+from .crypter import Crypter
 
 
 class DatumaroBinaryBase(DatumaroBase):
     """"""
+
+    def __init__(self, path: str, encryption_key: Optional[bytes] = None):
+        self._fp: Optional[BufferedWriter] = None
+        self._crypter = Crypter(encryption_key)
+        super().__init__(path)
+
+    def _load_impl(self, path: str) -> None:
+        """Actual implementation of loading Datumaro binary format."""
+        try:
+            with open(path, "rb") as fp:
+                self._fp = fp
+                self._check_signature()
+                self._check_encryption_field()
+                self._read_info()
+                self._read_categories()
+        finally:
+            self._fp = None
+
+        return
+
+    def _check_signature(self):
+        signature = self._fp.read(DatumaroBinaryPath.SIGNATURE_LEN).decode()
+        DatumaroBinaryPath.check_signature(signature)
+
+    def _check_encryption_field(self):
+        len_byte = self._fp.read(4)
+        _bytes = self._fp.read(struct.unpack("I", len_byte)[0])
+
+        extracted_key = self._crypter.decrypt(_bytes)
+
+        if not self._crypter.handshake(extracted_key):
+            raise DatasetImportError("Encryption key handshake fails. You give a wrong key.")
+
+    def _read_info(self):
+        len_byte = self._fp.read(4)
+        _bytes = self._fp.read(struct.unpack("I", len_byte)[0])
+        _bytes = self._crypter.decrypt(_bytes)
+
+        self._infos, _ = DictMapper.backward(_bytes)
+
+    def _read_categories(self):
+        len_byte = self._fp.read(4)
+        _bytes = self._fp.read(struct.unpack("I", len_byte)[0])
+        _bytes = self._crypter.decrypt(_bytes)
+
+        categories, _ = DictMapper.backward(_bytes)
+        self._categories = self._load_categories({"categories": categories})

datumaro/plugins/data_formats/datumaro_binary/crypter.py

@@ -0,0 +1,41 @@
+# Copyright (C) 2023 Intel Corporation
+#
+# SPDX-License-Identifier: MIT
+
+from typing import Optional
+
+from cryptography.fernet import Fernet
+
+
+class Crypter:
+    FERNET_KEY_LEN = 44
+
+    def __init__(self, key: Optional[bytes]) -> None:
+        if key is not None:
+            self._key = key
+            self._fernet = Fernet(self._key)
+        else:
+            self._key = None
+            self._fernet = None
+
+    @property
+    def key(self) -> Optional[bytes]:
+        return self._key
+
+    def decrypt(self, msg: bytes):
+        return self._fernet.decrypt(msg) if self._fernet is not None else msg
+
+    def encrypt(self, msg: bytes):
+        return self._fernet.encrypt(msg) if self._fernet is not None else msg
+
+    def handshake(self, key: bytes) -> bool:
+        if self._key is None and key == b"":
+            return True
+        if self._key is not None and self._key == key:
+            return True
+
+        return False
+
+    @staticmethod
+    def gen_key() -> bytes:
+        return Fernet.generate_key()

datumaro/plugins/data_formats/datumaro_binary/exporter.py

@@ -4,26 +4,100 @@
 
 # pylint: disable=no-self-use
 
-from io import TextIOWrapper
+import os.path as osp
+import struct
+from io import BufferedWriter
+from typing import Any, Optional
 
+from datumaro.components.dataset_base import IDataset
+from datumaro.components.exporter import ExportContext
 from datumaro.plugins.data_formats.datumaro.exporter import DatumaroExporter
 from datumaro.plugins.data_formats.datumaro.exporter import _SubsetWriter as __SubsetWriter
+from datumaro.plugins.data_formats.datumaro_binary.crypter import Crypter
+from datumaro.plugins.data_formats.datumaro_binary.mapper import DictMapper
 
 from .format import DatumaroBinaryPath
 
 
 class _SubsetWriter(__SubsetWriter):
     """"""
 
-    def _sign(self, fp: TextIOWrapper):
-        fp.write(DatumaroBinaryPath.SIGNATURE.encode("utf-8"))
+    def __init__(self, context: IDataset, ann_file: str, encryption_key: Optional[bytes] = None):
+        super().__init__(context, ann_file)
+        self._fp: Optional[BufferedWriter] = None
+        self._crypter = Crypter(encryption_key)
+
+    def _sign(self):
+        self._fp.write(DatumaroBinaryPath.SIGNATURE.encode())
+
+    def _dump_encryption_field(self) -> int:
+        if self._crypter.key is None:
+            msg = b""
+        else:
+            msg = self._crypter.key
+            msg = self._crypter.encrypt(msg)
+
+        return self._fp.write(struct.pack(f"I{len(msg)}s", len(msg), msg))
+
+    def _dump_header(self, header: Any):
+        msg = DictMapper.forward(header)
+
+        if self._crypter.key is not None:
+            msg = self._crypter.encrypt(msg)
+
+        length = struct.pack("I", len(msg))
+        return self._fp.write(length + msg)
+
+    def _dump_info(self):
+        self._dump_header(self.infos)
+
+    def _dump_categories(self):
+        self._dump_header(self.categories)
 
     def write(self):
-        with open(self.ann_file, "wb") as fp:
-            self._sign(fp)
+        try:
+            with open(self.ann_file, "wb") as fp:
+                self._fp = fp
+                self._sign()
+                self._dump_encryption_field()
+                self._dump_header(self.infos)
+                self._dump_header(self.categories)
+        finally:
+            self._fp = None
 
 
 class DatumaroBinaryExporter(DatumaroExporter):
     DEFAULT_IMAGE_EXT = DatumaroBinaryPath.IMAGE_EXT
-    WRITER_CLS = _SubsetWriter
     PATH_CLS = DatumaroBinaryPath
+
+    def __init__(
+        self,
+        extractor: IDataset,
+        save_dir: str,
+        *,
+        save_images=None,
+        save_media: Optional[bool] = None,
+        image_ext: Optional[str] = None,
+        default_image_ext: Optional[str] = None,
+        save_dataset_meta: bool = False,
+        ctx: Optional[ExportContext] = None,
+        encryption_key: Optional[bytes] = None,
+    ):
+        self._encryption_key = encryption_key
+        super().__init__(
+            extractor,
+            save_dir,
+            save_images=save_images,
+            save_media=save_media,
+            image_ext=image_ext,
+            default_image_ext=default_image_ext,
+            save_dataset_meta=save_dataset_meta,
+            ctx=ctx,
+        )
+
+    def create_writer(self, subset: str):
+        return _SubsetWriter(
+            context=self,
+            ann_file=osp.join(self._annotations_dir, subset + self.PATH_CLS.ANNOTATION_EXT),
+            encryption_key=self._encryption_key,
+        )

datumaro/plugins/data_formats/datumaro_binary/format.py

@@ -2,6 +2,10 @@
 #
 # SPDX-License-Identifier: MIT
 
+from datumaro.errors import DatasetImportError
+
+_SIGNATURE = "signature:datumaro_binary"
+
 
 class DatumaroBinaryPath:
     IMAGES_DIR = "images"
@@ -13,4 +17,12 @@ class DatumaroBinaryPath:
     ANNOTATION_EXT = ".datumaro"
     IMAGE_EXT = ".jpg"
     MASK_EXT = ".png"
-    SIGNATURE = "signature:datumaro_binary"
+    SIGNATURE = _SIGNATURE
+    SIGNATURE_LEN = len(_SIGNATURE)
+
+    @classmethod
+    def check_signature(cls, signature: str):
+        if signature != cls.SIGNATURE:
+            raise DatasetImportError(
+                f"Input signature={signature} is not aligned with the ground truth signature={cls.SIGNATURE}"
+            )