def apply_to_channels(single_channel_func):
    """Creates a new function which operates on each channel

    Parameters
    ----------
    single_channel_func: function
        Function that acts on a single color channel

    Returns
    -------
    channel_func: function
        The same function that operates on all color channels

    Example
    -------
    >>> from pymatting import *
    >>> import numpy as np
    >>> from scipy.signal import convolve2d
    >>> single_channel_fun = lambda x: convolve2d(x, np.ones((3, 3)), 'valid')
    >>> multi_channel_fun = apply_to_channels(single_channel_fun)
    >>> I = np.random.rand(480, 320, 3)
    >>> multi_channel_fun(I).shape
    (478, 318, 3)
    """

    @wraps(single_channel_func)
    def multi_channel_func(image, *args, **kwargs):
        if len(image.shape) == 2:
            return single_channel_func(image, *args, **kwargs)
        else:
            shape = image.shape
            image = image.reshape(shape[0], shape[1], -1)

            result = np.stack(
                [
                    single_channel_func(image[:, :, c].copy(), *args, **kwargs)
                    for c in range(image.shape[2])
                ],
                axis=2,
            )

            return result.reshape(list(result.shape[:2]) + list(shape[2:]))

    return multi_channel_func
Creates a new function which operates on each channel

    Parameters
    ----------
    single_channel_func: function
        Function that acts on a single color channel

    Returns
    -------
    channel_func: function
        The same function that operates on all color channels

    Example
    -------
    >>> from pymatting import *
    >>> import numpy as np
    >>> from scipy.signal import convolve2d
    >>> single_channel_fun = lambda x: convolve2d(x, np.ones((3, 3)), 'valid')
    >>> multi_channel_fun = apply_to_channels(single_channel_fun)
    >>> I = np.random.rand(480, 320, 3)
    >>> multi_channel_fun(I).shape
    (478, 318, 3)
    
pymatting/util/util.py