Source code for vlkit.geometry.distance_transform

import numpy as np
from scipy.ndimage.morphology import distance_transform_edt


def distance_transform(x):
    """distance transform

    Args:
        x (np.ndarray): a [h w] binary map

    Returns:
        (tuple): tuple containing:

            dist (np.ndarray): the distance matrix [h w] \n
            yxs (np.ndarray): the coordinates of nearest points [2 h w] \n
            field (np.ndarray): the directional vector field [2 h w]

    .. image:: _static/distance_transform.svg

    See `an example <https://github.com/vlkit/vlkit/blob/master/examples/distance_transform.ipynb>`_ here 
    """
    assert isinstance(x, np.ndarray)
    assert x.ndim == 2
    h, w = x.shape

    x = x.astype(bool)
    dist, yxs = distance_transform_edt(np.logical_not(x), return_distances=True, return_indices=True)

    ys = np.arange(h).reshape(-1, 1)
    xs = np.arange(w).reshape(1, -1)
    field = yxs - np.stack((np.tile(ys, (1, w)), np.tile(xs, (h, 1))))

    return dist, yxs, field