4060 Add PatchIter and PatchIterd transform (#4061)

* [DLMED] change PatchIter to be a transform

Signed-off-by: Nic Ma <[email protected]>

* [DLMED] add dict transform

Signed-off-by: Nic Ma <[email protected]>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* [DLMED] add unit tests

Signed-off-by: Nic Ma <[email protected]>

* [DLMED] store coords in dict

Signed-off-by: Nic Ma <[email protected]>

* [DLMED] update according to comments

Signed-off-by: Nic Ma <[email protected]>

* [DLMED] restore the doc-string

Signed-off-by: Nic Ma <[email protected]>

* [DLMED] store more info

Signed-off-by: Nic Ma <[email protected]>

* [DLMED] update according to comments

Signed-off-by: Nic Ma <[email protected]>

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: Nic-Ma

docs/source/data.rst

@@ -107,16 +107,23 @@ Patch-based dataset
 .. autoclass:: GridPatchDataset
   :members:
 
-`PatchIter`
-~~~~~~~~~~~
-.. autoclass:: PatchIter
-  :members:
-
 `PatchDataset`
 ~~~~~~~~~~~~~~
 .. autoclass:: PatchDataset
   :members:
 
+`PatchIter`
+"""""""""""
+.. autoclass:: PatchIter
+    :members:
+    :special-members: __call__
+
+`PatchIterd`
+""""""""""""
+.. autoclass:: PatchIterd
+    :members:
+    :special-members: __call__
+
 Image reader
 ------------
 

monai/data/__init__.py

@@ -32,7 +32,7 @@
     load_decathlon_properties,
 )
 from .folder_layout import FolderLayout
-from .grid_dataset import GridPatchDataset, PatchDataset, PatchIter
+from .grid_dataset import GridPatchDataset, PatchDataset, PatchIter, PatchIterd
 from .image_dataset import ImageDataset
 from .image_reader import ImageReader, ITKReader, NibabelReader, NumpyReader, PILReader, WSIReader
 from .image_writer import (

monai/data/grid_dataset.py

@@ -9,21 +9,26 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Callable, Dict, Iterable, Optional, Sequence, Union
+from copy import deepcopy
+from typing import Callable, Dict, Hashable, Iterable, Mapping, Optional, Sequence, Union
 
+import numpy as np
+
+from monai.config import KeysCollection
 from monai.data.dataset import Dataset
 from monai.data.iterable_dataset import IterableDataset
 from monai.data.utils import iter_patch
 from monai.transforms import apply_transform
-from monai.utils import NumpyPadMode, deprecated_arg, ensure_tuple, look_up_option
+from monai.utils import NumpyPadMode, deprecated_arg, ensure_tuple, first, look_up_option
 
-__all__ = ["PatchDataset", "GridPatchDataset", "PatchIter"]
+__all__ = ["PatchDataset", "GridPatchDataset", "PatchIter", "PatchIterd"]
 
 
 class PatchIter:
     """
-    A class to return a patch generator with predefined properties such as `patch_size`.
+    Return a patch generator with predefined properties such as `patch_size`.
     Typically used with :py:class:`monai.data.GridPatchDataset`.
+
     """
 
     def __init__(
@@ -42,7 +47,8 @@ def __init__(
                 ``"median"``, ``"minimum"``, ``"reflect"``, ``"symmetric"``, ``"wrap"``, ``"empty"``}
                 One of the listed string values or a user supplied function. Defaults to ``"wrap"``.
                 See also: https://numpy.org/doc/1.18/reference/generated/numpy.pad.html
-            pad_opts: padding options, see numpy.pad
+            pad_opts: other arguments for the `np.pad` function.
+                note that `np.pad` treats channel dimension as the first dimension.
 
         Note:
             The `patch_size` is the size of the
@@ -52,37 +58,89 @@ def __init__(
             specified by a `patch_size` of (10, 10, 10).
 
         """
-        self.patch_size = (None,) + tuple(patch_size)
+        self.patch_size = (None,) + tuple(patch_size)  # expand to have the channel dim
         self.start_pos = ensure_tuple(start_pos)
         self.mode: NumpyPadMode = look_up_option(mode, NumpyPadMode)
         self.pad_opts = pad_opts
 
-    def __call__(self, array):
+    def __call__(self, array: np.ndarray):
         """
         Args:
             array: the image to generate patches from.
+
         """
         yield from iter_patch(
             array,
-            patch_size=self.patch_size,  # expand to have the channel dim
+            patch_size=self.patch_size,  # type: ignore
             start_pos=self.start_pos,
             copy_back=False,
             mode=self.mode,
             **self.pad_opts,
         )
 
 
+class PatchIterd:
+    """
+    Dictionary-based wrapper of :py:class:`monai.data.PatchIter`.
+    Return a patch generator for dictionary data and the coordinate, Typically used
+    with :py:class:`monai.data.GridPatchDataset`.
+    Suppose all the expected fields specified by `keys` have same shape.
+
+    Args:
+        keys: keys of the corresponding items to iterate patches.
+        patch_size: size of patches to generate slices for, 0/None selects whole dimension
+        start_pos: starting position in the array, default is 0 for each dimension
+        mode: {``"constant"``, ``"edge"``, ``"linear_ramp"``, ``"maximum"``, ``"mean"``,
+            ``"median"``, ``"minimum"``, ``"reflect"``, ``"symmetric"``, ``"wrap"``, ``"empty"``}
+            One of the listed string values or a user supplied function. Defaults to ``"wrap"``.
+            See also: https://numpy.org/doc/1.18/reference/generated/numpy.pad.html
+        pad_opts: other arguments for the `np.pad` function.
+            note that `np.pad` treats channel dimension as the first dimension.
+
+    """
+
+    coords_key = "patch_coords"
+    original_spatial_shape_key = "original_spatial_shape"
+    start_pos_key = "start_pos"
+
+    def __init__(
+        self,
+        keys: KeysCollection,
+        patch_size: Sequence[int],
+        start_pos: Sequence[int] = (),
+        mode: Union[NumpyPadMode, str] = NumpyPadMode.WRAP,
+        **pad_opts,
+    ):
+        self.keys = ensure_tuple(keys)
+        self.patch_iter = PatchIter(patch_size=patch_size, start_pos=start_pos, mode=mode, **pad_opts)
+
+    def __call__(self, data: Mapping[Hashable, np.ndarray]):
+        d = dict(data)
+        original_spatial_shape = d[first(self.keys)].shape[1:]
+
+        for patch in zip(*[self.patch_iter(d[key]) for key in self.keys]):
+            coords = patch[0][1]  # use the coordinate of the first item
+            ret = {k: v[0] for k, v in zip(self.keys, patch)}
+            # fill in the extra keys with unmodified data
+            for k in set(d.keys()).difference(set(self.keys)):
+                ret[k] = deepcopy(d[k])
+            # also store the `coordinate`, `spatial shape of original image`, `start position` in the dictionary
+            ret[self.coords_key] = coords
+            ret[self.original_spatial_shape_key] = original_spatial_shape
+            ret[self.start_pos_key] = self.patch_iter.start_pos
+            yield ret, coords
+
+
 class GridPatchDataset(IterableDataset):
     """
-    Yields patches from images read from an image dataset.
-    Typically used with `PatchIter` so that the patches are chosen in a contiguous grid sampling scheme.
+    Yields patches from data read from an image dataset.
+    Typically used with `PatchIter` or `PatchIterd` so that the patches are chosen in a contiguous grid sampling scheme.
 
      .. code-block:: python
 
         import numpy as np
 
-        from monai.data import GridPatchDataset, DataLoader, PatchIter
-        from monai.transforms import RandShiftIntensity
+        from monai.data import GridPatchDataset, DataLoader, PatchIter, RandShiftIntensity
 
         # image-level dataset
         images = [np.arange(16, dtype=float).reshape(1, 4, 4),
@@ -109,7 +167,7 @@ class GridPatchDataset(IterableDataset):
         data: the data source to read image data from.
         patch_iter: converts an input image (item from dataset) into a iterable of image patches.
             `patch_iter(dataset[idx])` must yield a tuple: (patches, coordinates).
-            see also: :py:class:`monai.data.PatchIter`.
+            see also: :py:class:`monai.data.PatchIter` or :py:class:`monai.data.PatchIterd`.
         transform: a callable data transform operates on the patches.
         with_coordinates: whether to yield the coordinates of each patch, default to `True`.
 

monai/transforms/croppad/dictionary.py

@@ -70,6 +70,7 @@
     "RandCropByPosNegLabeld",
     "ResizeWithPadOrCropd",
     "BoundingRectd",
+    "RandCropByLabelClassesd",
     "SpatialPadD",
     "SpatialPadDict",
     "BorderPadD",
@@ -98,7 +99,6 @@
     "ResizeWithPadOrCropDict",
     "BoundingRectD",
     "BoundingRectDict",
-    "RandCropByLabelClassesd",
     "RandCropByLabelClassesD",
     "RandCropByLabelClassesDict",
 ]

tests/test_grid_dataset.py

@@ -14,8 +14,8 @@
 
 import numpy as np
 
-from monai.data import DataLoader, GridPatchDataset, PatchIter
-from monai.transforms import RandShiftIntensity
+from monai.data import DataLoader, GridPatchDataset, PatchIter, PatchIterd
+from monai.transforms import RandShiftIntensity, RandShiftIntensityd
 from monai.utils import set_determinism
 
 
@@ -76,6 +76,48 @@ def test_loading_array(self):
                 item[1], np.array([[[0, 1], [0, 2], [2, 4]], [[0, 1], [2, 4], [2, 4]]]), rtol=1e-5
             )
 
+    def test_loading_dict(self):
+        set_determinism(seed=1234)
+        # test sequence input data with dict
+        data = [
+            {
+                "image": np.arange(16, dtype=float).reshape(1, 4, 4),
+                "label": np.arange(16, dtype=float).reshape(1, 4, 4),
+                "metadata": "test string",
+            },
+            {
+                "image": np.arange(16, dtype=float).reshape(1, 4, 4),
+                "label": np.arange(16, dtype=float).reshape(1, 4, 4),
+                "metadata": "test string",
+            },
+        ]
+        # image level
+        patch_intensity = RandShiftIntensityd(keys="image", offsets=1.0, prob=1.0)
+        patch_iter = PatchIterd(keys=["image", "label"], patch_size=(2, 2), start_pos=(0, 0))
+        ds = GridPatchDataset(data=data, patch_iter=patch_iter, transform=patch_intensity, with_coordinates=True)
+        # use the grid patch dataset
+        for item in DataLoader(ds, batch_size=2, shuffle=False, num_workers=0):
+            np.testing.assert_equal(item[0]["image"].shape, (2, 1, 2, 2))
+            np.testing.assert_equal(item[0]["label"].shape, (2, 1, 2, 2))
+            self.assertListEqual(item[0]["metadata"], ["test string", "test string"])
+        np.testing.assert_allclose(
+            item[0]["image"],
+            np.array([[[[1.4965, 2.4965], [5.4965, 6.4965]]], [[[11.3584, 12.3584], [15.3584, 16.3584]]]]),
+            rtol=1e-4,
+        )
+        np.testing.assert_allclose(item[1], np.array([[[0, 1], [0, 2], [2, 4]], [[0, 1], [2, 4], [2, 4]]]), rtol=1e-5)
+        if sys.platform != "win32":
+            for item in DataLoader(ds, batch_size=2, shuffle=False, num_workers=2):
+                np.testing.assert_equal(item[0]["image"].shape, (2, 1, 2, 2))
+            np.testing.assert_allclose(
+                item[0]["image"],
+                np.array([[[[1.2548, 2.2548], [5.2548, 6.2548]]], [[[9.1106, 10.1106], [13.1106, 14.1106]]]]),
+                rtol=1e-3,
+            )
+            np.testing.assert_allclose(
+                item[1], np.array([[[0, 1], [0, 2], [2, 4]], [[0, 1], [2, 4], [2, 4]]]), rtol=1e-5
+            )
+
 
 if __name__ == "__main__":
     unittest.main()