"""Functions for image masking.
"""
import cv2
import numpy as np
import logging
logger = logging.getLogger('lavaflow')
# -----------------------------------------------------------------------------
# Image masking
[docs]class ColorMask(object):
"""Class for creating an image mask based on color matching.
Each pixel is converted into the desired color space for color matching using
``cv2.cvtColor`` based on the color conversion code and each color channel is
checked to see if the value falls within ``sensitivity`` of the ``center``
over the range ``[0, 255]``.
The LAB color space is good for color matching as it corresponds to colors as
they are perceived by the human eye, L for perceptual lightness and a* and b*
for the four unique colors of human vision red, green, blue, and yellow. This
means that a numerical change in LAB space corresponds to a similar perceived
change in color which makes it easier to choose the center and sensitivity.
"""
def __init__(self, center, sensitivity, code=None, output_type='mask'):
"""
Args:
center (np.ndarray): center color (same color space as the image)
sensitivity (np.ndarray): relative sensitivity (separate for each color channel)
code (int): color conversion code e.g. ``cv2.COLOR_BGR2LAB``
output_type (str): output type for map
"""
self.code = code
if self.code is not None:
self.center = cv2.cvtColor(np.array(center).reshape(1, 1, 3).astype(np.uint8), self.code).flatten().astype(np.uint8)
else:
self.center = np.array(center).astype(np.uint8)
self.sensitivity = np.array(sensitivity)
self.lower_bound = (self.center - 255 * self.sensitivity).clip(0, 255).astype(np.uint8)
self.upper_bound = (self.center + 255 * self.sensitivity).clip(0, 255).astype(np.uint8)
self.widths = self.upper_bound - self.lower_bound
self.distance_threshold = 255 * np.sqrt(3)
if output_type in ['mask', 'distance']:
self.output_type = output_type
else:
raise Exception('output type not supported')
[docs] def __call__(self, img):
"""Map image to image mask or distance map.
Args:
img (np.ndarray): image
Returns:
obj (np.ndarray): image mask or distance map
"""
if self.output_type == 'mask':
return self.mask(img)
elif self.output_type == 'distance':
return self.distance(img)
else:
raise Exception('output type not supported')
[docs] def mask(self, img):
"""Map image to image mask.
Args:
img (np.ndarray): image
Returns:
mask (np.ndarray): image mask
"""
if self.code is not None:
tmp = cv2.cvtColor(img, self.code)
else:
tmp = img
return cv2.inRange(tmp, self.lower_bound, self.upper_bound)
[docs] def distance(self, img, channels=(0, 1, 2)):
"""Map image ot distance map.
Args:
img (np.ndarray): image
Returns:
distance (np.ndarray): distance map
"""
if self.code is not None:
tmp = cv2.cvtColor(img, self.code)
else:
tmp = img
return np.linalg.norm((tmp[:, :, channels].astype(np.float64) - self.center) * 255 / self.widths, axis=2) / self.distance_threshold
[docs]class MultiColorMask(object):
"""Class for creating an image mask based on multiple color matching.
This class combines outputs from multiple :class:``lavaflow.masking.ColorMask``
either by taking the union of image masks or the minimum of the distance maps.
"""
def __init__(self, color_masks, output_type='mask'):
"""
Args:
color_masks (list): array of color masks
output_type (str): output type for map
Returns:
self (ColorMask): color mask
"""
self.color_masks = color_masks
if output_type in ['mask', 'distance']:
self.output_type = output_type
else:
raise Exception('output type not supported')
[docs] def __call__(self, img):
"""Map image to image mask or distance map.
Args:
img (np.ndarray): image
Returns:
obj (np.ndarray): image mask or distance map
"""
if self.output_type == 'mask':
return self.mask(img)
elif self.output_type == 'distance':
return self.distance(img)
else:
raise Exception('output type not supported')
[docs] def mask(self, img):
"""Mask image to image mask.
Args:
img (np.ndarray): image
Returns:
mask (np.ndarray): image mask
"""
return np.dstack([color_mask.mask(img) for color_mask in self.color_masks]).max(axis=2)
[docs] def distance(self, img, channels=(0, 1, 2)):
"""Map image ot distance map.
Args:
img (np.ndarray): image
Returns:
distance (np.ndarray): distance map
"""
return np.dstack([color_mask.distance(img) for color_mask in self.color_masks]).min(axis=2)