Repository: Project-MONAI/MONAI
File: monai/transforms/spatial/array.py (Python)

Code before:
        Returns:
            tuple[NdarrayOrTensor, numpy.ndarray]:  tuple of filtered patches and locations.
        """
        n_dims = len(image_np.shape)
        idx = argwhere(image_np.sum(tuple(range(1, n_dims))) < self.threshold).reshape(-1)
        idx_np = convert_data_type(idx, np.ndarray)[0]
        return image_np[idx], locations[idx_np]

    def filter_count(self, image_np: NdarrayOrTensor, locations: np.ndarray) -> tuple[NdarrayOrTensor, np.ndarray]:
        """
        Sort the patches based on the sum of their intensity or in random order, and just keep `self.num_patches` of them.

        Args:
            image_np: a numpy.ndarray or torch.Tensor representing a stack of patches.
            locations: a numpy.ndarray representing the stack of location of each patch.
        """
        if self.sort_fn is None:
            image_np = image_np[: self.num_patches]
            locations = locations[: self.num_patches]
        elif self.num_patches is not None:
            n_dims = len(image_np.shape)
            if self.sort_fn == GridPatchSort.MIN:
                idx = argsort(image_np.sum(tuple(range(1, n_dims))))
            elif self.sort_fn == GridPatchSort.MAX:
                idx = argsort(-image_np.sum(tuple(range(1, n_dims))))
            elif self.sort_fn == GridPatchSort.RANDOM:
                idx = self.patch_idx
            else:
                raise ValueError(f'`sort_fn` should be either "min", "max", or None! {self.sort_fn} provided!')
            idx = idx[: self.num_patches]
            idx_np = convert_data_type(idx, np.ndarray)[0]
            image_np = image_np[idx]
            locations = locations[idx_np]
        return image_np, locations

    def __call__(self, array: NdarrayOrTensor) -> MetaTensor:
        """
        Extract the patches (sweeping the entire image in a row-major sliding-window manner with possible overlaps).

        Args:
            array: a input image as `numpy.ndarray` or `torch.Tensor`

        Return:
            MetaTensor: the extracted patches as a single tensor (with patch dimension as the first dimension),
                with defined `PatchKeys.LOCATION` and `PatchKeys.COUNT` metadata.
        """
        # create the patch iterator which sweeps the image row-by-row
        patch_iterator = iter_patch(
            array,
            patch_size=(None,) + self.patch_size,  # expand to have the channel dim
            start_pos=(0,) + self.offset,  # expand to have the channel dim

Review comment (suggestion):
Hi @jmnolte, thanks for your quick update!
If you rewrite the `filter_count` method here, it might be a little weird. Theoretically we don't want the `GridPatch` to support random operations, so based on @drbeh  [comment](https://github.com/Project-MONAI/MONAI/pull/6701#pullrequestreview-1522897921), you should probably override this method inside `RandGridPatch`.
Free feel to ask if you have any question.
Thanks!

Code after:
            tuple[NdarrayOrTensor, numpy.ndarray]:  tuple of filtered patches and locations.
        """
        n_dims = len(image_np.shape)
        idx = argwhere(image_np.sum(tuple(range(1, n_dims))) < self.threshold).reshape(-1)
        idx_np = convert_data_type(idx, np.ndarray)[0]
        return image_np[idx], locations[idx_np]

    def filter_count(self, image_np: NdarrayOrTensor, locations: np.ndarray) -> tuple[NdarrayOrTensor, np.ndarray]:
        """
        Sort the patches based on the sum of their intensity, and just keep `self.num_patches` of them.

        Args:
            image_np: a numpy.ndarray or torch.Tensor representing a stack of patches.
            locations: a numpy.ndarray representing the stack of location of each patch.
        """
        if self.sort_fn is None:
            image_np = image_np[: self.num_patches]
            locations = locations[: self.num_patches]
        elif self.num_patches is not None:
            n_dims = len(image_np.shape)
            if self.sort_fn == GridPatchSort.MIN:
                idx = argsort(image_np.sum(tuple(range(1, n_dims))))
            elif self.sort_fn == GridPatchSort.MAX:
                idx = argsort(-image_np.sum(tuple(range(1, n_dims))))
            else:
                raise ValueError(f'`sort_fn` should be either "min", "max", or None! {self.sort_fn} provided!')
            idx = idx[: self.num_patches]
            idx_np = convert_data_type(idx, np.ndarray)[0]
            image_np = image_np[idx]
            locations = locations[idx_np]
        return image_np, locations

    def __call__(self, array: NdarrayOrTensor) -> MetaTensor:
        """
        Extract the patches (sweeping the entire image in a row-major sliding-window manner with possible overlaps).

        Args:
            array: a input image as `numpy.ndarray` or `torch.Tensor`

        Return:
            MetaTensor: the extracted patches as a single tensor (with patch dimension as the first dimension),
                with defined `PatchKeys.LOCATION` and `PatchKeys.COUNT` metadata.
        """
        # create the patch iterator which sweeps the image row-by-row
        patch_iterator = iter_patch(
            array,
            patch_size=(None,) + self.patch_size,  # expand to have the channel dim
            start_pos=(0,) + self.offset,  # expand to have the channel dim
            overlap=self.overlap,
            copy_back=False,
            mode=self.pad_mode,