utils.py 文件源码

def output_as_mask(f):
    """Decorator to add a return_mask option when image and mask are being
    returned"""

    def wrapper(*args, **kwargs):
        return_mask = kwargs.pop('return_mask', False)
        image, mask = f(*args, **kwargs)
        if return_mask:
            return mask
        else:
            return cv2.bitwise_and(image, image, mask=mask)

    wrapper.__name__ = f.__name__
    wrapper.__doc__ = """{}
    The binary mask can also be returned instead by setting
    return_mask to True.""".format(f.__doc__)
    return wrapper