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

Code before:
def distance_transform_edt(
    img: NdarrayOrTensor, sampling: None | float | list[float] = None, force_scipy: bool = False
) -> NdarrayOrTensor:
    """
    Euclidean distance transform, either GPU based with CuPy / cuCIM
    or CPU based with scipy.ndimage.
    Choice depends on cuCIM being available or scipy can be forced with the ``force_scipy`` flag.

    Note that the runtime running on the CPU may be really depending on the inputs size.

    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.
        sampling: Spacing of elements along each dimension. If a sequence, must be of length equal to the input rank;
            if a single number, this is used for all axes. If not specified, a grid spacing of unity is implied.
        force_scipy: Force the CPU based scipy implementation of the euclidean distance transform
    """
    distance_transform_edt, has_cucim = optional_import(
        "cucim.core.operations.morphology", name="distance_transform_edt"
    )

    if has_cp and has_cucim and not force_scipy:
        img_ = convert_to_cupy(img)
        if img_.ndim == 4:
            out = []
            for channel in img_:
                out.append(distance_transform_edt(channel, sampling=sampling))
            distance = cp.stack(out)
        else:
            distance = distance_transform_edt(img_, sampling=sampling)
    else:
        if not has_ndimage:
            raise RuntimeError("scipy.ndimage required if cupy is not available")
        img_ = convert_to_numpy(img)
        distance = ndimage.distance_transform_edt(img_, sampling=sampling)

    return convert_to_dst_type(distance, dst=img, dtype=distance.dtype)[0]


if __name__ == "__main__":
    print_transform_backends()


Review comment (suggestion):
For a 4D Tensor, this will lead to different results depending on whether you use force_scipy is True or False. Or even more important depending on whether you have CUDA available or not. I think we either need to use this logic for both cases or not support any input with dims > 3 here at all (see https://github.com/Project-MONAI/MONAI/pull/6981#issuecomment-1725969539) .

Code after:
def distance_transform_edt(
    img: NdarrayOrTensor,
    sampling: None | float | list[float] = None,
    return_distances: bool = True,
    return_indices: bool = False,
    distances: NdarrayOrTensor | None = None,
    indices: NdarrayOrTensor | None = None,
    *,
    block_params: tuple[int, int, int] | None = None,
    float64_distances: bool = False,
) -> None | NdarrayOrTensor | tuple[NdarrayOrTensor, NdarrayOrTensor]:
    """
    Euclidean distance transform, 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.

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

    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.
        return_distances: Whether to calculate the distance transform.
        return_indices: Whether to calculate the feature transform.
        distances: An output array to store the calculated distance transform, instead of returning it.
            `return_distances` must be True.
        indices: An output array to store the calculated feature transform, instead of returning it. `return_indicies` must be True.
        block_params: This parameter is specific to cuCIM and does not exist in SciPy. For details, look into `cuCIM`_.
        float64_distances: This parameter is specific to cuCIM and does not exist in SciPy.
            If True, use double precision in the distance computation (to match SciPy behavior).
            Otherwise, single precision will be used for efficiency.

    Returns:
        distances: The calculated distance transform. Returned only when `return_distances` is True and `distances` is not supplied.
            It will have the same shape as image. For cuCIM: Will have dtype torch.float64 if float64_distances is True,
            otherwise it will have dtype torch.float32. For SciPy: Will have dtype np.float64.
        indices: The calculated feature transform. It has an image-shaped array for each dimension of the image.
            Returned only when `return_indices` is True and `indices` is not supplied. dtype np.float64.

    """
    distance_transform_edt, has_cucim = optional_import(
        "cucim.core.operations.morphology", name="distance_transform_edt"
    )
    use_cp = has_cp and has_cucim and isinstance(img, torch.Tensor) and img.device.type == "cuda"