find nearest neighbor to each pixel in a map

Question

I am trying to find a simple method in python in which for each pixel in a 2dim mask, I can get the indices of the nearest non-zero neighbor. In Matlab there is a bwdist which returns exactly that. for example: if my input is:

array [[0 0 0 0 0 0 0]
       [0 1 0 0 0 0 0]
       [0 0 0 0 0 1 0]
       [0 0 0 0 0 0 0]]

my output should be:

array [[(1,1) (1,1) (1,1) (1,1) (2,5) (2,5) (2,5)]
       [(1,1) (1,1) (1,1) (1,1) (2,5) (2,5) (2,5)]
       [(1,1) (1,1) (1,1) (2,5) (2,5) (2,5) (2,5)]
       [(1,1) (1,1) (1,1) (2,5) (2,5) (2,5) (2,5)]]

The function can also return the absolute index (for 1dim array) like bwdist in Matlab.

Thanks!

EDIT: so far, I have tried some potential solutions relating to scipy such as distance_transform_edt, but it only finds the distance to the closest pixel and not the pixel itself. I also use OpenCV and VLfeat in other places in my code if its relevant.

Answer 1

This is actually a one-liner when using scipy.

If your input matrix is mat, the coordinates of the nearest nonzero value are given by:

import scipy.ndimage

nearest_neighbor = scipy.ndimage.morphology.distance_transform_edt(
    mat==0, return_distances=False, return_indices=True)

For the matrix given in the question, this results in the following index matrix, which is the correct answer:

[[[1 1 1 1 2 2 2]
  [1 1 1 1 2 2 2]
  [1 1 1 2 2 2 2]
  [1 1 1 2 2 2 2]]

 [[1 1 1 1 5 5 5]
  [1 1 1 1 5 5 5]
  [1 1 1 5 5 5 5]
  [1 1 1 5 5 5 5]]]

The index matrix is read as following: The point at 0,0's closest neighbor is at 1,1. The point at 0,4's closest neighbor is at 2,5.

Answer 2

From the docs:

OpenCV has distanceTransform() and distanceTransformWithLabels() functions which work sort of similarly, but there are some differences from this Matlab function. From the Matlab docs for bwdist:

D = bwdist(BW) computes the Euclidean distance transform of the binary image BW. For each pixel in BW, the distance transform assigns a number that is the distance between that pixel and the nearest nonzero pixel of BW.

Compare this to the OpenCV docs for distanceTransformWithLabels():

Calculates the distance to the closest zero pixel for each pixel of the source image.

So Matlab gives the distance to the nearest non-zero pixel, while OpenCV gives the distance to the nearest zero pixel. So you'll need to invert the image for OpenCV. Additionally, the optional output for Matlab with the labels gives the linear index corresponding to that closest pixel:

[D,idx] = bwdist(BW) also computes the closest-pixel map in the form of an index array, idx. Each element of idx contains the linear index of the nearest nonzero pixel of BW. The closest-pixel map is also called the feature map, feature transform, or nearest-neighbor transform.

With OpenCV, the label that gets output is not a coordinate of the image, nor an index. Instead it's just a numerical label, similar to a connected component label, that is not related to the pixel location/index at all.

This variant of the function does not only compute the minimum distance for each pixel (x,y) but also identifies the nearest connected component consisting of zero pixels (labelType==DIST_LABEL_CCOMP) or the nearest zero pixel (labelType==DIST_LABEL_PIXEL).

This means that you'll have to use this labeled image to mask your input and find the pixel that corresponds to that label (as far as I know, this is the best way to do it, at least).

Solution:

So just to wrap our heads around how to get where we want, let's take a look at where this function gets us (with an inverted image as input as stated before):

In [138]: img
Out[138]:
array([[  0,   0,   0,   0,   0,   0,   0],
       [  0, 255,   0,   0,   0,   0,   0],
       [  0,   0,   0,   0,   0, 255,   0],
       [  0,   0,   0,   0,   0,   0,   0]], dtype=uint8)

In [139]: dist, labels = cv2.distanceTransformWithLabels(~a, distanceType=cv2.DIST_L2, maskSize=3)

In [140]: print(dist)
[[1.3999939 1.        1.3999939 2.1968994 2.1968994 2.        2.1968994]
 [1.        0.        1.        2.        1.3999939 1.        1.3999939]
 [1.3999939 1.        1.3999939 2.        1.        0.        1.       ]
 [2.1968994 2.        2.1968994 2.1968994 1.3999939 1.        1.3999939]]

In [141]: print(labels)
[[1 1 1 1 2 2 2]
 [1 1 1 1 2 2 2]
 [1 1 1 2 2 2 2]
 [1 1 1 2 2 2 2]]

So ok, if we just loop through the unique values in labels, create a mask for each of them, mask the original image...and then find the white pixel inside that labeled area, we'll have the indices:

In [146]: for l in np.unique(labels):
     ...:     mask = label == l
     ...:     i = np.where(img * mask)
     ...:     print(i)
     ...:
(array([1]), array([1]))
(array([2]), array([5]))

This isn't the exact output you asked for, but it is a list of the indices, and you have the labels. So now we just need to map these. What I'll do is create an empty two-channel matrix to hold the index values, and then fill it based on the mask from the labels:

In [177]: index_img = np.zeros((*img.shape, 2), dtype=np.intp)

In [178]: for l in np.unique(labels):
     ...:     mask = label == l
     ...:     index_img[mask] = np.dstack(np.where(img * mask))

And this is a two-channel array with the info you want. The structure is a little different (not using tuples for each entry), but it is usually the structure you want for other OpenCV functions (two-channel array):

In [204]: index_img[:, :, 0]
Out[204]:
array([[1, 1, 1, 1, 2, 2, 2],
       [1, 1, 1, 1, 2, 2, 2],
       [1, 1, 1, 2, 2, 2, 2],
       [1, 1, 1, 2, 2, 2, 2]])

In [205]: index_img[:, :, 1]
Out[205]:
array([[1, 1, 1, 1, 5, 5, 5],
       [1, 1, 1, 1, 5, 5, 5],
       [1, 1, 1, 5, 5, 5, 5],
       [1, 1, 1, 5, 5, 5, 5]])

Putting it all together

Here's a function which does this, and has an option for spitting out this two channel output or just the linear output as Matlab does:

def bwdist(img, metric=cv2.DIST_L2, dist_mask=cv2.DIST_MASK_5, label_type=cv2.DIST_LABEL_CCOMP, ravel=True):
    """Mimics Matlab's bwdist function.

    Available metrics:
        https://docs.opencv.org/3.4/d7/d1b/group__imgproc__misc.html#gaa2bfbebbc5c320526897996aafa1d8eb
    Available distance masks:
        https://docs.opencv.org/3.4/d7/d1b/group__imgproc__misc.html#gaaa68392323ccf7fad87570e41259b497
    Available label types:
        https://docs.opencv.org/3.4/d7/d1b/group__imgproc__misc.html#ga3fe343d63844c40318ee627bd1c1c42f
    """
    flip = cv2.threshold(img, 0, 255, cv2.THRESH_BINARY_INV)[1]
    dist, labeled = cv2.distanceTransformWithLabels(flip, metric, dist_mask)

    # return linear indices if ravel == True (default)
    if ravel:  
        idx = np.zeros(img.shape, dtype=np.intp)  # np.intp type is for indices
        for l in np.unique(labeled):
            mask = labeled == l
            idx[mask] = np.flatnonzero(img * mask)
        return dist, idx

    # return two-channel indices if ravel == False
    idx = np.zeros((*img.shape, 2), dtype=np.intp)  
    for l in np.unique(labeled):
        mask = labeled == l
        idx[mask] = np.dstack(np.where(img * mask))
    return dist, idx

And with the example Matlab gives in the documentation:

In [241]: bw = np.zeros((5, 5), dtype=np.uint8)
     ...: bw[1, 1] = 1
     ...: bw[3, 3] = 1
     ...: print(bw)
     ...:
[[0 0 0 0 0]
 [0 1 0 0 0]
 [0 0 0 0 0]
 [0 0 0 1 0]
 [0 0 0 0 0]]

In [244]: d, idx = bwdist(bw)

In [245]: print(d)
[[1.3999939 1.        1.3999939 2.1968994 3.1968994]
 [1.        0.        1.        2.        2.1968994]
 [1.3999939 1.        1.3999939 1.        1.3999939]
 [2.1968994 2.        1.        0.        1.       ]
 [3.1968994 2.1968994 1.3999939 1.        1.3999939]]

In [246]: print(idx)
[[ 6  6  6  6 18]
 [ 6  6  6  6 18]
 [ 6  6  6 18 18]
 [ 6  6 18 18 18]
 [ 6 18 18 18 18]]