Merge pull request #229 from tryolabs/drawing-custom-objects

Enable users to draw custom objects

Author: DiegoFernandezC

norfair/drawing/__init__.py

@@ -3,5 +3,6 @@
 from .color import Color, ColorType, Palette
 from .draw_boxes import draw_boxes, draw_tracked_boxes
 from .draw_points import draw_points, draw_tracked_objects
+from .drawer import Drawable
 from .fixed_camera import FixedCamera
 from .path import AbsolutePaths, Paths

norfair/drawing/draw_boxes.py

@@ -132,7 +132,10 @@ def draw_boxes(
         text_color = parse_color(text_color)
 
     for obj in drawables:
-        d = Drawable(obj)
+        if not isinstance(obj, Drawable):
+            d = Drawable(obj)
+        else:
+            d = obj
 
         if color == "by_id":
             obj_color = Palette.choose_color(d.id)

norfair/drawing/draw_points.py

@@ -126,7 +126,10 @@ def draw_points(
         radius = int(round(max(max(frame.shape) * 0.002, 1)))
 
     for o in drawables:
-        d = Drawable(o)
+        if not isinstance(o, Drawable):
+            d = Drawable(o)
+        else:
+            d = o
 
         if hide_dead_points and not d.live_points.any():
             continue

norfair/drawing/drawer.py

@@ -1,4 +1,4 @@
-from typing import Optional, Sequence, Tuple, Union
+from typing import Any, Optional, Sequence, Tuple, Union
 
 import numpy as np
 
@@ -304,16 +304,36 @@ class Drawable:
 
     Parameters
     ----------
-    obj : Union[Detection, TrackedObject]
+    obj : Union[Detection, TrackedObject], optional
         A [Detection][norfair.tracker.Detection] or a [TrackedObject][norfair.tracker.TrackedObject]
+        that will be used to initialized the drawable.
+        If this parameter is passed, all other arguments are ignored
+    points : np.ndarray, optional
+        Points included in the drawable, shape is `(N_points, N_dimensions)`. Ignored if `obj` is passed
+    id : Any, optional
+        Id of this object. Ignored if `obj` is passed
+    label : Any, optional
+        Label specifying the class of the object. Ignored if `obj` is passed
+    scores : np.ndarray, optional
+        Confidence scores of each point, shape is `(N_points,)`. Ignored if `obj` is passed
+    live_points : np.ndarray, optional
+        Bolean array indicating which points are alive, shape is `(N_points,)`. Ignored if `obj` is passed
 
     Raises
     ------
     ValueError
         If obj is not an instance of the supported classes.
     """
 
-    def __init__(self, obj: Union[Detection, TrackedObject]) -> None:
+    def __init__(
+        self,
+        obj: Union[Detection, TrackedObject] = None,
+        points: np.ndarray = None,
+        id: Any = None,
+        label: Any = None,
+        scores: np.ndarray = None,
+        live_points: np.ndarray = None,
+    ) -> None:
         if isinstance(obj, Detection):
             self.points = obj.points
             self.id = None
@@ -331,6 +351,12 @@ def __init__(self, obj: Union[Detection, TrackedObject]) -> None:
             # it could be the scores of the last detection or some kind of moving average
             self.scores = None
             self.live_points = obj.live_points
+        elif obj is None:
+            self.points = points
+            self.id = id
+            self.label = label
+            self.scores = scores
+            self.live_points = live_points
         else:
             raise ValueError(
                 f"Extecting a Detection or a TrackedObject but received {type(obj)}"