Project-MONAI
Versatile distance transforms (Euclidean and Geodesic)
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:
- https://github.com/Project-MONAI/MONAI/pull/1329
- https://github.com/HiLab-git/UGIR
- https://arxiv.org/abs/1710.04043
- https://arxiv.org/abs/1707.00652
also used in monai metrics computation https://github.com/Project-MONAI/MONAI/blob/f01994d47403f920b5a51cfeae75418383d10178/monai/metrics/utils.py#L23
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.
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
As mentioned in #4603, @masadcv just released FastGeodis, a package to compute Geodesic and Euclidean distance maps in PyTorch.
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: https://github.com/rapidsai/cucim/pull/318
@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 https://github.com/masadcv/FastGeodis/issues/11 about that. I did not verify what impact it has, so it might still work as usual.
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.
As of https://github.com/Project-MONAI/MONAI/pull/6981, the Euclidean distance transform will be part of MONAI from version 1.3 ongoing. It has GPU support via cuCIM and runs a lot faster. Falls back to SciPy if necessary, look into the docs for details.
I will also soon publish an overhauled version of DeepEdit with much faster transforms based on torch and this GPU-based distance transform, which will hopefully be integrated into MONAI as well