common.py 文件源码

def mask_image(lower_mask, upper_mask, img, apply_mask=False):
    """" Masks an image according to the upper and lower bounds
    Parameters
    ----------
    lower_mask : ndarray
        lower mask to apply to image, length must match image channels
    upper_mask : ndarray
        upper mask to apply to image, length must match image channels
    img :  ndarray
        image to apply mask to
    apply_mask : bool
        returns the masked image instead of the mask itself
    """
    shape = np.array(img.shape).flatten()
    if len(np.array(img.shape).flatten()) == 3:
        shape_size = shape[-1]
    else:
        shape_size = 1
    assert (len(lower_mask) == shape_size)
    assert (len(upper_mask) == shape_size)
    color_min = np.array(lower_mask, np.uint8)
    color_max = np.array(upper_mask, np.uint8)
    new_img = cv2.inRange(img, color_min, color_max)
    if apply_mask:
        return cv2.bitwise_and(img, img, mask=new_img)
    return new_img