Skip to content

fiery-distmap

Euclidean distance transform in PyTorch.

fiery-distmap is a fiery match; it imports as fiery.distmap.

This is an implementation of the algorithm from the paper

"Distance Transforms of Sampled Functions"
Pedro F. Felzenszwalb & Daniel P. Huttenlocher
Theory of Computing (2012)

Although it is written in PyTorch, our implementation loops across voxels and is therefore quite slow. Moreover, it takes masks as input and therefore does not support backpropagation.

Installation

Dependency

  • torch >= 1.3

Pip

pip install fiery-distmap

Usage

from fiery.distmap import euclidean_distance_transform

Example

See our demo notebook

API

euclidean_distance_transform(x, ndim=None, vx=1, squared=False)
"""Compute the Euclidean distance transform of a binary image

Parameters
----------
x : (..., *spatial) tensor
    Input tensor. Zeros will stay zero, and the distance will
    be propagated into nonzero voxels.
ndim : int, default=`x.dim()`
    Number of spatial dimensions
vx : [sequence of] float, default=1
    Voxel size
squared : bool, default=False
    Return the squared distance map, skipping the final square root.

Returns
-------
d : (..., *spatial) tensor
    Distance map
"""

euclidean_signed_transform(x, ndim=None, vx=1, squared=False)
"""Compute the signed Euclidean distance transform of a binary image

Parameters
----------
x : (..., *spatial) tensor
    Input tensor.
    A negative distance will propagate into zero voxels and
    a positive distance will propagate into nonzero voxels.
ndim : int, default=`x.dim()`
    Number of spatial dimensions
vx : [sequence of] float, default=1
    Voxel size
squared : bool, default=False
    Return the squared distance map, skipping the final square root.

Returns
-------
d : (..., *spatial) tensor
    Signed distance map
"""

l1_distance_transform(x, ndim=None, vx=1)
"""Compute the L1 distance transform of a binary image

Parameters
----------
x : (..., *spatial) tensor
    Input tensor. Zeros will stay zero, and the distance will
    be propagated into nonzero voxels.
ndim : int, default=`x.dim()`
    Number of spatial dimensions
vx : [sequence of] float, default=1
    Voxel size

Returns
-------
d : (..., *spatial) tensor
    Distance map
"""

l1_signed_transform(x, ndim=None, vx=1)
"""Compute the signed L1 distance transform of a binary image

Parameters
----------
x : (..., *spatial) tensor
    Input tensor.
    A negative distance will propagate into zero voxels and
    a positive distance will propagate into nonzero voxels.
ndim : int, default=`x.dim()`
    Number of spatial dimensions
vx : [sequence of] float, default=1
    Voxel size

Returns
-------
d : (..., *spatial) tensor
    Signed distance map
"""

  • edt : a very fast CPU implementation of the same algorithm, written in C.

  • scipy.ndimage.distance_transform_edt : reference implementation, written in C, based on the paper

    "A linear time algorithm for computing exact euclidean distance transforms of binary images in arbitrary dimensions"
    C. R. Maurer, Jr., R. Qi, V. Raghavan
    IEEE Trans. PAMI 25, 265-270, (2003)