Import for BraTS dataset format

CHANGELOG.md

@@ -13,6 +13,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   (<https://github.com/openvinotoolkit/datumaro/pull/539>)
 - \[API\] A way to request dataset and extractor media type with `media_type`
   (<https://github.com/openvinotoolkit/datumaro/pull/539>)
+- BraTS format (import-only) (.npy and .nii.gz), new `MultiframeImage`
+  media type (<https://github.com/openvinotoolkit/datumaro/pull/628>)
 
 ### Changed
 - TBD

datumaro/components/media.py

@@ -500,3 +500,30 @@ 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):
+    def __init__(
+        self,
+        images: Optional[Iterable[Union[str, Image, np.ndarray, Callable[[str], np.ndarray]]]],
+        *,
+        path: Optional[str] = None,
+    ):
+        self._path = path
+
+        self._images = [None] * len(images or [])
+        for i, image in enumerate(images or []):
+            assert isinstance(image, (str, Image, np.ndarray)) or callable(image)
+
+            if isinstance(image, str):
+                image = Image(path=image)
+            elif isinstance(image, np.ndarray) or callable(image):
+                image = Image(data=image)
+
+            self._images[i] = image
+
+        assert self._path or self._images
+
+    @property
+    def data(self) -> List[Image]:
+        return self._images

datumaro/components/operations.py

@@ -42,7 +42,7 @@
     WrongGroupError,
 )
 from datumaro.components.extractor import CategoriesInfo, DatasetItem
-from datumaro.components.media import Image, MediaElement, PointCloud, Video
+from datumaro.components.media import Image, MediaElement, MultiframeImage, PointCloud, Video
 from datumaro.util import filter_dict, find
 from datumaro.util.annotation_util import (
     OKS,
@@ -187,6 +187,10 @@ def _merge_media(
             not item_b.media or isinstance(item_b.media, Video)
         ):
             media = cls._merge_videos(item_a, item_b)
+        elif (not item_a.media or isinstance(item_a.media, MultiframeImage)) and (
+            not item_b.media or isinstance(item_b.media, MultiframeImage)
+        ):
+            media = cls._merge_multiframe_images(item_a, item_b)
         elif (not item_a.media or isinstance(item_a.media, MediaElement)) and (
             not item_b.media or isinstance(item_b.media, MediaElement)
         ):
@@ -330,6 +334,38 @@ def _merge_videos(item_a: DatasetItem, item_b: DatasetItem) -> Video:
 
         return media
 
+    @staticmethod
+    def _merge_multiframe_images(item_a: DatasetItem, item_b: DatasetItem) -> MultiframeImage:
+        media = None
+
+        if isinstance(item_a.media, MultiframeImage) and isinstance(item_b.media, MultiframeImage):
+            if item_a.media.path and item_b.media.path and item_a.media.path != item_b.media.path:
+                raise MismatchingMediaPathError(
+                    (item_a.id, item_a.subset), item_a.media.path, item_b.media.path
+                )
+
+            if item_a.media.path or item_a.media.data:
+                media = item_a.media
+
+                if item_b.media.data:
+                    for image in item_b.media.data:
+                        if image not in media.data:
+                            media.data.append(image)
+            else:
+                media = item_b.media
+
+                if item_a.media.data:
+                    for image in item_a.media.data:
+                        if image not in media.data:
+                            media.data.append(image)
+
+        elif isinstance(item_a.media, MultiframeImage):
+            media = item_a.media
+        else:
+            media = item_b.media
+
+        return media
+
     @staticmethod
     def _merge_anno(a: Iterable[Annotation], b: Iterable[Annotation]) -> List[Annotation]:
         return merge_annotations_equal(a, b)

datumaro/plugins/brats_format.py

@@ -0,0 +1,106 @@
+# Copyright (C) 2022 Intel Corporation
+#
+# SPDX-License-Identifier: MIT
+
+import glob
+import os.path as osp
+
+import nibabel as nib
+import numpy as np
+
+from datumaro.components.annotation import AnnotationType, LabelCategories, Mask
+from datumaro.components.extractor import DatasetItem, Importer, SourceExtractor
+from datumaro.components.format_detection import FormatDetectionContext
+from datumaro.components.media import MultiframeImage
+
+
+class BratsPath:
+    IMAGES_DIR = "images"
+    LABELS = "labels"
+    DATA_EXT = ".nii.gz"
+
+
+class BratsExtractor(SourceExtractor):
+    def __init__(self, path):
+        if not osp.isdir(path):
+            raise FileNotFoundError("Can't read dataset directory '%s'" % path)
+
+        self._subset_suffix = osp.basename(path)[len(BratsPath.IMAGES_DIR) :]
+        subset = None
+        if self._subset_suffix == "Tr":
+            subset = "train"
+        elif self._subset_suffix == "Ts":
+            subset = "test"
+        super().__init__(subset=subset, media_type=MultiframeImage)
+
+        self._root_dir = osp.dirname(path)
+        self._categories = self._load_categories()
+        self._items = list(self._load_items(path).values())
+
+    def _load_categories(self):
+        label_cat = LabelCategories()
+
+        labels_path = osp.join(self._root_dir, BratsPath.LABELS)
+        if osp.isfile(labels_path):
+            with open(labels_path, encoding="utf-8") as f:
+                for line in f:
+                    label_cat.add(line.strip())
+
+        return {AnnotationType.label: label_cat}
+
+    def _load_items(self, path):
+        items = {}
+
+        for image_path in glob.glob(osp.join(path, f"*{BratsPath.DATA_EXT}")):
+            data = nib.load(image_path).get_fdata()
+
+            item_id = osp.basename(image_path)[: -len(BratsPath.DATA_EXT)]
+
+            images = [0] * data.shape[2]
+            for i in range(data.shape[2]):
+                images[i] = data[:, :, i]
+
+            items[item_id] = DatasetItem(
+                id=item_id, subset=self._subset, media=MultiframeImage(images, path=image_path)
+            )
+
+        masks_dir = osp.join(self._root_dir, BratsPath.LABELS + self._subset_suffix)
+        for mask in glob.glob(osp.join(masks_dir, f"*{BratsPath.DATA_EXT}")):
+            data = nib.load(mask).get_fdata()
+
+            item_id = osp.basename(image_path)[: -len(BratsPath.DATA_EXT)]
+
+            if item_id not in items:
+                items[item_id] = DatasetItem(id=item_id)
+
+            anno = []
+            for i in range(data.shape[2]):
+                classes = np.unique(data[:, :, i])
+                for class_id in classes:
+                    anno.append(
+                        Mask(
+                            image=self._lazy_extract_mask(data[:, :, i], class_id),
+                            label=class_id,
+                            attributes={"image_id": i},
+                        )
+                    )
+
+            items[item_id].annotations = anno
+
+        return items
+
+    @staticmethod
+    def _lazy_extract_mask(mask, c):
+        return lambda: mask == c
+
+
+class BratsImporter(Importer):
+    @classmethod
+    def detect(cls, context: FormatDetectionContext) -> None:
+        with context.require_any():
+            with context.alternative():
+                context.require_file(f"*/*{BratsPath.DATA_EXT}")
+
+    @classmethod
+    def find_sources(cls, path):
+        return cls._find_sources_recursive(path, "", "brats", filename=f"{BratsPath.IMAGES_DIR}*")

datumaro/plugins/brats_numpy_format.py

@@ -0,0 +1,106 @@
+# Copyright (C) 2022 Intel Corporation
+#
+# SPDX-License-Identifier: MIT
+
+import os.path as osp
+
+import numpy as np
+
+from datumaro.components.annotation import AnnotationType, Cuboid3d, LabelCategories, Mask
+from datumaro.components.extractor import DatasetItem, Importer, SourceExtractor
+from datumaro.components.format_detection import FormatDetectionContext
+from datumaro.components.media import MultiframeImage
+from datumaro.util.pickle_util import PickleLoader
+
+
+class BratsNumpyPath:
+    IDS_FILE = "val_ids.p"
+    BOXES_FILE = "val_brain_bbox.p"
+    LABELS_FILE = "labels"
+    DATA_SUFFIX = "_data_cropped"
+    LABEL_SUFFIX = "_label_cropped"
+
+
+class BratsNumpyExtractor(SourceExtractor):
+    def __init__(self, path):
+        if not osp.isfile(path):
+            raise FileNotFoundError("Can't read annotation file '%s'" % path)
+
+        super().__init__(media_type=MultiframeImage)
+
+        self._root_dir = osp.dirname(path)
+        self._categories = self._load_categories()
+        self._items = list(self._load_items(path).values())
+
+    def _load_categories(self):
+        label_cat = LabelCategories()
+
+        labels_path = osp.join(self._root_dir, BratsNumpyPath.LABELS_FILE)
+        if osp.isfile(labels_path):
+            with open(labels_path, encoding="utf-8") as f:
+                for line in f:
+                    label_cat.add(line.strip())
+
+        return {AnnotationType.label: label_cat}
+
+    def _load_items(self, path):
+        items = {}
+
+        with open(path, "rb") as f:
+            ids = PickleLoader.restricted_load(f)
+
+        boxes = None
+        boxes_file = osp.join(self._root_dir, BratsNumpyPath.BOXES_FILE)
+        if osp.isfile(boxes_file):
+            with open(boxes_file, "rb") as f:
+                boxes = PickleLoader.restricted_load(f)
+
+        for i, item_id in enumerate(ids):
+            image_path = osp.join(self._root_dir, item_id + BratsNumpyPath.DATA_SUFFIX + ".npy")
+            media = None
+            if osp.isfile(image_path):
+                data = np.load(image_path)[0].transpose()
+                images = [0] * data.shape[2]
+                for j in range(data.shape[2]):
+                    images[j] = data[:, :, j]
+
+                media = MultiframeImage(images, path=image_path)
+
+            anno = []
+            mask_path = osp.join(self._root_dir, item_id + BratsNumpyPath.LABEL_SUFFIX + ".npy")
+            if osp.isfile(mask_path):
+                mask = np.load(mask_path)[0].transpose()
+                for j in range(mask.shape[2]):
+                    classes = np.unique(mask[:, :, j])
+                    for class_id in classes:
+                        anno.append(
+                            Mask(
+                                image=self._lazy_extract_mask(mask[:, :, j], class_id),
+                                label=class_id,
+                                attributes={"image_id": j},
+                            )
+                        )
+
+            if boxes is not None:
+                box = boxes[i]
+                anno.append(Cuboid3d(position=list(box[0]), rotation=list(box[1])))
+
+            items[item_id] = DatasetItem(id=item_id, media=media, annotations=anno)
+
+        return items
+
+    @staticmethod
+    def _lazy_extract_mask(mask, c):
+        return lambda: mask == c
+
+
+class BratsNumpyImporter(Importer):
+    @classmethod
+    def detect(cls, context: FormatDetectionContext) -> None:
+        context.require_file(BratsNumpyPath.IDS_FILE)
+
+    @classmethod
+    def find_sources(cls, path):
+        return cls._find_sources_recursive(
+            path, "", "brats_numpy", filename=BratsNumpyPath.IDS_FILE
+        )

datumaro/plugins/cifar_format.py

@@ -8,7 +8,6 @@
 from collections import OrderedDict
 
 import numpy as np
-import numpy.core.multiarray
 
 from datumaro.components.annotation import AnnotationType, Label, LabelCategories
 from datumaro.components.converter import Converter
@@ -18,26 +17,7 @@
 from datumaro.components.media import Image
 from datumaro.util import cast
 from datumaro.util.meta_file_util import has_meta_file, parse_meta_file
-
-
-class RestrictedUnpickler(pickle.Unpickler):
-    def find_class(self, module, name):
-        if module == "numpy.core.multiarray" and name in PickleLoader.safe_numpy:
-            return getattr(numpy.core.multiarray, name)
-        elif module == "numpy" and name in PickleLoader.safe_numpy:
-            return getattr(numpy, name)
-        raise pickle.UnpicklingError("Global '%s.%s' is forbidden" % (module, name))
-
-
-class PickleLoader:
-    safe_numpy = {
-        "dtype",
-        "ndarray",
-        "_reconstruct",
-    }
-
-    def restricted_load(s):
-        return RestrictedUnpickler(s, encoding="latin1").load()
+from datumaro.util.pickle_util import PickleLoader
 
 
 class CifarPath:

datumaro/plugins/mnist_format.py

@@ -121,7 +121,11 @@ class MnistImporter(Importer):
     @classmethod
     def find_sources(cls, path):
         return cls._find_sources_recursive(
-            path, ".gz", "mnist", file_filter=lambda p: osp.basename(p).split("-")[1] == "labels"
+            path,
+            ".gz",
+            "mnist",
+            file_filter=lambda p: 1 < len(osp.basename(p).split("-"))
+            and osp.basename(p).split("-")[1] == "labels",
         )
 
 

datumaro/util/pickle_util.py

@@ -0,0 +1,27 @@
+# Copyright (C) 2022 Intel Corporation
+#
+# SPDX-License-Identifier: MIT
+
+import pickle  # nosec - disable B403:import_pickle check - fixed
+
+import numpy.core.multiarray
+
+
+class RestrictedUnpickler(pickle.Unpickler):
+    def find_class(self, module, name):
+        if module == "numpy.core.multiarray" and name in PickleLoader.safe_numpy:
+            return getattr(numpy.core.multiarray, name)
+        elif module == "numpy" and name in PickleLoader.safe_numpy:
+            return getattr(numpy, name)
+        raise pickle.UnpicklingError("Global '%s.%s' is forbidden" % (module, name))
+
+
+class PickleLoader:
+    safe_numpy = {
+        "dtype",
+        "ndarray",
+        "_reconstruct",
+    }
+
+    def restricted_load(s):
+        return RestrictedUnpickler(s, encoding="latin1").load()