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