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
评论列表
文章目录