diff --git a/src/ar_sandbox/examples/examples.py b/src/ar_sandbox/examples/examples.py index 2a762e0dbbff1b6e786e8b177788320bd6cc1046..83cb4f9fc3bed7527c92fe163bf0889a447bf121 100644 --- a/src/ar_sandbox/examples/examples.py +++ b/src/ar_sandbox/examples/examples.py @@ -1,51 +1,60 @@ +from typing import Any import numpy as np COLORS = [ - (0, [255, 0, 0]), - (0.25, [255, 255, 0]), - (0.5, [0, 255, 0]), - (0.75, [0, 255, 255]), - (1, [0, 0, 255]) + (0, (255, 0, 0)), + (0.25, (255, 255, 0)), + (0.5, (0, 255, 0)), + (0.75, (0, 255, 255)), + (1, (0, 0, 255)) ] class LevelDisplay(): - @staticmethod - def get_color(depth_matrix, frame, draw_lines=True, colormap=COLORS): + def __init__(self, colormap:list = COLORS) -> None: """ - Given a depth matrix between 0-1, return a color from a given colormap. + Inialize the level display with a given colormap. The colormap is a list of tuples containing the depth value, and the color for this level : - (depth_value, [RedVal, GreenVal, BlueVal]) + (depth_value, (RedVal, GreenVal, BlueVal)) The colormap should at least contain a color for the depths 0 and 1. See COLORS for an example colormap. - + + Arguments: + colormap : A list of tuple representing a mapping between a float value and + a triple of Red, Green and Blue values. + """ + self.colormap = colormap + self.points = np.array([c[0] for c in colormap]) + self.r = [c[1][0] for c in colormap] + self.g = [c[1][1] for c in colormap] + self.b = [c[1][2] for c in colormap] + + def __call__(self, *args: Any, **kwds: Any) -> Any: + return self.get_color(*args, **kwds) + + def get_color(self, depth_matrix: np.ndarray, frame: np.ndarray, draw_lines: bool=True) -> np.ndarray: + """ + Given a depth matrix between 0-1, return a color from the colormap + Arguments: depth_matrix : Normalized 1 channel numpy matrix frame: Unused. The frame captured by the camera draw_line: Should line be drawn between levels ? - colormap: Override the default colormap. Return: A 3 channel matrix in BGR frame of the same size as depth_matrix. """ - points = np.array([c[0] for c in colormap]) - r = [c[1][2] for c in colormap] - g = [c[1][1] for c in colormap] - b = [c[1][0] for c in colormap] - - # print(rval) - if draw_lines: line_mask = LevelDisplay.draw_lines(depth_matrix, points) - # print(line_mask.shape) - rval = np.interp(depth_matrix, points, r) - gval = np.interp(depth_matrix, points, g) - bval = np.interp(depth_matrix, points, b) + if draw_lines: line_mask = LevelDisplay.draw_lines(depth_matrix, self.points) + rval = np.interp(depth_matrix, self.points, self.r) + gval = np.interp(depth_matrix, self.points, self.g) + bval = np.interp(depth_matrix, self.points, self.b) res = np.dstack((bval, gval, rval)).astype(np.uint8) if draw_lines: res[line_mask,:] = 0 return res @staticmethod - def draw_lines(depth, points, width=0.025, between_levels=False): + def draw_lines(depth:np.ndarray, points:list, width: float=0.025, between_levels: bool=False) -> np.ndarray: """ Given the depth matrix and a list of points, return a mask for every value near the points. If between_levels is set to True, mask in @@ -73,6 +82,4 @@ class LevelDisplay(): halfwidth = width/2 for line in lines: mask = np.where(np.logical_or(mask, np.logical_and(depth >= line - halfwidth, depth <= line+halfwidth)), True, False) - # rpos_mline = np.interp([rpos_mline]) - # print(mask) return mask \ No newline at end of file