"""Utility functions for OpenCV.
"""
import cv2
import itertools
import numpy as np
from lavaflow.utils.geometry import roundPoint
import logging
logger = logging.getLogger('lavaflow')
# -----------------------------------------------------------------------------
# Masking
[docs]def applyMask(img, mask, offset=0, radius=0, spread=0):
"""Apply mask to image.
The mask can be dilated or eroded using ``cv2.morphologyEx`` and smoothed
using ``cv2.GaussianBlue``. A positive offset uses dilation and a negative
offset uses erosion to expand or contract the mask respectively.
Args:
img (np.ndarray): image
mask (np.ndarray): monocrhome mask
offset (int): number of pixels to offset mask
radius (int): blur radius (must be odd)
spread (float): blur spread
Returns:
img (np.ndarray): image with mask applied
"""
mask = mask.astype(float)
mask[mask > 0] = 255
if radius > 0 and spread > 0:
mask = cv2.GaussianBlur(mask, (radius, radius), spread, None, spread)
if offset > 0:
size = 1 + offset * 2
mask = cv2.morphologyEx(mask, cv2.MORPH_DILATE, np.ones((size, size), np.uint8))
if offset < 0:
size = 1 - offset * 2
mask = cv2.morphologyEx(mask, cv2.MORPH_ERODE, np.ones((size, size), np.uint8))
img = img * np.expand_dims((mask / 255), axis=2)
return img.astype(np.uint8)
# -----------------------------------------------------------------------------
# Visualization helpers
[docs]def ensureColor(img):
"""Ensure image is BGR not GRAY.
Args:
img (np.ndarray): image
Returns:
img (np.ndarray): image converted from GRAY to BGR
"""
if len(img.shape) < 3:
img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
if img.dtype != np.uint8:
img = img.astype(np.uint8)
return img
# -----------------------------------------------------------------------------
# Visualization wrappers
[docs]def immask(img, mask, color=(255, 255, 0), alpha=0.5):
"""Function to apply a transparent mask overlay to an image
Args:
img (np.ndarray): image
mask (np.ndarray): mask to overlay
color (tuple): mask color
alpha (float): mask transparency
Returns:
background (np.ndarray): img with mask applied
"""
if img.shape != mask.shape:
raise ValueError(f"img.shape {img.shape} != mask.shape {mask.shape}")
background = cv2.cvtColor(img, cv2.COLOR_BGR2BGRA)
foreground = np.dstack([np.multiply((cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR) / 255), np.array(color)), (mask * alpha)]).astype(np.uint8)
alpha_background = background[:,:,3] / 255.0
alpha_foreground = foreground[:,:,3] / 255.0
for i in range(0, 3):
background[:,:,i] = alpha_foreground * foreground[:,:,i] + \
alpha_background * background[:,:,i] * (1 - alpha_foreground)
background[:,:,3] = (1 - (1 - alpha_foreground) * (1 - alpha_background)) * 255
return background
[docs]def imgrid(array, shape=None, max_size=(2160, 3840)):
"""Function for composing multiple images into a grid.
Args:
array (list): list of images to compose
shape (tuple): grid shape (rows, columns)
max_size (tuple): resolution used to scale image grid
Returns:
grid (np.ndarray): image grid
"""
array = [ensureColor(img) for img in array]
n = len(array)
if shape is not None:
p, q = shape
m = p * q
else:
q = int(np.ceil(np.sqrt(n)))
p = int(np.ceil(n/q))
m = p * q
(h, w, d) = array[0].shape
grid = np.hstack(array + [np.zeros((h, w, d)).astype(np.uint8)] * (m - n)).reshape(h, -1, w*q, 3).swapaxes(0, 1).reshape(h * p, w * q, d)
maxh, maxw = max_size
ratio = min(maxw / (w*q), maxh / (h*p))
print(ratio)
if ratio < 1:
return cv2.resize(grid, None, fx=ratio, fy=ratio, interpolation=cv2.INTER_AREA)
else:
return grid
[docs]def imshow(img):
"""Wrapper for OpenCV imshow.
Args:
img (np.ndarray): image
"""
cv2.imshow('img', ensureColor(img))
cv2.waitKey(0)
cv2.destroyAllWindows()
# -----------------------------------------------------------------------------
# Drawing wrappers
[docs]def drawPoints(vis, v, color=(255, 255, 255), points=1, line=0, loop=1, size=3, thickness=1):
"""Draw points on an image.
Args:
vis (np.ndarray): visualization image
v (np.ndarray): array of points, shaped (N, 2)
color (tuple): color
points (int): flag to draw points
line (int): flag to connect points with lines
loop (int): flag to connect end point with start point (ignored if line = 0)
size (int): point size
thickness (int): line thickness
Returns:
vis (np.ndarray): visualization image
"""
v = v.reshape(-1, 2)
if line:
for i in range(1 - loop, len(v)):
cv2.line(vis, roundPoint(v[i-1]), roundPoint(v[i]), color, thickness)
if points:
for point in v:
cv2.circle(vis, roundPoint(point), size, color, -1)
[docs]def drawSegments(vis, s, color=(255, 255, 255), points=0, line=1, size=3, thickness=1):
"""Draw line segments on an image.
Args:
vis (np.ndarray): visualization image
s (np.ndarray): array of line segments, shaped (N, 2, 2)
color (tuple): color
points (int): flag to draw points
line (int): flag to connect points with lines
size (int): point size
thickness (int): line thickness
Returns:
vis (np.ndarray): visualization image
"""
if line:
for a, b in s:
cv2.line(vis, roundPoint(a), roundPoint(b), color, thickness)
if points:
for point in np.unique(s.reshape(-1, 2), axis=0):
cv2.circle(vis, roundPoint(point), size, color, -1)
[docs]def drawTriangles(vis, tri, color=(255, 255, 255), segment_color=(0, 0, 255), points=1, edges=1, segments=0, size=3, thickness=1):
"""Draw triangulation on an image.
Args:
vis (np.ndarray): visualization image
tri (dict): output from triangle.triangulate
color (tuple): color
segment_color (tuple): color for segment constraints
points (int): flag to draw triangle points
edges (int): flag to draw triangle edges
segments (int): flag to draw segment constraints
size (int): point size
thickness (int): line thickness
Returns:
vis (np.ndarray): visualization image
"""
if points:
for point in tri['vertices']:
cv2.circle(vis, roundPoint(point), size, color, -1)
t = tri['triangles']
s = np.unique(np.sort(np.vstack([np.hstack([t[:, i - 1, np.newaxis], t[:, i, np.newaxis]]) for i in range(0, 3)]), axis=1), axis=0)
if edges:
for a, b in tri['vertices'][s]:
cv2.line(vis, roundPoint(a), roundPoint(b), color, thickness)
if points:
cv2.circle(vis, roundPoint(a), size, color, -1)
cv2.circle(vis, roundPoint(b), size, color, -1)
if segments:
for a, b in tri['vertices'][tri['segments']]:
cv2.line(vis, roundPoint(a), roundPoint(b), segment_color, thickness)
if points:
cv2.circle(vis, roundPoint(a), size, segment_color, -1)
cv2.circle(vis, roundPoint(b), size, segment_color, -1)
[docs]def drawVelocityGrid(vis, x, y, u, v, color=(255, 255, 255), spacing=1, magnification=1, points=0, arrows=1, size=3, thickness=1, tipLength=0.05):
"""Draw (N, M) velocity grid on an image.
Args:
vis (np.ndarray): visualization image
x (np.ndarray): grid x coordinates, shaped (N, M)
y (np.ndarray): grid y coordinates, shaped (N, M)
u (np.ndarray): velocity horizontal components, shaped (N, M)
v (np.ndarray): velocity vertical components, shaped (N, M)
spacing (double): minimum spacing
magnification (float): magnification to apply to velocity
color (tuple): color
points (int): flag to draw points
arrows (int): flag to draw arrows
size (int): point size
thickness (int): line thickness
tipLength (int): arrow tip length
Returns:
vis (np.ndarray): visualization image
"""
if not arrows:
tipLength = 0
x_spacing = x[1, 1] - x[0, 0]
y_spacing = y[1, 1] - y[0, 0]
x_step = int(np.ceil(spacing / x_spacing))
y_step = int(np.ceil(spacing / y_spacing))
for i, j in itertools.product(np.arange(0, x.shape[0], x_step), np.arange(0, x.shape[1], y_step)):
a = (x[i, j], y[i, j])
b = ((x[i, j] + magnification * u[i, j]), (y[i, j] + magnification * v[i, j]))
if not np.isnan(b).any():
cv2.arrowedLine(vis, roundPoint(a), roundPoint(b), color, thickness, tipLength=tipLength)
if points:
cv2.circle(vis, roundPoint(a), size, color, -1)
[docs]def drawVelocityLine(vis, x, y, u, v, color=(255, 255, 255), spacing=1, magnification=1, points=0, arrows=1, size=3, thickness=1, tipLength=0.05):
"""Draw (N, M) velocity grid on an image.
Args:
vis (np.ndarray): visualization image
x (np.ndarray): line x coordinates, shaped (N)
y (np.ndarray): line y coordinates, shaped (N)
u (np.ndarray): velocity horizontal components, shaped (N)
v (np.ndarray): velocity vertical components, shaped (N)
spacing (double): minimum spacing
magnification (float): magnification to apply to velocity
color (tuple): color
points (int): flag to draw points
arrows (int): flag to draw arrows
size (int): point size
thickness (int): line thickness
tipLength (int): arrow tip length
Returns:
vis (np.ndarray): visualization image
"""
if not arrows:
tipLength = 0
line_spacing = np.linalg.norm((x[1] - x[0], y[1] - y[0]))
step = int(np.ceil(spacing / line_spacing))
for i in np.arange(0, x.shape[0], step):
a = (x[i], y[i])
b = ((x[i] + magnification * u[i]), (y[i] + magnification * v[i]))
if not np.isnan(b).any():
cv2.arrowedLine(vis, roundPoint(a), roundPoint(b), color, thickness, tipLength=tipLength)
if points:
cv2.circle(vis, roundPoint(a), size, color, -1)