EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE void EuclideanDistanceTransformSample()

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;
}