使用值作为索引沿新维度折叠一个numpy数组

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]`数组

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

目前，我有一个简单的方法，但是它的计算量非常大。 因为我需要经常执行此操作。

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

对于大小为[224,224]且元素位于{0, 1, 2, ..., 24} 64数组，上​​面的代码大约需要5s 。

有更快的方法吗？

Answer 1

以下内容对我来说非常快捷：

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)

Answer 2

绝对更优雅：使用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 3

我犯了一个错误，在输出数组中，所有非零都应该为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 

如果有任何问题，请在评论中发布。

多谢您的协助！