allow dictionary image only and test endianness

monai/transforms/io/array.py

@@ -33,6 +33,27 @@
 __all__ = ["LoadImage", "SaveImage"]
 
 
+def switch_endianness(data, old, new):
+    """
+    If any numpy arrays have `old` (e.g., ">"),
+    replace with `new` (e.g., "<").
+    """
+    if isinstance(data, np.ndarray):
+        if data.dtype.byteorder == old:
+            data = data.newbyteorder(new)
+    elif isinstance(data, tuple):
+        data = (switch_endianness(x, old, new) for x in data)
+    elif isinstance(data, list):
+        data = [switch_endianness(x, old, new) for x in data]
+    elif isinstance(data, dict):
+        data = {k: switch_endianness(v, old, new) for k, v in data.items()}
+    elif isinstance(data, (bool, str, float, int)):
+        pass
+    else:
+        raise AssertionError()
+    return data
+
+
 class LoadImage(Transform):
     """
     Load image file or files from provided path based on reader.
@@ -132,6 +153,9 @@ def __call__(
         if self.image_only:
             return img_array
         meta_data[Key.FILENAME_OR_OBJ] = ensure_tuple(filename)[0]
+        # make sure all elements in metadata are little endian
+        meta_data = switch_endianness(meta_data, ">", "<")
+
         return img_array, meta_data
 
 

monai/transforms/io/dictionary.py

@@ -59,6 +59,7 @@ def __init__(
         dtype: DtypeLike = np.float32,
         meta_key_postfix: str = "meta_dict",
         overwriting: bool = False,
+        image_only: bool = False,
         *args,
         **kwargs,
     ) -> None:
@@ -76,11 +77,13 @@ def __init__(
                 For example, load nifti file for `image`, store the metadata into `image_meta_dict`.
             overwriting: whether allow to overwrite existing meta data of same key.
                 default is False, which will raise exception if encountering existing key.
+            image_only: if True return dictionary containing just only the image volumes, otherwise return
+                dictionary containing image data array and header dict per input key.
             args: additional parameters for reader if providing a reader name.
             kwargs: additional parameters for reader if providing a reader name.
         """
         super().__init__(keys)
-        self._loader = LoadImage(reader, False, dtype, *args, **kwargs)
+        self._loader = LoadImage(reader, image_only, dtype, *args, **kwargs)
         if not isinstance(meta_key_postfix, str):
             raise TypeError(f"meta_key_postfix must be a str but is {type(meta_key_postfix).__name__}.")
         self.meta_key_postfix = meta_key_postfix
@@ -98,15 +101,20 @@ def __call__(self, data, reader: Optional[ImageReader] = None):
         d = dict(data)
         for key in self.keys:
             data = self._loader(d[key], reader)
-            if not isinstance(data, (tuple, list)):
-                raise ValueError("loader must return a tuple or list.")
-            d[key] = data[0]
-            if not isinstance(data[1], dict):
-                raise ValueError("metadata must be a dict.")
-            key_to_add = f"{key}_{self.meta_key_postfix}"
-            if key_to_add in d and not self.overwriting:
-                raise KeyError(f"Meta data with key {key_to_add} already exists and overwriting=False.")
-            d[key_to_add] = data[1]
+            if self._loader.image_only:
+                if not isinstance(data, np.ndarray):
+                    raise ValueError("loader must return a numpy array (because image_only=True was used).")
+                d[key] = data
+            else:
+                if not isinstance(data, (tuple, list)):
+                    raise ValueError("loader must return a tuple or list (because image_only=False was used).")
+                d[key] = data[0]
+                if not isinstance(data[1], dict):
+                    raise ValueError("metadata must be a dict.")
+                key_to_add = f"{key}_{self.meta_key_postfix}"
+                if key_to_add in d and not self.overwriting:
+                    raise KeyError(f"Meta data with key {key_to_add} already exists and overwriting=False.")
+                d[key_to_add] = data[1]
         return d
 
 

tests/test_nifti_endianness.py

@@ -0,0 +1,48 @@
+import tempfile
+import unittest
+from typing import TYPE_CHECKING, List, Tuple
+from unittest.case import skipUnless
+
+import numpy as np
+from parameterized import parameterized
+
+from monai.data import DataLoader, Dataset, create_test_image_2d
+from monai.transforms import LoadImage, LoadImaged
+from monai.utils.module import optional_import
+
+if TYPE_CHECKING:
+    import nibabel as nib
+
+    has_nib = True
+else:
+    nib, has_nib = optional_import("nibabel")
+
+TESTS: List[Tuple] = []
+for endianness in ["<", ">"]:
+    for use_array in [True, False]:
+        for image_only in [True, False]:
+            TESTS.append((endianness, use_array, image_only))
+
+
+class TestNiftiEndianness(unittest.TestCase):
+    def setUp(self):
+        self.im, _ = create_test_image_2d(100, 100)
+        self.fname = tempfile.NamedTemporaryFile(suffix=".nii.gz").name
+
+    @parameterized.expand(TESTS)
+    @skipUnless(has_nib, "Requires NiBabel")
+    def test_endianness(self, endianness, use_array, image_only):
+
+        hdr = nib.Nifti1Header(endianness=endianness)
+        nii = nib.Nifti1Image(self.im, np.eye(4), header=hdr)
+        nib.save(nii, self.fname)
+
+        data = [self.fname] if use_array else [{"image": self.fname}]
+        tr = LoadImage(image_only=image_only) if use_array else LoadImaged("image", image_only=image_only)
+        check_ds = Dataset(data, tr)
+        check_loader = DataLoader(check_ds, batch_size=1)
+        _ = next(iter(check_loader))
+
+
+if __name__ == "__main__":
+    unittest.main()