Distance transform on image using NumPy

Question

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 plt.figure() plt.imshow(dist_transform)