Methodology
Generalised Geodesic Distance Transform
A distance transform maps each image pixel/voxel into its smallest distance to regions of interest [FCTB08]. It is a fundamental operator with relevant applications in computer vision, graphics, shape analysis, pattern recognition and computational geometry [FCTB08].
Generalised Geodesic distance transform combines spatial as well as gradient based distance transform and is defined as [CSB08]:
The ability to incorporate image gradients with spatial distances has enabled application of Geodesic distance transforms in a number of areas, including image editing and filtering [CSB08]. Methods from recent years have shown effectiveness in applying Geodesic distance transform to interactively annotate 3D medical imaging data [CSB08, WZL+18], where it enables providing segmentation labels, i.e. voxel-wise labels, for different objects of interests.
Statement of Need
Despite existing open-source implementation of distance transforms [AAB+15, SL18, Wan20], open-source implementations of efficient Geodesic distance transform algorithms [CSS09, WDB+08] on CPU and GPU do not exist. However, efficient CPU [SL18] and GPU [AAB+15] implementations exist for Euclidean distance transform. To the best of our knowledge, FastGeodis is the first open-source implementation of efficient Geodesic distance transform [CSS09], achieving up to 20x speed-up on CPU and up to 74x speed-up on GPU as compared to existing open-source libraries [Wan20]. It also provides efficient implementation of Euclidean distance transform. In addition, it is the first open-source implementation of generalised Geodesic distance transform and Geodesic Symmetric Filtering (GSF) proposed in [CSB08].
The ability to efficiently compute Geodesic and Euclidean distance transforms can significantly enhance distance transform applications especially for training deep learning models that utilise distance transforms [WZL+18]. It will improve prototyping, experimentation, and deployment of such methods, where efficient computation of distance transforms has been a limiting factor. In 3D medical imaging problems, efficient computation of distance transforms will lead to significant speed-ups, enabling online learning applications for better processing/labelling/inference from volumetric datasets [AFV22]. In addition, FastGeodis provides efficient implementation for both CPUs and GPUs hardware and hence will enable efficient use of a wide range of hardware devices.
Implementation
FastGeodis implements an efficient distance transform algorithm from [CSS09], which provides parallelisable raster scans to compute distance transform. The implementation consists of data propagation passes parallelised using threads for elements across a line (2D) or plane (3D). Figure below shows these data propagation passes, where each pass consists of computing distance values for next row (2D) or plane (3D) by utilising parallel threads and data from previous row/plane, hence resulting in propagating distance values along the direction of pass. For 2D data, four distance propagation passes are required, top-bottom, bottom-top, left-right and right-left, whereas for 3D data six passes are required, front-back, back-front, top-bottom, bottom-top, left-right and right-left. The algorithm can be applied to efficiently compute both Geodesic and Euclidean distance transforms. In addition to this, FastGeodis also provides non-parallelisable raster scan based distance transform method from [Toi96], which is implemented using single-thread CPU and used for comparison.
2D images: 1 of 4 passes |
3D volumes: 1 of 6 passes |
|
|
FastGeodis package is implemented using PyTorch [PGM+19] utilising OpenMP for CPU and CUDA for GPU parallelisation of the algorithm. It is accessible as a python package, that can be installed across different operating systems and devices. A comprehensive documentation and a range of examples are provided for understanding the usage of the package on 2D and 3D data using CPU or GPU. The provided examples include 2D/3D examples for Geodesic, Euclidean, Signed Geodesic distance transform as well as computing Geodesic symmetric filtering (GSF) that is essential first step in implementing interactive segmentation method from [CSB08].
In the following section, more details of each implemented algorithm is presented.
Criminisi et al’s Parallelisable Generalised Geodesic Distance Transform
We implement both 2D and 3D parallelisable generalised Geodesic distance transform algorithms from [CSS09] on both CPU (OpenMP) and GPU (CUDA). The 2D algorithm works by computing distance propagation in one row at a time. This is a hard constraint because compute for each row is dependent on computed distances in previous rows from previous invocation. Consider an example of an image with 4 x 6 dimension. Then a successful full iteration of our method would involve going through the rows one by one as follows:
As can be seen, this involves top-down and left-right passes. How we implement left-right is by reusing top-down code and transposing the data instead. Please note that for each step, only the pixels highlighted in green color can be computed (as they have the data available from previous row). It is this row that we split into multiple threads using an underlying hardware (CPU or GPU). This parallelisation enables speed-up as compared to non-parallelisable CPU implementations, e.g. in [Wan20] which implements raster scan algorithm from [Toi96].
Going beyond 2D images, we also implement the parallelisable Geodesic distance transform for 3D data. We provide both CPU (OpenMP) and GPU (CUDA) optimised implementations. Our 3D implementation operates on the same principle (we can process one plane at a time). However, in 3D case, since we have more data, we can utilise more compute on GPU and process a plane in parallel, however we still have the data dependency constraints that prevent us from processing all planes together.
[WDB+08] presents a further optimised approach for computing Geodesic distance transforms on GPUs, however this method is protected by multiple patents and hence is not suitable for open-source implementation in FastGeodis package.
Toivanen et al’s Non-parallelisable Generalised Geodesic Distance Transform
In addition to parallelisable algorithm from [CSS09], we also implement non-parallelisable Geodesic distance transform method from [Toi96] using CPU. This method is used for comparison of accuracy as well as execution of the parallelised Geodesic distance transform algorithm presented above.
In this method, a 2D pass operates as a raster scan that is sequentially applied in two passes: forward and backward. Figure below, shows an overview of this method, as described in [CSB08]:
For both forward and backward pass, an L shaped kernel is used in a single raster scan pass to propagate distance in forward and backward direction. In particular, this L shaped kernel along with the sequential pass limits this method to non-parallelisable CPU implementation, which we include in FastGeodis for comparison purposes.
Sethian’s Fast Marching-based Non-parallelisable Generalised Geodesic Distance Transform
In addition to the above methods, we also implement fast marching based non-parallelisable Geodesic distance transform method from [Set99] using CPU. This method is used for comparison of accuracy as well as execution of the parallelised Geodesic distance transform algorithm presented above. It serves as golden reference as this method provides accuracy distance transform calculation. While being accurate, it is computationally expensive taking orders of magnitude more time and compute.
Performance Improvements
FastGeodis (CPU/GPU) is compared with existing GeodisTK (https://github.com/taigw/GeodisTK) in terms of execution speed as well as accuracy. All our experiments were evaluated on Nvidia GeForce Titan X (12 GB) with 6-Core Intel Xeon E5-1650 CPU. We present our results below:
Execution
The following figures summarise execution speed comparison of FastGeodis with [Wan20].
Execution Speed for 2D images |
Execution Speed for 3D volumes |
|
|
The above results are further summarised in tables below, along with calculated speed-ups for FastGeodis functions vs [Wan20]:
Device |
CPU |
CPU |
GPU |
CPU |
GPU |
|---|---|---|---|---|---|
Spatial Size |
GeodisTK |
FastGeodis |
FastGeodis |
Speed-up |
Speed-up |
64 x 64 |
0.00401 s |
0.00178 s |
0.00245 s |
2.3 x |
1.6 x |
128 x 128 |
0.01555 s |
0.00279 s |
0.00487 s |
5.6 x |
3.2 x |
256 x 256 |
0.06412 s |
0.00705 s |
0.00961 s |
9.1 x |
6.7 x |
512 x 512 |
0.25041 s |
0.01873 s |
0.01735 s |
13.4 x |
14.4 x |
1024 x 1024 |
1.00589 s |
0.04581 s |
0.03461 s |
22 x |
29.1 x |
2048 x 2048 |
4.01677 s |
0.18565 s |
0.07161 s |
21.6 x |
56.1 x |
Device |
CPU |
CPU |
GPU |
CPU |
GPU |
|---|---|---|---|---|---|
Spatial Size |
GeodisTK |
FastGeodis |
FastGeodis |
Speed-up |
Speed-up |
64 x 64 x64 |
0.13098 s |
0.03933 s |
0.00724 s |
3.3 x |
18.1 x |
128 x 128 x 128 |
1.09478 s |
0.30489 s |
0.01783 s |
3.6 x |
61.4 x |
256 x 256 x 256 |
8.60261 s |
2.53498 s |
0.12443 s |
3.4 x |
69.1 x |
512 x 512 x 512 |
69.08086 s |
22.24171 s |
0.92998 s |
3.1 x |
74.3 x |
It can be observed that for 2D images, FastGeodis leads to a speed-up of upto 20x on CPU and upto 55x on GPU. For 3D images, FastGeodis leads to a speed-up of upto 3x on CPU and upto 74x on GPU.
Accuracy
For accuracy, we use Fast Marching-based implementation as golden reference. These results are visualised below for visual comparison as well as quantitative comparison using joint histograms.
2D Image Data
3D Image Data
Martín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Greg S. Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Ian Goodfellow, Andrew Harp, Geoffrey Irving, Michael Isard, Yangqing Jia, Rafal Jozefowicz, Lukasz Kaiser, Manjunath Kudlur, Josh Levenberg, Dandelion Mané, Rajat Monga, Sherry Moore, Derek Murray, Chris Olah, Mike Schuster, Jonathon Shlens, Benoit Steiner, Ilya Sutskever, Kunal Talwar, Paul Tucker, Vincent Vanhoucke, Vijay Vasudevan, Fernanda Viégas, Oriol Vinyals, Pete Warden, Martin Wattenberg, Martin Wicke, Yuan Yu, and Xiaoqiang Zheng. TensorFlow: large-scale machine learning on heterogeneous systems. 2015. Software available from tensorflow.org. URL: https://www.tensorflow.org/.
Muhammad Asad, Lucas Fidon, and Tom Vercauteren. Econet: efficient convolutional online likelihood network for scribble-based interactive segmentation. arXiv preprint arXiv:2201.04584, 2022.
Antonio Criminisi, Toby Sharp, and Andrew Blake. Geos: geodesic image segmentation. In European Conference on Computer Vision, 99–112. Springer, 2008. doi:10.1007/978-3-540-88682-2_9.
Antonio Criminisi, Toby Sharp, and Khan Siddiqui. Interactive geodesic segmentation of n-dimensional medical images on the graphics processor. In Radiological Society of North America (RSNA). December 2009. URL: https://www.microsoft.com/en-us/research/publication/interactive-geodesic-segmentation-of-n-dimensional-medical-images-on-the-graphics-processor/.
Ricardo Fabbri, Luciano Da F Costa, Julio C Torelli, and Odemir M Bruno. 2d euclidean distance transform algorithms: a comparative survey. ACM Computing Surveys (CSUR), 40(1):1–44, 2008.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: an imperative style, high-performance deep learning library. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d\textquotesingle Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems 32, pages 8024–8035. Curran Associates, Inc., 2019. URL: http://papers.neurips.cc/paper/9015-pytorch-an-imperative-style-high-performance-deep-learning-library.pdf.
James A Sethian. Fast marching methods. SIAM review, 41(2):199–235, 1999.
Seung-Lab. Multi-label anisotropic 3d euclidean distance transform (mlaedt-3d). 2018. URL: https://github.com/seung-lab/euclidean-distance-transform-3d.
Pekka J Toivanen. New geodosic distance transforms for gray-scale images. Pattern Recognition Letters, 17(5):437–450, 1996. doi:10.1016/0167-8655(96)00010-4.
Guotai Wang. Geodistk: geodesic distance transform toolkit for 2d and 3d images. 2020. URL: https://github.com/taigw/GeodisTK.
Guotai Wang, Maria A Zuluaga, Wenqi Li, Rosalind Pratt, Premal A Patel, Michael Aertsen, Tom Doel, Anna L David, Jan Deprest, Sébastien Ourselin, and others. Deepigeos: a deep interactive geodesic framework for medical image segmentation. IEEE transactions on pattern analysis and machine intelligence, 41(7):1559–1572, 2018.
Ofir Weber, Yohai S Devir, Alexander M Bronstein, Michael M Bronstein, and Ron Kimmel. Parallel algorithms for approximation of distance maps on parametric surfaces. ACM Transactions on Graphics (TOG), 27(4):1–16, 2008. doi:10.1145/1409625.1409626.