How to efficiently calculate the distance to the nearest 1 in a mask in numpy?

Question

In numpy, I have a 2d array of 1s and 0s. I need to calculate a new array (same sizes), where each element contains the distance to the nearest 1 from the corresponding point in the mask array.

eg.

a=np.array(
[[1,1,0],
[1,0,0],
[1,0,0]])

I need b to look like this:

array([[0,0,1],
       [0,1,1.41],
       [0,1,2]])

PS. I will do this on very large arrays, so the more efficient the better! Thank!

Answer 1

You are looking for the equivalent of MATLAB bwdist; check out this SO question for more details. The short answer is to use scipy.ndimage.morphology.distance_transform_edt.