api.py 文件源码

def color_mask(image, color, tolerance=0):
    """Extract a mask of image according to color under a certain
    tolerance level (defaults to 0)."""
    if tolerance > 100:
        tolerance = 100
    elif tolerance < 0:
        tolerance = 0
    tolerance = int(tolerance * 255 / 100)
    red, green, blue = color
    bgr_color = np.uint8([[[blue, green, red]]])
    hsv_color = cv2.cvtColor(bgr_color, cv2.COLOR_BGR2HSV)[0][0]
    mask_image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
    lower_range = hsv_color - np.array([tolerance, 0, 0])
    lower_range[lower_range > 255] = 255
    lower_range[lower_range < 0] = 0
    upper_range = hsv_color + np.array([tolerance, 0, 0])
    upper_range[upper_range > 255] = 255
    upper_range[upper_range < 0] = 0
    mask = cv2.inRange(mask_image, lower_range, upper_range)
    return mask