I would like to find the find the distance transform of a binary image in the fastest way possible without using the scipy function distance_transform_edt(). The image is 256 by 256. The reason I don't want to use scipy is because using it is difficult in Tensorflow. Every time I want to use this package I need to start a new session and this takes a lot of time. So I would like to make a custom function that only utilizes NumPy.

My approach is as follows: Find the coordinated for all the ones and all the zeros in the image. Find the euclidian distance between each of the zero pixels (a) and the one pixels (b) and then the value at each (a) position is the minimum distance to a (b) pixel. I do this for each 0 pixel. The resultant image has the same dimensions as the original binary map. My attempt at doing this is shown below.

I tried to do this as fast as possible using no loops and only vectorization. But my function still can't work as fast as the scipy package can. When I timed the code it looks like the assignment to the variable "a" is taking the longest time. But I do not know if there is a way to speed this up.

If anyone has any other suggestions for different algorithms to solve this problem of distance transforms or can direct me to other implementations in python, it would be very appreciated.

def get_dst_transform_img(og): #og is a numpy array of original image
   ones_loc = np.where(og == 1)
   ones = np.asarray(ones_loc).T # coords of all ones in og
   zeros_loc = np.where(og == 0)
   zeros = np.asarray(zeros_loc).T # coords of all zeros in og

   a = -2 * np.dot(zeros, ones.T) 
   b = np.sum(np.square(ones), axis=1) 
   c = np.sum(np.square(zeros), axis=1)[:,np.newaxis]
   dists = a + b + c
   dists = np.sqrt(dists.min(axis=1)) # min dist of each zero pixel to one pixel
   x = og.shape[0]
   y = og.shape[1]
   dist_transform = np.zeros((x,y))
   dist_transform[zeros[:,0], zeros[:,1]] = dists 

  • 2
    \$\begingroup\$ This post is double-posted on SO, and answered there. \$\endgroup\$ – Cris Luengo Dec 12 '18 at 5:55
  • \$\begingroup\$ You could just try to copy the code from the scipy function. There are two imports used as far as I can tell, one is a normal Python file in the same folder, the other one is written in C and needs to be compiled, though. \$\endgroup\$ – Graipher Dec 12 '18 at 10:45

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.