Merge pull request #219 from moooises/draw_score

Added the flag draw_scores to the function draw_boxes.

Author: javiber

demos/motmetrics4norfair/src/motmetrics4norfair.py

@@ -3,7 +3,7 @@
 
 import numpy as np
 
-from norfair import drawing, metrics, Tracker, video
+from norfair import Tracker, drawing, metrics, video
 from norfair.camera_motion import MotionEstimator
 from norfair.filter import FilterPyKalmanFilterFactory
 

norfair/drawing/draw_boxes.py

@@ -26,7 +26,8 @@ def draw_boxes(
     detections: Sequence["Detection"] = None,  # Deprecated
     line_color: Optional[ColorLike] = None,  # Deprecated
     line_width: Optional[int] = None,  # Deprecated
-    label_size: Optional[int] = None,  # Deprecated
+    label_size: Optional[int] = None,  # Deprecated´
+    draw_scores: bool = False,
 ) -> np.ndarray:
     """
     Draw bounding boxes corresponding to Detections or TrackedObjects.
@@ -63,6 +64,9 @@ def draw_boxes(
     draw_labels : bool, optional
         If set to True, the label is added to a title that is drawn on top of the box.
         If an object doesn't have a label this parameter is ignored.
+    draw_scores : bool, optional
+        If set to True, the score is added to a title that is drawn on top of the box.
+        If an object doesn't have a label this parameter is ignored.
     text_size : Optional[float], optional
         Size of the title, the value is used as a multiplier of the base size of the font.
         By default the size is scaled automatically based on the frame size.
@@ -148,7 +152,9 @@ def draw_boxes(
                 thickness=thickness,
             )
 
-        text = _build_text(d, draw_labels=draw_labels, draw_ids=draw_ids)
+        text = _build_text(
+            d, draw_labels=draw_labels, draw_ids=draw_ids, draw_scores=draw_scores
+        )
         if text:
             if text_color is None:
                 obj_text_color = obj_color

norfair/drawing/draw_points.py

@@ -26,6 +26,7 @@ def draw_points(
     hide_dead_points: bool = True,
     detections: Sequence["Detection"] = None,  # deprecated
     label_size: Optional[int] = None,  # deprecated
+    draw_scores: bool = False,
 ) -> np.ndarray:
     """
     Draw the points included in a list of Detections or TrackedObjects.
@@ -61,6 +62,9 @@ def draw_points(
     draw_labels : bool, optional
         If set to True, the label is added to a title that is drawn on top of the box.
         If an object doesn't have a label this parameter is ignored.
+    draw_scores : bool, optional
+        If set to True, the score is added to a title that is drawn on top of the box.
+        If an object doesn't have a label this parameter is ignored.
     text_size : Optional[int], optional
         Size of the title, the value is used as a multiplier of the base size of the font.
         By default the size is scaled automatically based on the frame size.
@@ -152,10 +156,12 @@ def draw_points(
                         thickness=thickness,
                     )
 
-        if draw_labels or draw_ids:
+        if draw_labels or draw_ids or draw_scores:
             position = d.points[d.live_points].mean(axis=0)
             position -= radius
-            text = _build_text(d, draw_labels=draw_labels, draw_ids=draw_ids)
+            text = _build_text(
+                d, draw_labels=draw_labels, draw_ids=draw_ids, draw_scores=draw_scores
+            )
 
             Drawer.text(
                 frame,
@@ -165,6 +171,7 @@ def draw_points(
                 color=obj_text_color,
                 thickness=text_thickness,
             )
+
     return frame
 
 

norfair/drawing/utils.py

@@ -13,12 +13,16 @@ def _centroid(tracked_points: np.ndarray) -> Tuple[int, int]:
     return int(sum_x / num_points), int(sum_y / num_points)
 
 
-def _build_text(drawable: "Drawable", draw_labels, draw_ids):
+def _build_text(drawable: "Drawable", draw_labels, draw_ids, draw_scores):
     text = ""
     if draw_labels and drawable.label is not None:
         text = str(drawable.label)
     if draw_ids and drawable.id is not None:
         if len(text) > 0:
             text += "-"
         text += str(drawable.id)
+    if draw_scores and drawable.scores is not None:
+        if len(text) > 0:
+            text += "-"
+        text += str(np.round(np.mean(drawable.scores), 4))
     return text