imageutil.py 文件源码

def add_channel(image, channelfirst):
    """
    Add channel if missing and make first axis if requested.

    >>> import numpy as np
    >>> image = np.ones((10, 20))
    >>> image = add_channel(image, True)
    >>> shapestr(image)
    '1x10x20'

    :param ndarray image: RBG (h,w,3) or gray-scale image (h,w).
    :param bool channelfirst: If True, make channel first axis
    :return: Numpy array with channel (as first axis if makefirst=True)
    :rtype: numpy.array
    """
    if not 2 <= image.ndim <= 3:
        raise ValueError('Image must be 2 or 3 channel!')
    if image.ndim == 2:  # gray-scale image
        image = np.expand_dims(image, axis=-1)  # add channel axis
    return np.rollaxis(image, 2) if channelfirst else image