Repository: Project-MONAI/MONAI

Issue:
IO formats 
Hello

How do I perform multi class semantic (instance) segmentation over my set of 2D PNG/TIF grayscale images and their corresponding masks file (NPY numpy array, TIF or any other) matrix representation with # channels corresponding to the classes?
Do I need to use NIFTI_interface? If so, how do I convert my PNG/TIF/NPY files to this format?

Thank you

Discussion:
Nic-Ma: /black

Patch:
diff --git a/docs/source/transforms.rst b/docs/source/transforms.rst
index 692108f9..d5bb39bb 100644
--- a/docs/source/transforms.rst
+++ b/docs/source/transforms.rst
@@ -183,6 +183,12 @@ IO
     :members:
     :special-members: __call__
 
+`LoadNumpy`
+"""""""""""
+.. autoclass:: LoadNumpy
+    :members:
+    :special-members: __call__
+
 Post-processing
 ^^^^^^^^^^^^^^^
 
@@ -570,6 +576,12 @@ IO (Dict)
     :members:
     :special-members: __call__
 
+`LoadNumpyd`
+""""""""""""
+.. autoclass:: LoadNumpyd
+    :members:
+    :special-members: __call__
+
 Post-processing (Dict)
 ^^^^^^^^^^^^^^^^^^^^^^
 
diff --git a/monai/transforms/io/array.py b/monai/transforms/io/array.py
index bcaf91f2..6ba1acff 100644
--- a/monai/transforms/io/array.py
+++ b/monai/transforms/io/array.py
@@ -17,7 +17,7 @@ from typing import Optional
 
 import numpy as np
 from torch.utils.data._utils.collate import np_str_obj_array_pattern
-
+from monai.config import KeysCollection
 from monai.data.utils import correct_nifti_header_if_necessary
 from monai.transforms.compose import Transform
 from monai.utils import optional_import, ensure_tuple
@@ -111,6 +111,8 @@ class LoadNifti(Transform):
 class LoadPNG(Transform):
     """
     Load common 2D image format (PNG, JPG, etc. using PIL) file or files from provided path.
+    If loading a list of files, stack them together and add a new dimension as first dimension,
+    and use the meta data of the first image to represent the stacked result.
     It's based on the Image module in PIL library:
     https://pillow.readthedocs.io/en/stable/reference/Image.html
     """
@@ -138,6 +140,10 @@ class LoadPNG(Transform):
             if self.dtype:
                 data = data.astype(self.dtype)
             img_array.append(data)
+
+            if self.image_only:
+                continue
+
             meta = dict()
             meta["filename_or_obj"] = name
             meta["spatial_shape"] = data.shape[:2]
@@ -146,10 +152,6 @@ class LoadPNG(Transform):
             meta["width"] = img.width
             meta["height"] = img.height
             meta["info"] = img.info
-
-            if self.image_only:
-                continue
-
             if not compatible_meta:
                 compatible_meta = meta
             else:
@@ -159,3 +161,73 @@ class LoadPNG(Transform):
 
         img_array = np.stack(img_array, axis=0) if len(img_array) > 1 else img_array[0]
         return img_array if self.image_only else (img_array, compatible_meta)
+
+
+class LoadNumpy(Transform):
+    """
+    Load arrays or pickled objects from .npy, .npz or pickled files, file or files are from provided path.
+    A typical usage is to load the `mask` data for classification task.
+    If loading a list of files or laoding npz file, stack results together and add a new dimension as first dimension,
+    and use the meta data of the first file to represent the stacked result.
+    It can load part of the npz file with specified `npz_keys`.
+    It's based on the Numpy load/read API:
+    https://numpy.org/doc/stable/reference/generated/numpy.load.html
+
+    """
+
+    def __init__(
+        self, data_only: bool = False, dtype: Optional[np.dtype] = np.float32, npz_keys: Optional[KeysCollection] = None
+    ) -> None:
+        """
+        Args:
+            data_only: if True return only the data array, otherwise return data array and metadata.
+            dtype: if not None convert the loaded data to this data type.
+            npz_keys: if loading npz file, only load the specified keys, if None, load all the items.
+                stack the loaded items together to construct a new first dimension.
+
+        """
+        self.data_only = data_only
+        self.dtype = dtype
+        if npz_keys is not None:
+            npz_keys = ensure_tuple(npz_keys)
+        self.npz_keys = npz_keys
+
+    def __call__(self, filename):
+        """
+        Args:
+            filename (str, list, tuple, file): path file or file-like object or a list of files.
+        """
+        if isinstance(filename, (tuple, list)):
+            for name in filename:
+                if name.endswith(".npz"):
+                    raise TypeError("can not load a list of npz file.")
+        filename = ensure_tuple(filename)
+        data_array = list()
+        compatible_meta = None
+
+        def _save_data_meta(data_array, name, data, compatible_meta):
+            data_array.append(data if self.dtype is None else data.astype(self.dtype))
+            if not self.data_only:
+                meta = dict()
+                meta["filename_or_obj"] = name
+                meta["spatial_shape"] = data.shape
+                if not compatible_meta:
+                    compatible_meta = meta
+                else:
+                    assert np.allclose(
+                        meta["spatial_shape"], compatible_meta["spatial_shape"]
+                    ), "all the data in the list should have same shape."
+            return compatible_meta
+
+        for name in filename:
+            data = np.load(name, allow_pickle=True)
+            if name.endswith(".npz"):
+                # load expected items from NPZ file
+                npz_keys = [f"arr_{i}" for i in range(len(data))] if self.npz_keys is None else self.npz_keys
+                for k in npz_keys:
+                    compatible_meta = _save_data_meta(data_array, name, data[k], compatible_meta)
+            else:
+                compatible_meta = _save_data_meta(data_array, name, data, compatible_meta)
+
+        data_array = np.stack(data_array, axis=0) if len(data_array) > 1 else data_array[0]
+        return data_array if self.data_only else (data_array, compatible_meta)
diff --git a/monai/transforms/io/dictionary.py b/monai/transforms/io/dictionary.py
index faf4060a..5a04b440 100644
--- a/monai/transforms/io/dictionary.py
+++ b/monai/transforms/io/dictionary.py
@@ -21,7 +21,7 @@ import numpy as np
 
 from monai.config import KeysCollection
 from monai.transforms.compose import MapTransform
-from monai.transforms.io.array import LoadNifti, LoadPNG
+from monai.transforms.io.array import LoadNifti, LoadPNG, LoadNumpy
 
 
 class LoadNiftid(MapTransform):
@@ -47,8 +47,8 @@ class LoadNiftid(MapTransform):
             keys: keys of the corresponding items to be transformed.
                 See also: :py:class:`monai.transforms.compose.MapTransform`
             as_closest_canonical: if True, load the image as closest to canonical axis format.
-            dtype: if not None convert the loaded image to this data type.
-            meta_key_postfix: use `key_{postfix}` to to store meta data of the nifti image,
+            dtype: if not None convert the loaded image data to this data type.
+            meta_key_postfix: use `key_{postfix}` to store the metadata of the nifti image,
                 default is `meta_dict`. The meta data is a dictionary object.
                 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.
@@ -95,10 +95,10 @@ class LoadPNGd(MapTransform):
         Args:
             keys: keys of the corresponding items to be transformed.
                 See also: :py:class:`monai.transforms.compose.MapTransform`
-            dtype: if not None convert the loaded image to this data type.
-            meta_key_postfix: use `key_{postfix}` to to store meta data of the nifti image,
+            dtype: if not None convert the loaded image data to this data type.
+            meta_key_postfix: use `key_{postfix}` to store the metadata of the PNG image,
                 default is `meta_dict`. The meta data is a dictionary object.
-                For example, load nifti file for `image`, store the metadata into `image_meta_dict`.
+                For example, load PNG 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.
 
@@ -127,5 +127,57 @@ class LoadPNGd(MapTransform):
         return d
 
 
+class LoadNumpyd(MapTransform):
+    """
+    Dictionary-based wrapper of :py:class:`monai.transforms.LoadNumpy`.
+    """
+
+    def __init__(
+        self,
+        keys: KeysCollection,
+        dtype: Optional[np.dtype] = np.float32,
+        npz_keys: Optional[KeysCollection] = None,
+        meta_key_postfix: str = "meta_dict",
+        overwriting: bool = False,
+    ) -> None:
+        """
+        Args:
+            keys: keys of the corresponding items to be transformed.
+                See also: :py:class:`monai.transforms.compose.MapTransform`
+            dtype: if not None convert the loaded data to this data type.
+            npz_keys: if loading npz file, only load the specified keys, if None, load all the items.
+                stack the loaded items together to construct a new first dimension.
+            meta_key_postfix: use `key_{postfix}` to store the metadata of the Numpy data,
+                default is `meta_dict`. The meta data is a dictionary object.
+                For example, load Numpy file for `mask`, store the metadata into `mask_meta_dict`.
+            overwriting: whether allow to overwrite existing meta data of same key.
+                default is False, which will raise exception if encountering existing key.
+
+        Raises:
+            ValueError: meta_key_postfix must be a string.
+
+        """
+        super().__init__(keys)
+        self.loader = LoadNumpy(data_only=False, dtype=dtype, npz_keys=npz_keys)
+        if not isinstance(meta_key_postfix, str):
+            raise ValueError("meta_key_postfix must be a string.")
+        self.meta_key_postfix = meta_key_postfix
+        self.overwriting = overwriting
+
+    def __call__(self, data):
+        d = dict(data)
+        for key in self.keys:
+            data = self.loader(d[key])
+            assert isinstance(data, (tuple, list)), "loader must return a tuple or list."
+            d[key] = data[0]
+            assert isinstance(data[1], dict), "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.")
+            d[key_to_add] = data[1]
+        return d
+
+
 LoadNiftiD = LoadNiftiDict = LoadNiftid
 LoadPNGD = LoadPNGDict = LoadPNGd
+LoadNumpyD = LoadNumpyDict = LoadNumpyd


Test patch:
diff --git a/tests/test_load_numpy.py b/tests/test_load_numpy.py
new file mode 100644
index 00000000..6332e624
--- /dev/null
+++ b/tests/test_load_numpy.py
@@ -0,0 +1,89 @@
+# Copyright 2020 MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+import os
+import shutil
+import numpy as np
+import tempfile
+from monai.transforms import LoadNumpy
+
+
+class TestLoadNumpy(unittest.TestCase):
+    def test_npy(self):
+        test_data = np.random.randint(0, 256, size=[3, 4, 4])
+        tempdir = tempfile.mkdtemp()
+        filepath = os.path.join(tempdir, "test_data.npy")
+        np.save(filepath, test_data)
+
+        result = LoadNumpy()(filepath)
+        self.assertTupleEqual(result[1]["spatial_shape"], test_data.shape)
+        self.assertTupleEqual(result[0].shape, test_data.shape)
+        np.testing.assert_allclose(result[0], test_data)
+
+        shutil.rmtree(tempdir)
+
+    def test_npz1(self):
+        test_data1 = np.random.randint(0, 256, size=[3, 4, 4])
+        tempdir = tempfile.mkdtemp()
+        filepath = os.path.join(tempdir, "test_data.npy")
+        np.save(filepath, test_data1)
+
+        result = LoadNumpy()(filepath)
+        self.assertTupleEqual(result[1]["spatial_shape"], test_data1.shape)
+        self.assertTupleEqual(result[0].shape, test_data1.shape)
+        np.testing.assert_allclose(result[0], test_data1)
+
+        shutil.rmtree(tempdir)
+
+    def test_npz2(self):
+        test_data1 = np.random.randint(0, 256, size=[3, 4, 4])
+        test_data2 = np.random.randint(0, 256, size=[3, 4, 4])
+        tempdir = tempfile.mkdtemp()
+        filepath = os.path.join(tempdir, "test_data.npz")
+        np.savez(filepath, test_data1, test_data2)
+
+        result = LoadNumpy()(filepath)
+        self.assertTupleEqual(result[1]["spatial_shape"], test_data1.shape)
+        self.assertTupleEqual(result[0].shape, (2, 3, 4, 4))
+        np.testing.assert_allclose(result[0], np.stack([test_data1, test_data2]))
+
+        shutil.rmtree(tempdir)
+
+    def test_npz3(self):
+        test_data1 = np.random.randint(0, 256, size=[3, 4, 4])
+        test_data2 = np.random.randint(0, 256, size=[3, 4, 4])
+        tempdir = tempfile.mkdtemp()
+        filepath = os.path.join(tempdir, "test_data.npz")
+        np.savez(filepath, test1=test_data1, test2=test_data2)
+
+        result = LoadNumpy(npz_keys=["test1", "test2"])(filepath)
+        self.assertTupleEqual(result[1]["spatial_shape"], test_data1.shape)
+        self.assertTupleEqual(result[0].shape, (2, 3, 4, 4))
+        np.testing.assert_allclose(result[0], np.stack([test_data1, test_data2]))
+
+        shutil.rmtree(tempdir)
+
+    def test_npy_pickle(self):
+        test_data = {"test": np.random.randint(0, 256, size=[3, 4, 4])}
+        tempdir = tempfile.mkdtemp()
+        filepath = os.path.join(tempdir, "test_data.npy")
+        np.save(filepath, test_data, allow_pickle=True)
+
+        result = LoadNumpy(data_only=True, dtype=None)(filepath).item()
+        self.assertTupleEqual(result["test"].shape, test_data["test"].shape)
+        np.testing.assert_allclose(result["test"], test_data["test"])
+
+        shutil.rmtree(tempdir)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/test_load_numpyd.py b/tests/test_load_numpyd.py
new file mode 100644
index 00000000..4e68e1b3
--- /dev/null
+++ b/tests/test_load_numpyd.py
@@ -0,0 +1,89 @@
+# Copyright 2020 MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+import os
+import shutil
+import numpy as np
+import tempfile
+from monai.transforms import LoadNumpyd
+
+
+class TestLoadNumpyd(unittest.TestCase):
+    def test_npy(self):
+        test_data = np.random.randint(0, 256, size=[3, 4, 4])
+        tempdir = tempfile.mkdtemp()
+        filepath = os.path.join(tempdir, "test_data.npy")
+        np.save(filepath, test_data)
+
+        result = LoadNumpyd(keys="mask")({"mask": filepath})
+        self.assertTupleEqual(result["mask_meta_dict"]["spatial_shape"], test_data.shape)
+        self.assertTupleEqual(result["mask"].shape, test_data.shape)
+        np.testing.assert_allclose(result["mask"], test_data)
+
+        shutil.rmtree(tempdir)
+
+    def test_npz1(self):
+        test_data1 = np.random.randint(0, 256, size=[3, 4, 4])
+        tempdir = tempfile.mkdtemp()
+        filepath = os.path.join(tempdir, "test_data.npy")
+        np.save(filepath, test_data1)
+
+        result = LoadNumpyd(keys="mask")({"mask": filepath})
+        self.assertTupleEqual(result["mask_meta_dict"]["spatial_shape"], test_data1.shape)
+        self.assertTupleEqual(result["mask"].shape, test_data1.shape)
+        np.testing.assert_allclose(result["mask"], test_data1)
+
+        shutil.rmtree(tempdir)
+
+    def test_npz2(self):
+        test_data1 = np.random.randint(0, 256, size=[3, 4, 4])
+        test_data2 = np.random.randint(0, 256, size=[3, 4, 4])
+        tempdir = tempfile.mkdtemp()
+        filepath = os.path.join(tempdir, "test_data.npz")
+        np.savez(filepath, test_data1, test_data2)
+
+        result = LoadNumpyd(keys="mask")({"mask": filepath})
+        self.assertTupleEqual(result["mask_meta_dict"]["spatial_shape"], test_data1.shape)
+        self.assertTupleEqual(result["mask"].shape, (2, 3, 4, 4))
+        np.testing.assert_allclose(result["mask"], np.stack([test_data1, test_data2]))
+
+        shutil.rmtree(tempdir)
+
+    def test_npz3(self):
+        test_data1 = np.random.randint(0, 256, size=[3, 4, 4])
+        test_data2 = np.random.randint(0, 256, size=[3, 4, 4])
+        tempdir = tempfile.mkdtemp()
+        filepath = os.path.join(tempdir, "test_data.npz")
+        np.savez(filepath, test1=test_data1, test2=test_data2)
+
+        result = LoadNumpyd(keys="mask", npz_keys=["test1", "test2"])({"mask": filepath})
+        self.assertTupleEqual(result["mask_meta_dict"]["spatial_shape"], test_data1.shape)
+        self.assertTupleEqual(result["mask"].shape, (2, 3, 4, 4))
+        np.testing.assert_allclose(result["mask"], np.stack([test_data1, test_data2]))
+
+        shutil.rmtree(tempdir)
+
+    def test_npy_pickle(self):
+        test_data = {"test": np.random.randint(0, 256, size=[3, 4, 4])}
+        tempdir = tempfile.mkdtemp()
+        filepath = os.path.join(tempdir, "test_data.npy")
+        np.save(filepath, test_data, allow_pickle=True)
+
+        result = LoadNumpyd(keys="mask", dtype=None)({"mask": filepath})["mask"].item()
+        self.assertTupleEqual(result["test"].shape, test_data["test"].shape)
+        np.testing.assert_allclose(result["test"], test_data["test"])
+
+        shutil.rmtree(tempdir)
+
+
+if __name__ == "__main__":
+    unittest.main()