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):
Thanks for the PR, I would suggest all the input should be channel-first since most of the monai transform follows instead of onlt treating 4D with channel-first which may cause confusion with CHW and HWD just like @john-zielke-snkeos said.
https://github.com/Project-MONAI/MONAI/blob/e5f933781b1a5f4eb7d278e06389d531bff38dca/monai/transforms/utils.py#L969
Thanks!

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"