Cleanup examples, optimize a little

src/ar_sandbox/examples/examples.py

+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