numpy：在图像通道上映射函数的惯用方式

Question

给出一个形状为(h, w, c)的数字数组，即具有一定高度，宽度和通道数的图像，以及一个接受c输入参数并产生c + k输出参数的函数，如何将这个函数映射到数组上以生成形状为(h, w, c + k)的输出数组，其中输入图像中的每个“像素”都被给定了该像素的通道值的函数输出所取代。

作为一个例子，我想做以下事情（概括为更复杂的功能）：

a = np.ones((10, 10, 3))

def f(r, b, g):
    return r, g, b, r + g + b

# TODO: map f over a resulting in np.dstack((a, np.full((10, 10), 3)))

Answer 1

您似乎正在寻找apply_along_axis：

np.apply_along_axis(f, 2, a)

请注意，f()将需要更改：

def f(col):
    r, g, b = col
    return r, g, b, r + g + b