Versatile distance transforms (Euclidean and Geodesic)

[tvercaut]

Is your feature request related to a problem? Please describe.
Several segmentation pipelines, and in particular interactive segementation ones, rely on some form of distance transform (e.g. Euclidean, Chamfer or Geodesic Distance Transform).

At the moment there is no off-the-shelf pytorch function to compute these.

Describe the solution you'd like
A MONAI-based implementation of distance transforms (ideally differentiable) would be fantastic. It woul accelerate uptake of such methods.

Describe alternatives you've considered

• Relying on scipy functions such as distance_transform_edt and distance_transform_cdt
• Rely on external code such as GeodisTK

Additional context
Example use cases of distance transforms in segmentaiton:

• [WIP] Support Deepgrow 2D/3D Training #1329
• https://github.com/HiLab-git/UGIR
• https://arxiv.org/abs/1710.04043
• https://arxiv.org/abs/1707.00652

[wyli]

also used in monai metrics computation

MONAI/monai/metrics/utils.py

Line 23 in f01994d

distance_transform_cdt, _ = optional_import("scipy.ndimage.morphology", name="distance_transform_cdt")

[ddrobny]

There is a recent paper using a convolution-based and differential approximation for a distance transform:
https://link.springer.com/chapter/10.1007%2F978-3-030-71278-5_31
There doesn't seem to be related code published yet but it might be worth looking into.

[tvercaut]

Thanks @ddrobny.

On a related note, tensorflow has a GPU implementation of the Felzenszwalb & Huttenlocher Distance transform:

• https://www.tensorflow.org/addons/api_docs/python/tfa/image/euclidean_dist_transform
• https://github.com/tensorflow/addons/blob/master/tensorflow_addons/custom_ops/image/cc/kernels/euclidean_distance_transform_op_gpu.cu.cc
  Note sure if it's differentiable but it would be neat to at least feature match it in MONAI

For Euclidean distance, we can compare the performance (and read the interesting comments) with an established (but GPL so don't look at the code) implementation:

• https://github.com/seung-lab/euclidean-distance-transform-3d

[tvercaut]

As mentioned in #4603, @masadcv just released FastGeodis, a package to compute Geodesic and Euclidean distance maps in PyTorch.

[grlee77]

We plan to have an equivalent of SciPy's distance_transform_edt (Euclidean distance transform) for CuPy arrays in the upcoming cuCIM 22.08 release. The GPU implementation is currently in 2D and 3D only as compared to the general nD one in SciPy.

see: rapidsai/cucim#318

[matt3o]

@grlee77 Thanks a lot! I just tried it on the DeepEdit Code and it appears to be 100% compatible to what scipy does, really awesome.
For everyone else trying this out: The FastGeodis Code yields different results than the scipy one, also see masadcv/FastGeodis#11 about that. I did not verify what impact it has, so it might still work as usual.

[matt3o]

Also for anyone wondering what the impact of this would be, I did some profiling on a full 3D volume. Here is the according snippet: https://gist.github.com/matt3o/0de8a24a10135ed24de6f77ab89e15b1

And the results were:

INFO:root:torch.Size([400, 400, 284])
INFO:root:distance_transform_cdt_scipy: 20 runs took 6.403534 seconds, which means 0.320177 seconds per run
INFO:root:distance_transform_edt_scipy: 20 runs took 99.965921 seconds, which means 4.998296 seconds per run
INFO:root:distance_transform_edt_cupy: 20 runs took 0.162181 seconds, which means 0.008109 seconds per run.

Edit: Added the cdt transform too, since that is the one used in DeepEdit. CDT runs a lot faster on CPU, still it's a huge difference for big volumes.