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

Question

给定一个形状为(h, w, c) 的numpy 数组，即具有特定高度、宽度和通道数的图像，以及一个接受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