in tensorflow_addons/custom_ops/image/cc/kernels/euclidean_distance_transform_op.h [70:126]
EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE void EuclideanDistanceTransformSample(
const uint8* input, T* output, int k, int c, int height, int width,
int channels) {
for (int i = 0; i < height; i++) {
for (int j = 0; j < width; j++) {
int index = GET_INDEX(i, j, k, c);
if (input[index] == 0) {
output[index] = static_cast<T>(0);
} else {
output[index] = Eigen::NumTraits<T>::highest();
}
}
}
int max = Eigen::numext::maxi(height, width);
T* f = new T[max];
T* d = new T[max];
// locations of parabolas in lower envelope
int* vw = new int[width];
int* vh = new int[height];
// locations of boundaries between parabolas
T* zw = new T[width + 1];
T* zh = new T[height + 1];
for (int i = 0; i < height; i++) {
for (int j = 0; j < width; j++) {
int index = GET_INDEX(i, j, k, c);
f[j] = output[index];
}
Distance<T>(f, d, vw, zw, width);
for (int j = 0; j < width; j++) {
int index = GET_INDEX(i, j, k, c);
if (Eigen::numext::isinf(d[j])) {
d[j] = Eigen::NumTraits<T>::highest();
}
output[index] = d[j];
}
}
for (int j = 0; j < width; j++) {
for (int i = 0; i < height; i++) {
int index = GET_INDEX(i, j, k, c);
f[i] = output[index];
}
Distance<T>(f, d, vh, zh, height);
for (int i = 0; i < height; i++) {
int index = GET_INDEX(i, j, k, c);
if (Eigen::numext::isinf(d[i])) {
d[i] = Eigen::NumTraits<T>::highest();
}
output[index] = Eigen::numext::sqrt(d[i]);
}
}
delete f;
delete d;
delete vh;
delete vw;
delete zh;
delete zw;
}