Support PointCloud dataset by DatumaroBinary format

datumaro/components/media.py

@@ -22,7 +22,7 @@
 
 class MediaType(IntEnum):
     NONE = 0
-    UNKNOWN = 1
+    MEDIA_ELEMENT = 1
     IMAGE = 2
     BYTE_IMAGE = 3
     VIDEO_FRAME = 4
@@ -34,7 +34,7 @@ class MediaType(IntEnum):
 
 
 class MediaElement:
-    MEDIA_TYPE = MediaType.UNKNOWN
+    _type = MediaType.MEDIA_ELEMENT
 
     def __init__(self, path: str) -> None:
         assert path, "Path can't be empty"
@@ -56,9 +56,13 @@ def __eq__(self, other: object) -> bool:
             return False
         return self._path == other._path
 
+    @property
+    def type(self) -> MediaType:
+        return self._type
+
 
 class Image(MediaElement):
-    MEDIA_TYPE = MediaType.IMAGE
+    _type = MediaType.IMAGE
 
     def __init__(
         self,
@@ -187,7 +191,7 @@ def save(self, path):
 
 
 class ByteImage(Image):
-    MEDIA_TYPE = MediaType.BYTE_IMAGE
+    _type = MediaType.BYTE_IMAGE
 
     _FORMAT_MAGICS = (
         (b"\x89PNG\r\n\x1a\n", ".png"),
@@ -255,7 +259,7 @@ def save(self, path):
 
 
 class VideoFrame(Image):
-    MEDIA_TYPE = MediaType.VIDEO_FRAME
+    _type = MediaType.VIDEO_FRAME
 
     def __init__(self, video: Video, index: int):
         self._video = video
@@ -355,7 +359,7 @@ def _navigate_to(self, idx: int) -> VideoFrame:
 
 
 class Video(MediaElement, Iterable[VideoFrame]):
-    MEDIA_TYPE = MediaType.VIDEO
+    _type = MediaType.VIDEO
 
     """
     Provides random access to the video frames.
@@ -525,7 +529,7 @@ def __hash__(self):
 
 
 class PointCloud(MediaElement):
-    MEDIA_TYPE = MediaType.POINT_CLOUD
+    _type = MediaType.POINT_CLOUD
 
     def __init__(self, path: str, extra_images: Optional[List[Image]] = None):
         self._path = path
@@ -534,7 +538,7 @@ def __init__(self, path: str, extra_images: Optional[List[Image]] = None):
 
 
 class MultiframeImage(MediaElement):
-    MEDIA_TYPE = MediaType.MULTIFRAME_IMAGE
+    _type = MediaType.MULTIFRAME_IMAGE
 
     def __init__(
         self,
@@ -566,7 +570,7 @@ def data(self) -> List[Image]:
 
 
 class RoIImage(Image):
-    MEDIA_TYPE = MediaType.ROI_IMAGE
+    _type = MediaType.ROI_IMAGE
 
     def __init__(
         self,
@@ -610,7 +614,7 @@ def save(self, path):
 
 
 class MosaicImage(Image):
-    MEDIA_TYPE = MediaType.MOSAIC_IMAGE
+    _type = MediaType.MOSAIC_IMAGE
 
     def __init__(self, imgs: List[ImageWithRoI], size: Tuple[int, int]) -> None:
         def _get_mosaic_img(_) -> np.ndarray:

datumaro/plugins/data_formats/datumaro/base.py

@@ -136,9 +136,7 @@ def _load_items(self, parsed):
                     related_images = [
                         Image(
                             size=ri.get("size"),
-                            path=osp.join(
-                                self._related_images_dir, self._subset, item_id, ri.get("path")
-                            ),
+                            path=osp.join(ri.get("path")),
                         )
                         for ri in ri_info
                     ]

datumaro/plugins/data_formats/datumaro/exporter.py

@@ -7,6 +7,7 @@
 import os
 import os.path as osp
 import shutil
+from contextlib import contextmanager
 
 import numpy as np
 import pycocotools.mask as mask_utils
@@ -29,7 +30,7 @@
     _Shape,
 )
 from datumaro.components.dataset import ItemStatus
-from datumaro.components.dataset_base import DEFAULT_SUBSET_NAME, DatasetItem, IDataset
+from datumaro.components.dataset_base import DEFAULT_SUBSET_NAME, DatasetItem
 from datumaro.components.exporter import Exporter
 from datumaro.components.media import Image, MediaElement, PointCloud
 from datumaro.util import cast, dump_json_file
@@ -38,7 +39,7 @@
 
 
 class _SubsetWriter:
-    def __init__(self, context: IDataset, ann_file: str):
+    def __init__(self, context: Exporter, ann_file: str):
         self._context = context
 
         self._data = {
@@ -64,68 +65,89 @@ def items(self):
     def is_empty(self):
         return not self.items
 
-    def add_item(self, item: DatasetItem):
-        annotations = []
-        item_desc = {
-            "id": item.id,
-            "annotations": annotations,
-        }
-
-        if item.attributes:
-            item_desc["attr"] = item.attributes
-
-        if isinstance(item.media, Image):
+    @contextmanager
+    def context_save_media(self, item: DatasetItem):
+        """Implicitly change the media path and save it if save_media=True.
+        When done, revert it's path as before.
+        """
+        if item.media is None:
+            yield
+        elif isinstance(item.media, Image):
             image = item.media_as(Image)
             path = image.path
+
             if self._context._save_media:
-                path = self._context._make_image_filename(item)
-                self._context._save_image(
-                    item, osp.join(self._context._images_dir, item.subset, path)
+                # Temporarily update image path and save it.
+                image._path = osp.join(
+                    self._context._images_dir, item.subset, self._context._make_image_filename(item)
                 )
+                self._context._save_image(item, image.path)
 
-            item_desc["image"] = {
-                "path": path,
-            }
-            if item.media.has_size:  # avoid occasional loading
-                item_desc["image"]["size"] = item.media.size
+            yield
+            image._path = path
         elif isinstance(item.media, PointCloud):
             pcd = item.media_as(PointCloud)
             path = pcd.path
+
             if self._context._save_media:
-                path = self._context._make_pcd_filename(item)
-                self._context._save_point_cloud(
-                    item, osp.join(self._context._pcd_dir, item.subset, path)
+                # Temporarily update pcd path and save it.
+                pcd._path = osp.join(
+                    self._context._pcd_dir, item.subset, self._context._make_pcd_filename(item)
                 )
+                self._context._save_point_cloud(item, pcd.path)
+
+                # Temporarily update pcd related images paths and save them.
+                for i, img in enumerate(sorted(pcd.extra_images, key=lambda v: v.path)):
+                    img.__path = img.path
+                    img._path = osp.join(
+                        self._context._related_images_dir,
+                        item.subset,
+                        item.id,
+                        f"image_{i}{self._context._find_image_ext(img)}",
+                    )
 
-            item_desc["point_cloud"] = {"path": path}
+                    if img.has_data:
+                        img.save(img.path)
 
-            images = sorted(pcd.extra_images, key=lambda v: v.path)
+            yield
+            pcd._path = path
             if self._context._save_media:
-                related_images = []
-                for i, img in enumerate(images):
-                    ri_desc = {}
+                for img in pcd.extra_images:
+                    img._path = img.__path
+                    del img.__path
+        else:
+            raise NotImplementedError
 
-                    # Images can have completely the same names or don't
-                    # have them at all, so we just rename them
-                    ri_desc["path"] = f"image_{i}{self._context._find_image_ext(img)}"
+    def add_item(self, item: DatasetItem):
+        annotations = []
+        item_desc = {
+            "id": item.id,
+            "annotations": annotations,
+        }
 
-                    if img.has_data:
-                        img.save(
-                            osp.join(
-                                self._context._related_images_dir,
-                                item.subset,
-                                item.id,
-                                ri_desc["path"],
-                            )
-                        )
-                    if img.has_size:
-                        ri_desc["size"] = img.size
-                    related_images.append(ri_desc)
-            else:
-                related_images = [{"path": img.path} for img in images]
+        if item.attributes:
+            item_desc["attr"] = item.attributes
+
+        with self.context_save_media(item):
+            if isinstance(item.media, Image):
+                image = item.media_as(Image)
+                item_desc["image"] = {
+                    "path": image.path,
+                }
+                if item.media.has_size:  # avoid occasional loading
+                    item_desc["image"]["size"] = image.size
+            elif isinstance(item.media, PointCloud):
+                pcd = item.media_as(PointCloud)
+
+                item_desc["point_cloud"] = {"path": pcd.path}
+
+                related_images = [
+                    {"path": img.path, "size": img.size} if img.has_size else {"path": img.path}
+                    for img in pcd.extra_images
+                ]
 
-            if related_images:
-                item_desc["related_images"] = related_images
+                if related_images:
+                    item_desc["related_images"] = related_images
 
         if isinstance(item.media, MediaElement):
             item_desc["media"] = {"path": item.media.path}

datumaro/plugins/data_formats/datumaro_binary/base.py

@@ -7,6 +7,7 @@
 from typing import Optional
 
 from datumaro.components.errors import DatasetImportError
+from datumaro.components.media import Image, MediaElement, MediaType, PointCloud
 from datumaro.plugins.data_formats.datumaro_binary.format import DatumaroBinaryPath
 from datumaro.plugins.data_formats.datumaro_binary.mapper import DictMapper
 from datumaro.plugins.data_formats.datumaro_binary.mapper.dataset_item import DatasetItemMapper
@@ -32,6 +33,7 @@ def _load_impl(self, path: str) -> None:
                 self._check_encryption_field()
                 self._read_info()
                 self._read_categories()
+                self._read_media_type()
                 self._read_items()
         finally:
             self._fp = None
@@ -49,21 +51,31 @@ def _check_encryption_field(self):
         if not self._crypter.handshake(extracted_key):
             raise DatasetImportError("Encryption key handshake fails. You give a wrong key.")
 
-    def _read_info(self):
+    def _read_header(self):
         len_byte = self._fp.read(4)
         _bytes = self._fp.read(struct.unpack("I", len_byte)[0])
         _bytes = self._crypter.decrypt(_bytes)
+        header, _ = DictMapper.backward(_bytes)
+        return header
 
-        self._infos, _ = DictMapper.backward(_bytes)
+    def _read_info(self):
+        self._infos = self._read_header()
 
     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)
+        categories = self._read_header()
         self._categories = self._load_categories({"categories": categories})
 
+    def _read_media_type(self):
+        media_type = self._read_header()["media_type"]
+        if media_type == MediaType.IMAGE:
+            self._media_type = Image
+        elif media_type == MediaType.POINT_CLOUD:
+            self._media_type = PointCloud
+        elif media_type == MediaType.MEDIA_ELEMENT:
+            self._media_type = MediaElement
+        else:
+            raise NotImplementedError(f"media_type={media_type} is currently not supported.")
+
     def _read_items(self):
         (n_items,) = struct.unpack("I", self._fp.read(4))
         offset = 0

datumaro/plugins/data_formats/datumaro_binary/exporter.py

@@ -6,13 +6,11 @@
 
 import os.path as osp
 import struct
-from contextlib import contextmanager
 from io import BufferedWriter
 from typing import Any, Optional
 
 from datumaro.components.dataset_base import DatasetItem, IDataset
-from datumaro.components.exporter import ExportContext
-from datumaro.components.media import Image
+from datumaro.components.exporter import ExportContext, Exporter
 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
@@ -25,12 +23,14 @@
 class _SubsetWriter(__SubsetWriter):
     """"""
 
-    def __init__(self, context: IDataset, ann_file: str, encryption_key: Optional[bytes] = None):
+    def __init__(self, context: Exporter, ann_file: str, encryption_key: Optional[bytes] = None):
         super().__init__(context, ann_file)
         self._fp: Optional[BufferedWriter] = None
         self._crypter = Crypter(encryption_key)
         self._data["items"] = bytearray()
         self._item_cnt = 0
+        media_type = context._extractor.media_type()
+        self._media_type = {"media_type": media_type._type}
 
     def _sign(self):
         self._fp.write(DatumaroBinaryPath.SIGNATURE.encode())
@@ -60,29 +60,10 @@ def _dump_categories(self):
         self._dump_header(self.categories)
 
     def add_item(self, item: DatasetItem):
-        with self.media_context(item):
+        with self.context_save_media(item):
             self.items.extend(DatasetItemMapper.forward(item))
         self._item_cnt += 1
 
-    @contextmanager
-    def media_context(self, item: DatasetItem):
-        if item.media is None:
-            yield
-        elif isinstance(item.media, Image):
-            image = item.media_as(Image)
-            _path = image.path
-
-            if self._context._save_media:
-                path = self._context._make_image_filename(item)
-                full_path = osp.join(self._context._images_dir, item.subset, path)
-                self._context._save_image(item, full_path)
-                image._path = full_path
-
-            yield
-            image._path = _path
-        else:
-            raise NotImplementedError
-
     def _dump_items(self):
         items_bytes = bytes(self.items)
         n_items_bytes = len(items_bytes)
@@ -96,6 +77,7 @@ def write(self):
                 self._dump_encryption_field()
                 self._dump_header(self.infos)
                 self._dump_header(self.categories)
+                self._dump_header(self._media_type)
                 self._dump_items()
         finally:
             self._fp = None