SIMD implementation of 2d euclidean distance. For each non-zero input pixel, returns the exact squared distance to the closest zero input pixel.

euclidean-distance-transform-2d.py

import numpy

def euclidean_distance_transform(F):
	ret = numpy.empty_like(F, dtype = int)

	def get_parabola_matrix(N):
		X = numpy.arange(N, dtype = int)
		G = numpy.empty((N, N), dtype = int)
		G[:] = X
		G -= X[:, None]
		G *= G
		return G

	# Horizontal pass
	G = get_parabola_matrix(F.shape[1])
	for i in range(F.shape[0]):
		numpy.amin(G + (F.shape[1] ** 2) * (F[i][:, None] != 0), axis = 0, out = ret[i, :])		

	# Vertical pass
	G = get_parabola_matrix(F.shape[0])	
	for i in range(F.shape[1]):
		numpy.amin(G + ret[:,i][:, None], axis = 0, out = ret[:, i])		

	# Job done
	return ret