如何实现沿轴的numpy.where()等效操作？优化图像栈最大强度投影

我需要对图像栈进行最大强度投影（maximum intensity projections），即保留每个像素在所有图层中的最亮（数值最高）点。原先通过遍历每个像素的循环实现，但速度较慢，希望利用NumPy数组级函数优化代码。

示例图像栈 #

import numpy as np

testLayer1 = [
    [0,0,0],
    [0,0,0],
    [0,0,0]
]
testLayer2 = [
    [1,1,1],
    [1,1,1],
    [1,1,1]
]
testLayer3 = [
    [0,1,2],
    [3,4,5],
    [6,7,8]
]
testLayer4 = [
    [0,0,0],
    [0,0,0],
    [0,0,10]
]
testStack = np.array([testLayer1,testLayer2,testLayer3,testLayer4])

原循环实现（速度慢） #

def MakeMaxIntProj(stack):
    maxIntProj = np.zeros((2,stack.shape[1],stack.shape[2]))

    # 遍历所有像素
    for i in range(stack.shape[1]):
        for j in range(stack.shape[2]):
            maxIntProj[0,i,j] = np.max(stack[:,i,j]) # 获取像素(i,j)的最大亮度值
            maxIntProj[1,i,j] = np.where(stack[:,i,j] == maxIntProj[0,i,j])[0][0] # 获取对应图层索引
    
    return maxIntProj

原方法输出 #

In [1]: MakeMaxIntProj(testStack)
Out[1]: array([[[ 1.,  1.,  2.],
                [ 3.,  4.,  5.],
                [ 6.,  7., 10.]],

               [[ 1.,  1.,  2.],
                [ 2.,  2.,  2.],
                [ 2.,  2.,  3.]]])

优化需求 #

已通过np.max(stack, axis=0)快速得到每个像素的最大值，但需要高效获取对应最大值的图层索引，替代循环部分提升效率。当前未完成的优化代码：

def MakeMaxIntProj(stack):
    maxIntProj = np.zeros((2,stack.shape[1],stack.shape[2]))
    
    maxIntProj[0,:,:] = np.max(stack,axis=0)
    maxIntProj[1,:,:] = 0 # 此处需要替换为高效获取索引的方法
    
    return maxIntProj

解决方案：使用np.argmax() #

NumPy的np.argmax()函数可以直接沿指定轴返回最大值的索引，完全匹配原代码中找第一个最大值对应图层的逻辑，且是数组级操作，速度远快于循环。

优化后的完整代码 #

import numpy as np

def MakeMaxIntProj(stack):
    maxIntProj = np.zeros((2, stack.shape[1], stack.shape[2]))
    
    # 获取每个像素的最大强度值
    maxIntProj[0, :, :] = np.max(stack, axis=0)
    # 获取对应最大值所在的图层索引
    maxIntProj[1, :, :] = np.argmax(stack, axis=0)
    
    return maxIntProj

验证结果 #

运行优化后的代码，输出与原循环方法完全一致：

In [3]: MakeMaxIntProj(testStack)
Out[3]: array([[[ 1.,  1.,  2.],
                [ 3.,  4.,  5.],
                [ 6.,  7., 10.]],

               [[ 1.,  1.,  2.],
                [ 2.,  2.,  2.],
                [ 2.,  2.,  3.]]])

说明 #

np.argmax()默认返回第一个出现最大值的索引，和原代码中np.where(...)[0][0]的行为完全一致。如果多个图层有相同的最大值，会取索引最小的那个图层，与原逻辑匹配。

内容的提问来源于stack exchange，提问作者Nathan