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 (performance):
When we implement non-spatial dimensions and we care about performance here, I think we should use the option to pass in a preallocated torch tensor/numpyarray with the final dimension using the [distances and indices parameter](https://docs.rapids.ai/api/cucim/nightly/api/#cucim.core.operations.morphology.distance_transform_edt) (applies to [scipy](https://docs.scipy.org/doc/scipy/reference/generated/scipy.ndimage.distance_transform_edt.html#scipy-ndimage-distance-transform-edt) as well). This way we reduce the number of allocations and copies compared to stacking. We then also need to make sure to match the proper datatypes. Keep in mind that we should support the float64_distances parameter for cucim here as well, since the float32 should be faster on GPU and the precision is probably not necessary in most cases. I am not 100% sure, but I think in case of cucim we should allocate a torch tensor and then use convert_to_cupy() to create a temporary view on this tensor cupy to pass to the transform. This way the result should get allocated using the pytorch memory pool

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"