Векторизация кода 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]])

Answer 2

Используйте mapper

Вы можете создать массив отображения, а затем просто индексирование с помощью входного массива даст нам желаемый результат -

mapper = np.array([0,0,1,2,2])
out = mapper[a] # a is input array

Сроки для всех опубликованных решений с одинаковой заданной формой - (601, 462, 951)) -

In [60]: np.random.seed(0)
    ...: a = np.random.randint(0,5,(601, 462, 951))

# @yatu's soln
In [61]: %timeit np.select([a==1, a==2, (a==3)|(a==4)], [0,1,2])
1 loop, best of 3: 5 s per loop

# Posted in this post
In [62]: %%timeit
    ...: mapper = np.array([0,0,1,2,2])
    ...: out = mapper[a]
1 loop, best of 3: 849 ms per loop

# @Austin's soln
In [63]: %%timeit
    ...: a[a == 1] = 0
    ...: a[a == 2] = 1
    ...: a[(a == 3) | (a == 4)] = 2
1 loop, best of 3: 1.04 s per loop

Дальнейшее усиление с d-типом с низкой точностью

Поскольку выходные данные будут иметь метки - 0,1,2, мы можем смело использовать UINT8 в качестве выходного типа данных и получить огромный прирост производительности. Итак, это будет -

mapper = np.array([0,0,1,2,2],dtype=np.uint8)
out = mapper[a]

Сроки -

In [66]: np.random.seed(0)
    ...: a = np.random.randint(0,5,(601, 462, 951))

In [67]: %%timeit
    ...: mapper = np.array([0,0,1,2,2],dtype=np.uint8)
    ...: out = mapper[a]
1 loop, best of 3: 380 ms per loop

Answer 3

Вы можете сделать:

a[a == 1] = 0
a[a == 2] = 1
a[(a == 3) | (a == 4)] = 2

, где a - ваш массивный массив.