python - 向量化用于修改图像分割中使用的掩码的 python 代码

Question

我正在应对图像分割挑战。我的面具有 5 个标签(0、1、2、3、4)，其中一个面具(2D 矩阵)的布局是:

[0 0 0 0 0 0
 0 0 1 1 2 2
 1 1 1 1 2 2
 2 2 2 2 2 2
 3 3 3 3 3 3
 4 4 4 4 4 4]

我想合并几个类，使修改后的掩码看起来像:

[0 0 0 0 0 0
 0 0 0 0 1 1
 0 0 0 0 1 1
 1 1 1 1 1 1
 2 2 2 2 2 2
 2 2 2 2 2 2]

将 0 和 1 合并为 0。 2 变成 1。 合并 3 和 4，将它们更改为 2。

我已经实现了一个循环版本，它花费了很多时间，因为我的面具尺寸是 (601, 462, 951)。

for i in range(0, dim.shape[0]):
  for j in range(0, dim.shape[1]):
    for k in range(0, dim.shape[2]):
      if dim[i, j, k] in (0, 2):
        dim[i, j, k] = 1

      if dim[i, j, k] == 3:
        dim[i, j, k] = 2

      if dim[i, j, k] in (4, 8):
        dim[i, j, k] = 3

      if dim[i, j, k] == 9:
        dim[i, j, k] = 4

我找不到任何方法来矢量化我的代码，以便我可以删除循环。

Answer 1

您可以使用 np.select一个简洁的解决方案，它允许您在给定条件列表的情况下从 choicelist 中进行选择:

np.select([a==1, a==2, (a==3)|(a==4)], [0,1,2])

array([[0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 1, 1],
       [0, 0, 0, 0, 1, 1],
       [1, 1, 1, 1, 1, 1],
       [2, 2, 2, 2, 2, 2],
       [2, 2, 2, 2, 2, 2]])