Given a numpy array of shape (h, w, c), i.e. an image with a certain height, width and number of channels, and a function which takes c input arguments and produces c + k output arguments, how can I map this function over the array to produce an output array of shape (h, w, c + k) in which each "pixel" in the input image is replaced by the function output given the channel values at that pixel.

As an example, I'd like to do something like the following (generalized to more complicated functions):

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)))

1 Answers

Marat On Best Solutions

It looks like you're looking for apply_along_axis:

np.apply_along_axis(f, 2, a)

Note that f() will need to be changed:

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