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

Code before:
        self.pad_mode = pad_mode
        self.pad_kwargs = pad_kwargs
        self.overlap = overlap
        self.num_patches = num_patches
        self.sort_fn = sort_fn.lower() if sort_fn else None
        self.threshold = threshold

    def filter_threshold(self, image_np: NdarrayOrTensor, locations: np.ndarray) -> tuple[NdarrayOrTensor, np.ndarray]:
        """
        Filter the patches and their locations according to a threshold.

        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.

        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))))
            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).

Review comment (suggestion):
```suggestion
        Sort the patches based on the sum of their intensity, and just keep `self.num_patches` of them.
```

Code after:
        self.pad_mode = pad_mode
        self.pad_kwargs = pad_kwargs
        self.overlap = overlap
        self.num_patches = num_patches
        self.sort_fn = sort_fn.lower() if sort_fn else None
        self.threshold = threshold

    def filter_threshold(self, image_np: NdarrayOrTensor, locations: np.ndarray) -> tuple[NdarrayOrTensor, np.ndarray]:
        """
        Filter the patches and their locations according to a threshold.

        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.

        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, 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).