"""Pipes for improving boundary detection and masking.
"""
import cv2
from lavaflow.pipes.core import AbstractPipe
import logging
logger = logging.getLogger('lavaflow')
# -----------------------------------------------------------------------------
# Pipes to render a visual representation of outputs
[docs]class VisualizeBoundaryPipe(AbstractPipe):
def __init__(self, x, boundary_detector, contours):
"""Constructor.
Args:
x (str): input key
boundary_detector (dict): boundary detector to draw
contours (list): list of contours to draw
Returns:
self (ImagePipe): image pipe
"""
self.x = x
self.boundary_detector = boundary_detector
self.contours = contours
[docs] def map(self, d):
"""Apply image processing.
Args:
d (np.ndarray): image to be processed
Returns:
d (np.ndarray): processed image
"""
vis = d[self.x].copy()
if self.boundary_detector.flow_source == 'point':
cv2.circle(vis, self.boundary_detector.point, self.boundary_detector.radius, (0, 255, 0), 2)
elif self.boundary_detector.flow_source == 'edge':
raise Exception('flow source edge not yet implemented')
else:
raise Exception('flow source type not supported')
index = d['index']
frame = d['frame']
cv2.putText(vis, f'index: {index}, frame: {frame}', (10, 20), cv2.FONT_HERSHEY_PLAIN, 1, (0, 0, 255), 1, cv2.LINE_8)
for k, color, thickness in self.contours:
cv2.drawContours(vis, [d[k]], -1, color, thickness)
d['vis'] = vis
return d