Repository: Project-MONAI/MONAI
File: monai/data/grid_dataset.py (Python)

Code before:
                    raise ValueError("transform must be an instance of monai.transforms.Compose.")

                first_random = self.patch_transform.get_index_of_first(
                    lambda t: isinstance(t, RandomizableTrait) or not isinstance(t, Transform)
                )
                if first_random is not None:
                    data = deepcopy(data) if self.copy_cache is True else data
                    for out_patch, others in zip(data, other):
                        if self.patch_transform is not None:
                            out_patch = self.patch_transform(out_patch, start=first_random)
                        if (
                            self.with_coordinates and len(others) > 0
                        ):  # patch_iter to yield at least 2 items: patch, coords
                            yield out_patch, others
                        else:
                            yield out_patch
        else:
            for image in super().__iter__():
                for patch, *others in self.patch_iter(image):
                    out_patch = patch
                    if self.patch_transform is not None:
                        out_patch = apply_transform(self.patch_transform, patch, map_items=False)
                    if self.with_coordinates and len(others) > 0:  # patch_iter to yield at least 2 items: patch, coords
                        yield out_patch, others[0]
                    else:
                        yield out_patch


class PatchDataset(IterableDataset):
    """
    returns a patch from an image dataset.
    The patches are generated by a user-specified callable `patch_func`,
    and are optionally post-processed by `transform`.
    For example, to generate random patch samples from an image dataset:

    .. code-block:: python

        import numpy as np

        from monai.data import PatchDataset, DataLoader
        from monai.transforms import RandSpatialCropSamples, RandShiftIntensity

        # image dataset
        images = [np.arange(16, dtype=float).reshape(1, 4, 4),
                  np.arange(16, dtype=float).reshape(1, 4, 4)]
        # image patch sampler
        n_samples = 5
        sampler = RandSpatialCropSamples(roi_size=(3, 3), num_samples=n_samples,
                                         random_center=True, random_size=False)
        # patch-level intensity shifts
        patch_intensity = RandShiftIntensity(offsets=1.0, prob=1.0)

Review comment (suggestion):
This loop appears multiple times in essentially the same form. We could define the following private method to be used here and wherever else we're applying transforms to generate patches:

```python
def _generate_patches(self, src, **apply_args):
    for patch, *others in src:
        out_patch = patch
        if self.patch_transform is not None:
            out_patch = apply_transform(self.patch_transform, patch, **apply_args)
        if self.with_coordinates and len(others) > 0:  # patch_iter to yield at least 2 items: patch, coords
            yield out_patch, others[0]
        else:
            yield out_patch
```

I think this captures the commonality in all the places where we're yield patches after transforms are applied. You'd have to double check my logic but I think we can reduce repeating the same pattern and make it something possibly to use later in subclasses.

Code after:

    def __iter__(self):
        if self.cache:
            cache_index = None
            for image in super().__iter__():
                key = self.hash_func(image)
                if key in self._hash_keys:
                    # if existing in cache, try to get the index in cache
                    cache_index = self._hash_keys.index(key)
                if cache_index is None:
                    # no cache for this index, execute all the transforms directly
                    yield from self._generate_patches(self.patch_iter(image))
                else:
                    if self._cache is None:
                        raise RuntimeError(
                            "Cache buffer is not initialized, please call `set_data()` before epoch begins."
                        )
                    data = self._cache[cache_index]  # type: ignore
                    other = self._cache_other[cache_index]  # type: ignore

                    # load data from cache and execute from the first random transform
                    data = deepcopy(data) if self.copy_cache else data
                    yield from self._generate_patches(zip(data, other), start=self.first_random)
        else:
            for image in super().__iter__():
                yield from self._generate_patches(self.patch_iter(image))


class PatchDataset(IterableDataset):
    """
    Yields patches from data read from an image dataset.
    The patches are generated by a user-specified callable `patch_func`,
    and are optionally post-processed by `transform`.
    For example, to generate random patch samples from an image dataset:

    .. code-block:: python

        import numpy as np

        from monai.data import PatchDataset, DataLoader
        from monai.transforms import RandSpatialCropSamples, RandShiftIntensity

        # image dataset
        images = [np.arange(16, dtype=float).reshape(1, 4, 4),
                  np.arange(16, dtype=float).reshape(1, 4, 4)]
        # image patch sampler
        n_samples = 5
        sampler = RandSpatialCropSamples(roi_size=(3, 3), num_samples=n_samples,
                                         random_center=True, random_size=False)
        # patch-level intensity shifts
        patch_intensity = RandShiftIntensity(offsets=1.0, prob=1.0)