Ускорение индексации целочисленных массивов для глубины

Question

Предположим, у меня есть массив

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

и я хочу преобразовать его в тензор вида

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

Где каждый слой глубины (индекс i) - это двоичная маска, показывающая, где на входе появляется i.

Я написал код для этого, который работает правильно, но слишком медленно для любого использования. Могу ли я заменить цикл в этой функции другой векторизованной операцией?

def im2segmap(im, depth):
    tensor = np.zeros((im.shape[0], im.shape[1], num_classes))

    for c in range(depth):
        rows, cols = np.argwhere(im==c).T
        tensor[c, rows, cols] = 1

    return tensor

Answer 1

Использование broadcasting -

(a==np.arange(num_classes)[:,None,None]).astype(int)

Или с builtin внешним сравнением -

(np.equal.outer(range(num_classes),a)).astype(int)

Используйте uint8, если вам нужно использовать int dtype или оставить как boolean, пропустив конверсию int для дальнейшего повышения.

Пробный прогон -

In [42]: a = np.array([[0,2,1],[1,0,1],[2,1,1]])

In [43]: num_classes = 3 # or depth

In [44]: (a==np.arange(num_classes)[:,None,None]).astype(int)
Out[44]: 
array([[[1, 0, 0],
        [0, 1, 0],
        [0, 0, 0]],

       [[0, 0, 1],
        [1, 0, 1],
        [0, 1, 1]],

       [[0, 1, 0],
        [0, 0, 0],
        [1, 0, 0]]])

Чтобы иметь depth/num_classes в качестве третьего затемнения, расширьте входной массив, а затем сравните с массивом диапазона -

(a[...,None]==np.arange(num_classes)).astype(int)
(np.equal.outer(im, range(num_classes))).astype(int)
(np.equal.outer(im, range(num_classes))).astype(np.uint8) # lower prec