Generic media support

datumaro/cli/__main__.py

@@ -6,6 +6,7 @@
 import logging as log
 import os.path as osp
 import sys
+import warnings
 
 from ..util.telemetry_utils import (
     close_telemetry_session, init_telemetry_session,
@@ -39,6 +40,10 @@ def init_logger(cls, args=None):
         log.basicConfig(format='%(asctime)s %(levelname)s: %(message)s',
             level=args.loglevel)
 
+        # Suppress own deprecation warnings
+        warnings.filterwarnings('ignore', category=DeprecationWarning,
+            module=r'datumaro\..*')
+
     @staticmethod
     def _define_loglevel_option(parser):
         parser.add_argument('--loglevel', type=loglevel, default='info',

datumaro/cli/commands/download.py

@@ -47,7 +47,7 @@ def build_parser(parser_ctor=argparse.ArgumentParser):
         |n
         Examples:|n
         - Download the MNIST dataset:|n
-        |s|s%(prog)s -i tfds:mnist -- --save-images|n
+        |s|s%(prog)s -i tfds:mnist -- --save-media|n
         |n
         - Download the VOC 2012 dataset, saving only the annotations in the COCO
           format into a specific directory:|n

datumaro/cli/contexts/project/__init__.py

@@ -371,10 +371,10 @@ def filter_command(args):
 
                 # Source might be missing in the working dir, so we specify
                 # the output directory.
-                # We specify save_images here as a heuristic. It can probably
+                # We specify save_media here as a heuristic. It can probably
                 # be improved by checking if there are images in the dataset
                 # directory.
-                dataset.save(project.source_data_dir(target), save_images=True)
+                dataset.save(project.source_data_dir(target), save_media=True)
 
             log.info("Finished")
         else:
@@ -389,7 +389,7 @@ def filter_command(args):
             dst_dir = osp.abspath(dst_dir)
 
             dataset.filter(filter_expr, *filter_args)
-            dataset.save(dst_dir, save_images=True)
+            dataset.save(dst_dir, save_media=True)
 
             log.info("Results have been saved to '%s'" % dst_dir)
 
@@ -557,10 +557,10 @@ def transform_command(args):
 
                 # Source might be missing in the working dir, so we specify
                 # the output directory
-                # We specify save_images here as a heuristic. It can probably
+                # We specify save_media here as a heuristic. It can probably
                 # be improved by checking if there are images in the dataset
                 # directory.
-                dataset.save(project.source_data_dir(target), save_images=True)
+                dataset.save(project.source_data_dir(target), save_media=True)
 
             log.info("Finished")
         else:
@@ -575,7 +575,7 @@ def transform_command(args):
             dst_dir = osp.abspath(dst_dir)
 
             dataset.transform(args.transform, **extra_args)
-            dataset.save(dst_dir, save_images=True)
+            dataset.save(dst_dir, save_media=True)
 
             log.info("Results have been saved to '%s'" % dst_dir)
 

datumaro/cli/contexts/project/diff.py

@@ -14,6 +14,8 @@
 import cv2
 import numpy as np
 
+from datumaro.components.media import Image
+
 with warnings.catch_warnings():
     warnings.simplefilter("ignore")
     import tensorboardX as tb
@@ -121,7 +123,10 @@ def save(self, a: IDataset, b: IDataset):
             self.update_mask_confusion(mask_diff)
 
             self.save_item_label_diff(item_a, item_b, label_diff)
-            self.save_item_bbox_diff(item_a, item_b, bbox_diff)
+
+            if a.media_type() and issubclass(a.media_type(), Image) and \
+                    b.media_type() and issubclass(b.media_type(), Image):
+                self.save_item_bbox_diff(item_a, item_b, bbox_diff)
 
         if len(self.label_confusion_matrix) != 0:
             self.save_conf_matrix(self.label_confusion_matrix,
@@ -243,11 +248,11 @@ def save_item_bbox_diff(self, item_a, item_b, diff):
         _, mispred, a_unmatched, b_unmatched = diff
 
         if 0 < len(a_unmatched) + len(b_unmatched) + len(mispred):
-            if not item_a.has_image or not item_a.image.has_data:
+            if not isinstance(item_a.media, Image) or not item_a.media.has_data:
                 log.warning("Item %s: item has no image data, "
                     "it will be skipped" % (item_a.id))
                 return
-            img_a = item_a.image.data.copy()
+            img_a = item_a.media.data.copy()
             img_b = img_a.copy()
             for a_bbox, b_bbox in mispred:
                 self.draw_bbox(img_a, a_bbox, self.get_a_label(a_bbox.label),

datumaro/components/converter.py

@@ -8,16 +8,17 @@
 import os
 import os.path as osp
 import shutil
+import warnings
 
 from attrs import define, field
 import attr
 
 from datumaro.components.cli_plugin import CliPlugin
 from datumaro.components.errors import (
-    AnnotationExportError, DatumaroError, ItemExportError,
+    AnnotationExportError, DatasetExportError, DatumaroError, ItemExportError,
 )
 from datumaro.components.extractor import DatasetItem, IExtractor
-from datumaro.components.media import Image
+from datumaro.components.media import Image, PointCloud
 from datumaro.components.progress_reporting import (
     NullProgressReporter, ProgressReporter,
 )
@@ -89,8 +90,10 @@ class Converter(CliPlugin):
     @classmethod
     def build_cmdline_parser(cls, **kwargs):
         parser = super().build_cmdline_parser(**kwargs)
-        parser.add_argument('--save-images', action='store_true',
-            help="Save images (default: %(default)s)")
+        parser.add_argument('--save-images', action='store_true', default=None,
+            help="Save images (default: %s)" % (None))
+        parser.add_argument('--save-media', action='store_true', default=None,
+            help="Save media (default: %s)" % (None))
         parser.add_argument('--image-ext', default=None,
             help="Image extension (default: keep or use format default%s)" % \
                 (' ' + cls.DEFAULT_IMAGE_EXT if cls.DEFAULT_IMAGE_EXT else ''))
@@ -138,7 +141,8 @@ def apply(self):
         raise NotImplementedError("Should be implemented in a subclass")
 
     def __init__(self, extractor: IExtractor, save_dir: str, *,
-            save_images: bool = False,
+            save_images = None,
+            save_media: bool = None,
             image_ext: Optional[str] = None,
             default_image_ext: Optional[str] = None,
             save_dataset_meta: bool = False,
@@ -147,7 +151,20 @@ def __init__(self, extractor: IExtractor, save_dir: str, *,
         assert default_image_ext
         self._default_image_ext = default_image_ext
 
-        self._save_images = save_images
+        if save_images is not None and save_media is not None:
+            raise DatasetExportError("Can't use 'save-media' and "
+                "save-images together")
+
+        if save_media is not None:
+            self._save_media = save_media
+        elif save_images is not None:
+            self._save_media = save_images
+            warnings.warn("'save-images' is deprecated and will be "
+                "removed in future. Use 'save-media' instead.",
+                DeprecationWarning, stacklevel=2)
+        else:
+            self._save_media = False
+
         self._image_ext = image_ext
 
         self._extractor = extractor
@@ -168,8 +185,8 @@ def __init__(self, extractor: IExtractor, save_dir: str, *,
     def _find_image_ext(self, item: Union[DatasetItem, Image]):
         src_ext = None
 
-        if isinstance(item, DatasetItem) and item.has_image:
-            src_ext = item.image.ext
+        if isinstance(item, DatasetItem) and isinstance(item.media, Image):
+            src_ext = item.media.ext
         elif isinstance(item, Image):
             src_ext = item.ext
 
@@ -192,7 +209,7 @@ def _save_image(self, item, path=None, *,
         assert not ((subdir or name or basedir) and path), \
             "Can't use both subdir or name or basedir and path arguments"
 
-        if not item.has_image or not item.image.has_data:
+        if not isinstance(item.media, Image) or not item.media.has_data:
             log.warning("Item '%s' has no image", item.id)
             return
 
@@ -201,14 +218,14 @@ def _save_image(self, item, path=None, *,
             self._make_image_filename(item, name=name, subdir=subdir))
         path = osp.abspath(path)
 
-        item.image.save(path)
+        item.media.save(path)
 
     def _save_point_cloud(self, item=None, path=None, *,
             name=None, subdir=None, basedir=None):
         assert not ((subdir or name or basedir) and path), \
             "Can't use both subdir or name or basedir and path arguments"
 
-        if not item.point_cloud:
+        if not item.media or not isinstance(item.media, PointCloud):
             log.warning("Item '%s' has no pcd", item.id)
             return
 
@@ -218,9 +235,9 @@ def _save_point_cloud(self, item=None, path=None, *,
         path = osp.abspath(path)
 
         os.makedirs(osp.dirname(path), exist_ok=True)
-        if item.point_cloud and osp.isfile(item.point_cloud):
-            if item.point_cloud != path:
-                shutil.copyfile(item.point_cloud, path)
+        if item.media and osp.isfile(item.media.path):
+            if item.media.path != path:
+                shutil.copyfile(item.media.path, path)
 
     def _save_meta_file(self, path):
         save_meta_file(path, self._extractor.categories())

datumaro/components/dataset.py

@@ -27,7 +27,7 @@
 )
 from datumaro.components.environment import Environment
 from datumaro.components.errors import (
-    CategoriesRedefinedError, ConflictingCategoriesError,
+    CategoriesRedefinedError, ConflictingCategoriesError, MediaTypeError,
     MultipleFormatsMatchError, NoMatchingFormatsError, RepeatedItemError,
     UnknownFormatError,
 )
@@ -36,6 +36,7 @@
     ImportContext, ImportErrorPolicy, ItemTransform, Transform, _ImportFail,
 )
 from datumaro.components.launcher import Launcher, ModelTransform
+from datumaro.components.media import Image, MediaElement
 from datumaro.components.progress_reporting import (
     NullProgressReporter, ProgressReporter,
 )
@@ -159,9 +160,11 @@ def categories(self):
             return self.parent.categories()
 
 
-    def __init__(self, parent: DatasetItemStorage, categories: CategoriesInfo):
+    def __init__(self, parent: DatasetItemStorage, categories: CategoriesInfo,
+            media_type: Optional[Type[MediaElement]]):
         self._parent = parent
         self._categories = categories
+        self._media_type = media_type
 
     def __iter__(self):
         yield from self._parent
@@ -184,6 +187,8 @@ def subsets(self):
     def get(self, id, subset=None):
         return self._parent.get(id, subset=subset)
 
+    def media_type(self):
+        return self._media_type
 
 class ItemStatus(Enum):
     added = auto()
@@ -195,7 +200,7 @@ class DatasetPatchWrapper(DatasetItemStorageDatasetView):
         # The purpose of this class is to indicate that the input dataset is
         # a patch and autofill patch info in Converter
         def __init__(self, patch: DatasetPatch, parent: IDataset):
-            super().__init__(patch.data, parent.categories())
+            super().__init__(patch.data, parent.categories(), parent.media_type())
             self.patch = patch
 
         def subsets(self):
@@ -261,19 +266,25 @@ def subsets(self):
     def categories(self):
         return self.parent.categories()
 
+    def media_type(self):
+        return self.parent.media_type()
+
     def as_dataset(self) -> Dataset:
         return Dataset.from_extractors(self, env=self.parent.env)
 
 
 class DatasetStorage(IDataset):
     def __init__(self, source: Union[IDataset, DatasetItemStorage] = None,
-            categories: CategoriesInfo = None):
+            categories: CategoriesInfo = None,
+            media_type: Optional[Type[MediaElement]] = None):
         if source is None and categories is None:
             categories = {}
         elif isinstance(source, IDataset) and categories is not None:
             raise ValueError("Can't use both source and categories")
         self._categories = categories
 
+        self._media_type = media_type
+
         # Possible combinations:
         # 1. source + storage
         #      - Storage contains a patch to the Source data.
@@ -392,6 +403,11 @@ def _update_status(item_id, new_status: ItemStatus):
 
         i = -1
         for i, item in enumerate(source):
+            if source.media_type():
+                if item.media and not isinstance(item.media, source.media_type()):
+                    raise MediaTypeError("Dataset elements must have a '%s' " \
+                        "media type" % source.media_type())
+
             if transform and transform.is_local:
                 old_id = (item.id, item.subset)
                 item = transform.transform_item(item)
@@ -478,7 +494,8 @@ def _merged(self) -> IDataset:
             return self._source
         elif self._source is not None:
             self.init_cache()
-        return DatasetItemStorageDatasetView(self._storage, self._categories)
+        return DatasetItemStorageDatasetView(self._storage,
+            self._categories, self._media_type)
 
     def __len__(self) -> int:
         if self._length is None:
@@ -502,6 +519,14 @@ def define_categories(self, categories: CategoriesInfo):
             raise CategoriesRedefinedError()
         self._categories = categories
 
+    def media_type(self):
+        if self.is_cache_initialized():
+            return self._media_type
+        elif self._media_type is not None:
+            return self._media_type
+        else:
+            return self._source.media_type()
+
     def put(self, item):
         is_new = self._storage.put(item)
 
@@ -624,7 +649,8 @@ class Dataset(IDataset):
     @classmethod
     def from_iterable(cls, iterable: Iterable[DatasetItem],
             categories: Union[CategoriesInfo, List[str], None] = None,
-            env: Optional[Environment] = None) -> Dataset:
+            env: Optional[Environment] = None,
+            media_type: Type = Image) -> Dataset:
         if isinstance(categories, list):
             categories = { AnnotationType.label:
                 LabelCategories.from_iterable(categories)
@@ -636,7 +662,8 @@ def from_iterable(cls, iterable: Iterable[DatasetItem],
         class _extractor(Extractor):
             def __init__(self):
                 super().__init__(length=len(iterable) \
-                    if hasattr(iterable, '__len__') else None)
+                        if hasattr(iterable, '__len__') else None,
+                    media_type=media_type)
 
             def __iter__(self):
                 return iter(iterable)
@@ -657,19 +684,23 @@ def from_extractors(*sources: IDataset,
             source = ExactMerge.merge(*sources)
             categories = ExactMerge.merge_categories(
                 s.categories() for s in sources)
-            dataset = Dataset(source=source, categories=categories, env=env)
+            media_type=ExactMerge.merge_media_types(sources)
+            dataset = Dataset(source=source, categories=categories,
+                media_type=media_type, env=env)
         return dataset
 
     def __init__(self, source: Optional[IDataset] = None, *,
             categories: Optional[CategoriesInfo] = None,
+            media_type: Optional[Type[MediaElement]] = None,
             env: Optional[Environment] = None) -> None:
         super().__init__()
 
         assert env is None or isinstance(env, Environment), env
         self._env = env
 
         self.eager = None
-        self._data = DatasetStorage(source, categories=categories)
+        self._data = DatasetStorage(source, categories=categories,
+            media_type=media_type)
         if self.is_eager:
             self.init_cache()
 
@@ -698,6 +729,9 @@ def subsets(self) -> Dict[str, DatasetSubset]:
     def categories(self) -> CategoriesInfo:
         return self._data.categories()
 
+    def media_type(self) -> Optional[Type[MediaElement]]:
+        return self._data.media_type()
+
     def get(self, id: str, subset: Optional[str] = None) \
             -> Optional[DatasetItem]:
         return self._data.get(id, subset)
@@ -1125,7 +1159,6 @@ def detect(path: str, *,
         if 1 < len(matches):
             raise MultipleFormatsMatchError(matches)
 
-
 @contextmanager
 def eager_mode(new_mode: bool = True, dataset: Optional[Dataset] = None) -> None:
     if dataset is not None: