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

Code before:
    return img.reshape(orig_shape)


class Fourier:
    """
    Helper class storing Fourier mappings
    """

    @staticmethod
    def shift_fourier(x: NdarrayOrTensor, spatial_dims: int) -> NdarrayOrTensor:
        """
        Applies fourier transform and shifts the zero-frequency component to the
        center of the spectrum. Only the spatial dimensions get transformed.

        Args:
            x: Image to transform.
            spatial_dims: Number of spatial dimensions.

        Returns
            k: K-space data.
        """
        dims = tuple(range(-spatial_dims, 0))
        k: NdarrayOrTensor
        if isinstance(x, torch.Tensor):
            if hasattr(torch.fft, "fftshift"):  # `fftshift` is new in torch 1.8.0
                k = torch.fft.fftshift(torch.fft.fftn(x, dim=dims), dim=dims).contiguous()
            else:
                # if using old PyTorch, will convert to numpy array and return
                k = np.ascontiguousarray(np.fft.fftshift(np.fft.fftn(x.cpu().numpy(), axes=dims), axes=dims))
        else:
            k = np.ascontiguousarray(np.fft.fftshift(np.fft.fftn(x, axes=dims), axes=dims))
        return k

    @staticmethod
    def inv_shift_fourier(k: NdarrayOrTensor, spatial_dims: int, n_dims: int | None = None) -> NdarrayOrTensor:
        """
        Applies inverse shift and fourier transform. Only the spatial
        dimensions are transformed.

        Args:
            k: K-space data.
            spatial_dims: Number of spatial dimensions.

        Returns:
            x: Tensor in image space.
        """
        dims = tuple(range(-spatial_dims, 0))
        out: NdarrayOrTensor
        if isinstance(k, torch.Tensor):
            if hasattr(torch.fft, "ifftshift"):  # `ifftshift` is new in torch 1.8.0
                out = torch.fft.ifftn(torch.fft.ifftshift(k, dim=dims), dim=dims, norm="backward").real.contiguous()

Review comment (performance):
Thanks for pointing and your PR!
Calling contiguous() results in a new tensor, meaning memory copying. For large tensors, this can increase memory consumption and computational overhead.
If contiguous() is called frequently, especially on large datasets, it can impact performance. Therefore, use this method only when necessary.

I would prefer if we can add an arg here `as_contiguous` such as what we did here: https://github.com/Project-MONAI/MONAI/blob/f1ef3e88f4bc594779c0c4188883f78bf6d1efff/monai/data/grid_dataset.py#L214

And also, we can use `ascontiguousarray` to only change one line during return.
https://github.com/Project-MONAI/MONAI/blob/f1ef3e88f4bc594779c0c4188883f78bf6d1efff/monai/transforms/utils_pytorch_numpy_unification.py#L398

What do you think?
cc @Nic-Ma @ericspod 

Code after:
    rescale_max: float = 1.0,
) -> torch.Tensor:
    """
    Please refer to :py:class:`monai.transforms.AddExtremePointsChannel` for the usage.

    Applies a gaussian filter to the extreme points image. Then the pixel values in points image are rescaled
    to range [rescale_min, rescale_max].

    Args:
        points: Extreme points of the object/organ.
        label: label image to get extreme points from. Shape must be
            (1, spatial_dim1, [, spatial_dim2, ...]). Doesn't support one-hot labels.
        sigma: if a list of values, must match the count of spatial dimensions of input data,
            and apply every value in the list to 1 spatial dimension. if only 1 value provided,
            use it for all spatial dimensions.
        rescale_min: minimum value of output data.
        rescale_max: maximum value of output data.
    """
    # points to image
    # points_image = torch.zeros(label.shape[1:], dtype=torch.float)
    points_image = torch.zeros_like(torch.as_tensor(label[0]), dtype=torch.float)
    for p in points:
        points_image[p] = 1.0

    if isinstance(sigma, Sequence):
        sigma = [torch.as_tensor(s, device=points_image.device) for s in sigma]
    else:
        sigma = torch.as_tensor(sigma, device=points_image.device)

    # add channel and add batch
    points_image = points_image.unsqueeze(0).unsqueeze(0)
    gaussian_filter = GaussianFilter(label.ndim - 1, sigma=sigma)
    points_image = gaussian_filter(points_image).squeeze(0).detach()

    # rescale the points image to [rescale_min, rescale_max]
    min_intensity = points_image.min()
    max_intensity = points_image.max()
    points_image = (points_image - min_intensity) / (max_intensity - min_intensity)
    return points_image * (rescale_max - rescale_min) + rescale_min


def map_spatial_axes(
    img_ndim: int, spatial_axes: Sequence[int] | int | None = None, channel_first: bool = True
) -> list[int]:
    """
    Utility to map the spatial axes to real axes in channel first/last shape.
    For example:
    If `channel_first` is True, and `img` has 3 spatial dims, map spatial axes to real axes as below:
    None -> [1, 2, 3]
    [0, 1] -> [1, 2]
    [0, -1] -> [1, -1]