[FEATURE] GPU implementation of accurate Euclidean Distance Transforms

[tvercaut]

Is your feature request related to a problem? Please describe.
The package currently provide GPU-based approximations of Geodesic and Euclidean distance transforms.

Following #42 fast marching now allows for more accurate computations of Euclidean distance transforms but this is limited to CPU operations and remains an approximation.

Describe the solution you'd like
Having a GPU implementation of an exact Euclidean Distance Transform would strengthen the package.

Describe alternatives you've considered
Rely on third party libs for this. e.g. CUCIM, Tensorflow or PBA+
https://docs.rapids.ai/api/cucim/stable/api/#cucim.core.operations.morphology.distance_transform_edt
https://www.tensorflow.org/addons/api_docs/python/tfa/image/euclidean_dist_transform
https://github.com/orzzzjq/Parallel-Banding-Algorithm-plus

Having this integrated in FastGeodis withouth the need for external dependencies would nonetheless be nice and make the pytorch link tighter.

Additional context
This feature request has been voice earlier, e.g.:
#40 (comment)
Project-MONAI/MONAI#1332 (comment)

The CUCIM code, based on PBA+ is released under MIT, seems fairly standalone, and could thus be copy-pasted in this repo for convenience:
https://github.com/rapidsai/cucim/tree/main/python/cucim/src/cucim/core/operations/morphology

[tvercaut]

In case a CPU fallback is interesting, I have identified and modernised a suitable standalone package here:
https://github.com/cai4cai/distance_transform

It is still missing support for anisotropic voxels though:
cai4cai/distance_transform#1

[masadcv]

HI @tvercaut ,
Many thanks for highlighting potential implementations that we could use.
I will have a look into the GPU implementation.
For the CPU, how is this method different from Eikonal solver that we have for Fast Marching? Does it make sense to replace Fast March or keep it and add this as additional method for CPU based Euclidean distance?

[tvercaut]

Fast marching is more generic but only provides an approximation. Felzenszwalb & Huttenlocher and PBA are both restricted to Euclidean distances but should be exact. I would thus keep both options in.

[masadcv]

okay, many thanks!, let me look into this... I will update with a PR in the coming days...

For GPU, I need more time to understand and then implement the approach linked above.

[tvercaut]

I just took a stab at supporting anisotrpic voxels in the CPU algorithm for exact Euclidean distance transforms. You can check it in this branch:
https://github.com/cai4cai/distance_transform/tree/anisotropy
This would certainly benefit from proper testing but should be suitable to give a go (or serve as a basis for re-implementation with libtorch).

[tvercaut]

On a related note, we could also consider implementing the Jump Flooding Algorithm (JFA):
https://www.comp.nus.edu.sg/~tants/jfa.html

It provides a competitive lternative to PBA according to this reference and it seems easy to implement:
https://urn.kb.se/resolve?urn=urn:nbn:se:liu:diva-163923

Code for JFA can be found in several locations:

• https://urn.kb.se/resolve?urn=urn:nbn:se:liu:diva-163923 (Appendix C)
• https://www.shadertoy.com/view/dsXSD8
• http://shadertoy.com/view/cdlXRl
• https://blog.demofox.org/2016/02/29/fast-voronoi-diagrams-and-distance-dield-textures-on-the-gpu-with-the-jump-flooding-algorithm/
• https://www.shadertoy.com/view/WlGyR3

[tvercaut]

A quick note to mention that the PBA+ algorithm can now also be used from cupy (it has been ported from CUCIM to cupy):
https://docs.cupy.dev/en/stable/reference/generated/cupyx.scipy.ndimage.distance_transform_edt.html#cupyx.scipy.ndimage.distance_transform_edt