   // voronoi diagram computation in thrust
   {
     // compute_temp_vector = compute_vector;
-    jfa_cuda << <dimGrid, dimBlock>>>
-        (thrust::raw_pointer_cast(compute_vector.data()), 1, image_size,
-         image_size);
+    distanceTransformStep<<<dimGrid, dimBlock>>>(
+        thrust::raw_pointer_cast(compute_vector.data()), 1, image_size,
+        image_size);
     // compute_vector.swap(compute_temp_vector);
 
     // cudaThreadSynchronize();
     for (int k = (image_size / 2); k > 0; k /= 2)
     {
-      jfa_cuda<<<dimGrid, dimBlock>>>(
+      distanceTransformStep<<<dimGrid, dimBlock>>>(
           thrust::raw_pointer_cast(compute_vector.data()), k, image_size,
           image_size);
       // cudaThreadSynchronize();

  data[index] = pixelValue >= 0.99f ? index : outlierValue;
}

__global__ void jfa_cuda(int *data, unsigned int step, int w, int h)
{
  int x = blockIdx.x * blockDim.x + threadIdx.x;
  int y = blockIdx.y * blockDim.y + threadIdx.y;
  if (x >= w || y >= h)
    return;

  int index = y * w + x;

  int cur_nearest = data[index];
  int cur_y = cur_nearest / w;
  int cur_x = cur_nearest - cur_y * w;
  int cur_dist = (x - cur_x) * (x - cur_x) + (y - cur_y) * (y - cur_y);

// if (x>253 && x<257 && y==255)
//  if ((x==255 && y==255) || (x==10 && y==10))
//  {
//    printf("%d %d: step: %d cur: %d %d %d, %d %d\n", x, y, step, cur_nearest,
//    cur_x, cur_y, w, h);
//  }

#pragma unroll
  for (int i = -1; i <= 1; i++)
  {
    int u = x + i * step;
    if (u < 0 || u >= w)
      continue;
#pragma unroll
    for (int j = -1; j <= 1; j += 2 - i * i)
    {
      int v = y + j * step;
      if (v < 0 || v >= h)
        continue;

      int newindex = v * w + u;
      int new_nearest = data[newindex];
      int new_y = new_nearest / w;
      int new_x = new_nearest - new_y * w;
      int new_dist = (x - new_x) * (x - new_x) + (y - new_y) * (y - new_y);

      if (new_dist < cur_dist)
      {
        cur_dist = new_dist;
        cur_nearest = new_nearest;
      }
    }
  }

  data[index] = cur_nearest;
}

__global__ void jfa_dtf_compute_distance_kernel(int image_size, int *index_ptr,


  int index = y * w + x;

  int currentNearest = data[index];
  int currentY = currentNearest / w;
  int currentX = currentNearest - currentY * w;
  float curr_w = (currentNearest < w * h) ? occupancy[currentNearest] : 0.0f;

  float currentDistance =
      sqrtf(float((x - currentX) * (x - currentX) + (y - currentY) * (y - currentY))) -
      curr_w;

#pragma unroll
