O(n)的可并行EDT欧式距离变换算法（python实现）

基于论文《A General Algorithm for Computing Distance Transforms in Linear Time》，添加了最近的坐标点记录的功能，精度和scipy.ndimage中的distance_transform_edt保持一致。

运行效果,其中左图为“蓝=indices-x,绿=indices-y,红=border”，右图为“”：

算法伪代码如下：

• 初始化参数：

  • 输入矩阵 b 的大小为 m 行 n 列。
  • 创建并初始化矩阵 g 和 g_indices，元素初始值为 -1。
  • 创建并初始化矩阵 indices，用于存储最近边界点的坐标，初始值为 -1。
  • 创建并初始化矩阵 distances，用于存储最近边界点的距离，初始值为 inf。
  • 创建并初始化数组 s 和 t，用于动态规划过程中的临时存储，初始值为 0。

• 遍历列进行计算（这里是行并行的，每一行都要进行该操作）：

  • 对于每一列 y 从左到右：

    • 如果当前列 b 的元素等于 0，则将对应的 g 元素设置为 n，否则为当前列的 g 元素加 1。

  • 对于每一列 y 从右到左：

    • 如果右边列的 g 元素小于当前列的 g 元素，则将当前列的 g 元素设置为右边列的 g 元素加 1，并更新 g_indices。
    • 否则保持当前列的 g 元素不变。

• 遍历列计算距离和索引（这里是列并行的）：

  • 对于每一列 y：

    • 初始化变量 q 为 0，s[0] 为 0，t[0] 为 0。

    • 定义函数 f(x, i) 计算平方距离，定义函数 Sep(i, u) 计算分离点。

    • 对于每一行 u 从 1 到 m-1：

      • 更新 q，使得 f(t[q], s[q]) 小于等于 f(t[q], u)。
      • 计算分离点 w，如果 w 小于 m，则更新 q、s[q] 和 t[q]。

    • 对于每一行 u 从 m-1 到 0：

      • 计算 distances[u, y] 和 indices[u, y]。
      • 如果 u 等于 t[q]，则递减 q。

• 返回结果：

  • 返回 distances 的平方根和 indices。

代码如下（仅包含算法部分，不含测试代码）：

def our_edtv4(b):
    m, n = b.shape
    g = np.full((m, n), -1, dtype='int32')
    g_indices = np.full((m, n), -1, dtype='int32')
    indices = np.full((m, n, 2), -1, dtype='int32')  # 用于存储最近边界点的坐标
    distances = np.full((m, n), np.inf, dtype='float64')  # 用于存储最近边界点的距离

    g[:, 0] = (b[:, 0] == 0) * n
    # :的部分可以并行
    for y in range(1, n):
        g[:, y] = (b[:, y] == 0) * (1 + g[:, y - 1])

    # :的部分可以并行
    for y in range(n - 2, -1, -1):
        mask = g[:, y + 1] < g[:, y]
        g[:, y] = mask * (1 + g[:, y + 1]) + (1 - mask) * g[:, y]
        g_indices[:, y] = y + (mask * 2 - 1) * g[:, y]

    s, t = np.zeros(m, dtype='int32'), np.zeros(m, dtype='int32')
    # 这里的 y 可以并行
    for y in range(n):
        q = 0
        s[0] = 0
        t[0] = 0

        f = lambda x, i: (x - i) ** 2 + g[i, y] ** 2
        Sep = lambda i, u: (u ** 2 - i ** 2 + g[u, y] ** 2 - g[i, y] ** 2) / (2 * (u - i))

        for u in range(1, m):
            while q >= 0 and f(t[q], s[q]) > f(t[q], u):
                q -= 1
            if q < 0:
                q = 0
                s[0] = u
            else:
                w = 1 + Sep(s[q], u)
                if w < m:
                    q = q + 1
                    s[q] = u
                    t[q] = w
        for u in range(m - 1, -1, -1):
            distances[u, y] = f(u, s[q])
            indices[u, y] = [s[q], g_indices[s[q], y]]
            if u == t[q]:
                q -= 1
    return distances ** 0.5, indices