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

Code before:
        grads = convert_to_dst_type(grads.squeeze(0), image_tensor)[0]

        return grads


class DistanceTransformEDT(Transform):
    """
    Applies the Euclidean distance transform on the input.

    Either GPU based with CuPy / cuCIM or CPU based with scipy.ndimage.
    Choice only depends on cuCIM being available.
    Note that the calculations can deviate, for details look into the cuCIM about distance_transform_edt().
    """

    backend = [TransformBackends.NUMPY, TransformBackends.CUPY]

    def __init__(self, sampling: None | float | list[float] = None, force_scipy: bool = False) -> None:
        super().__init__()
        self.force_scipy = force_scipy
        self.sampling = sampling

    def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor:
        """
        Args:
            img: Input image on which the distance transform shall be run.
                For CuPy shape must be ([channel, ]spatial_dim1, spatial_dim2 [, spatial_dim3]), so 2D, 3D and 4D input only.
                If you need to run the transform on other shapes, use the ``force_scipy`` flag.
                4D input gets passed channel-wise to cupy.

        Returns:
            An array with the same shape and data type as img
        """
        return distance_transform_edt(img=img, sampling=self.sampling, force_scipy=self.force_scipy)


Review comment (suggestion):
```suggestion
                channel first array, must have shape: (num_channels, H[, W, ..., ])
```

Code after:
        grads = convert_to_dst_type(grads.squeeze(0), image_tensor)[0]

        return grads


class DistanceTransformEDT(Transform):
    """
    Applies the Euclidean distance transform on the input.
    Either GPU based with CuPy / cuCIM or CPU based with scipy.
    To use the GPU implementation, make sure cuCIM is available and that the data is a `torch.tensor` on a GPU device.

    Note that the results of the libraries can differ, so stick to one if possible.
    For details, check out the `SciPy`_ and `cuCIM`_ documentation and / or :func:`monai.transforms.utils.distance_transform_edt`.

    .. _SciPy: https://docs.scipy.org/doc/scipy/reference/generated/scipy.ndimage.distance_transform_edt.html
    .. _cuCIM: https://docs.rapids.ai/api/cucim/nightly/api/#cucim.core.operations.morphology.distance_transform_edt
    """

    backend = [TransformBackends.NUMPY, TransformBackends.CUPY]

    def __init__(self, sampling: None | float | list[float] = None) -> None:
        super().__init__()
        self.sampling = sampling

    def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor:
        """
        Args:
            img: Input image on which the distance transform shall be run.
                Has to be a channel first array, must have shape: (num_channels, H, W [,D]).
                Can be of any type but will be converted into binary: 1 wherever image equates to True, 0 elsewhere.
                Input gets passed channel-wise to the distance-transform, thus results from this function will differ
                from directly calling ``distance_transform_edt()`` in CuPy or SciPy.
            sampling: Spacing of elements along each dimension. If a sequence, must be of length equal to the input rank -1;
                if a single number, this is used for all axes. If not specified, a grid spacing of unity is implied.

        Returns:
            An array with the same shape and data type as img
        """
        return distance_transform_edt(img=img, sampling=self.sampling)  # type: ignore