Разверните массив пустот по новому измерению, используя значения в качестве индекса

Question

У меня есть массив [m,m] numpy с элементом в {0, 1, 2, ..., 24}, теперь я хочу разделить каждое число в третьем измерении, чтобы получить массив [m,m,24].

простой пример, [5,5] массив с элементом в {0, 1, 2, 3}

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

Теперь мне нужно получить массив `[5,5,3]`

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

В настоящее время у меня есть простой метод, но он очень дорогой для вычислений. Потому что мне нужно часто делать эту операцию.

img = np.expand_dims(img, axis=2)
for i in range(24):
    img_norm[..., i] = (img[..., 0] == (i + np.ones(shape=img[..., 0].shape)))

Для массива 64 с размером [224,224] и элементом в {0, 1, 2, ..., 24}, приведенный выше код занимает около 5s.

Есть ли более быстрый способ сделать это?

Answer 1

Определенно более элегантно: используйте np.ndenumerate().

for (i,j), val in np.ndenumerate(img):
    img_norm[val-1,i,j] = val

Похоже, это должно быть быстрее, чем у вас, потому что O (N ^ 2), а не O (N ^ 3). Давайте попробуем это на массиве с размером и содержанием, как вы описываете:

def getnorm_ndenumerate(img):
    img_norm = np.zeros([np.max(img), *img.shape])
    for (i,j), val in np.ndenumerate(img):
        img_norm[val-1,i,j] = val  
    return img_norm

b = np.int64(25*np.random.rand(224, 224)) 

%timeit getnorm_ndenumerate(b)

дает

47.8 ms ± 1.38 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

Это действительно быстрее, чем у вас. Но элегантность имеет свою цену, потому что она медленнее, чем метод acdr .

Answer 2

Я сделал ошибку, в выходном массиве все ненулевые значения должны быть равны 1. Извините за мою глупую ошибку.

спасибо за вашу помощь. Я протестировал три метода выше, включая код из Jean-François Corbett, acdr + Jean-François Corbett и мой. Оказывается, метод из acdr + Jean-François Corbett самый быстрый.

Вот мой тестовый код

def test_time():
    def func1(img, max_num):
        w, h = img.shape
        img_norm = np.zeros([w, h, max_num], np.float32)
        for (i, j), val in np.ndenumerate(img):
            # img_norm[i, j, val - 1] = val
            img_norm[i, j, val - 1] = 0 if val == 0 else 1
        return img_norm

    def func2(img, max_num):
        w, h = img.shape
        img_norm = np.zeros([w, h, max_num], np.float32)
        for idx in range(1, max_num + 1):
            # img_norm[:, :, idx - 1] = idx*(img == idx)
            img_norm[:, :, idx - 1] = (img == idx)
        return img_norm

    def func3(img, max_num):
        w, h = img.shape
        img_norm = np.zeros([w, h, max_num], np.float32)
        for idx in range(max_num):
            # img_norm[:, :, idx] = (idx+1) * (img[:, :, 0] == (idx + np.ones(shape=img[:, :, 0].shape)))
            img_norm[:, :, idx] = (img == (idx + np.ones(shape=img.shape)))
        return img_norm

    import cv2
    img_tmp = cv2.imread('dat.png', cv2.IMREAD_UNCHANGED)
    img_tmp = np.asarray(img_tmp, np.int)

    # img_tmp = np.array([
    #     [0, 0, 1, 0, 0],
    #     [2, 0, 3, 0, 1],
    #     [0, 2, 3, 1, 0],
    #     [0, 0, 1, 0, 0],
    #     [1, 0, 2, 0, 1],
    # ])

    img_bkp = np.array(img_tmp, copy=True)
    print(img_bkp.shape)
    import time
    cnt = 100
    maxnum = 24
    start_time = time.time()
    for i in range(cnt):
        _ = func1(img_tmp, maxnum)
    print('1 total time =', time.time() - start_time)

    start_time = time.time()
    for i in range(cnt):
        _ = func2(img_tmp, maxnum)
    print('2 total time =', time.time() - start_time)

    start_time = time.time()
    for i in range(cnt):
        _ = func3(img_tmp, maxnum)
    print('3 total time =', time.time() - start_time)

    print((img_tmp == img_bkp).all())
    img1 = func1(img_tmp, maxnum)
    img2 = func2(img_tmp, maxnum)
    img3 = func3(img_tmp, maxnum)
    print(img1.shape, img2.shape, img3.shape)
    print((img1 == img2).all())
    print((img2 == img3).all())
    print((img1 == img3).all())
    # print(type(img1[0, 0, 0]), type(img2[0, 0, 0]), type(img3[0, 0, 0]))
    # print('img1\n', img1[:, :, 2])
    # print('img3\n', img3[:, :, 2])

Выход

    (224, 224)
    1 total time = 4.738261938095093
    2 total time = 0.7725710868835449
    3 total time = 1.5980615615844727
    True
    (224, 224, 24) (224, 224, 24) (224, 224, 24)
    True
    True
    True

Если есть какие-либо проблемы, пожалуйста, оставьте это в комментариях.

Спасибо за вашу помощь!

Answer 3

Следующее довольно быстро для меня:

import numpy as np
max_num = 3
img = np.array([
    [0,0,1,0,0],
    [2,0,3,0,1],
    [0,2,3,1,0],
    [0,0,1,0,0],
    [1,0,2,0,1],
    ])

img_norm = np.zeros(img.shape + (max_num,))
for idx in range(1, max_num + 1):
    img_norm[idx-1,:,:]=idx*(img == idx)

Тестирование со случайным массивом указанного вами размера;

max_num = 24
img = np.int64((max_num+1)*np.random.rand(224, 224)) # Random array

img_norm = np.zeros(img.shape + (max_num,))
for idx in range(1, max_num + 1):
    img_norm[idx-1,:,:]=img*(img == idx)

На моей машине вряд ли понадобится какое-то время.

def getnorm_acdr(img):
    max_num = np.max(img)
    img_norm = np.zeros([max_num, *img.shape])    
    for idx in range(1, max_num + 1):
        img_norm[idx-1,:,:]=img*(img == idx)

img = np.int64((max_num+1)*np.random.rand(224, 224))

%timeit getnorm_acdr(img)

Дает:

11.9 ms ± 536 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)