Add detection visualization

Author: spencer@primus

avapi/visualize/tracking.py

@@ -0,0 +1,96 @@
+import numpy as np
+from PIL import ImageDraw
+
+
+def draw_detections(image, detections, show_class=False, show_score=False):
+    """Draw detections on an image
+
+    Parameters
+    ----------
+    image: :class:`PIL.Image`
+        Image on which to draw the detections
+    detections: : set of AVstack type
+        A set of detections from AVstack
+    show_class: bool
+        Whether to draw the class of the object. Default is ``False``
+    show_score: bool
+        Whether to draw the score of the object. Default is ``False``
+
+    Returns
+    -------
+    : :class:`PIL.Image`
+        Image with detections drawn
+    """
+    draw = ImageDraw.Draw(image)
+    for detection in detections:
+        x0, y0, x1, y1 = detection.box.box2d
+        # x0, y0, w, h = np.array(detection.state_vector).reshape(4)
+        # x1, y1 = (x0 + w, y0 + h)
+        draw.rectangle([x0, y0, x1, y1], outline=(0, 255, 0), width=1)
+        class_ = detection.obj_type
+        score = detection.score if detection.score else "N/A"
+        # class_ = detection.metadata['class']['name']
+        # score = round(float(detection.metadata['score']),2)
+        if show_class and show_score:
+            draw.text((x0, y1 + 2), "{}:{}".format(class_, score), fill=(0, 255, 0))
+        elif show_class:
+            draw.text((x0, y1 + 2), "{}".format(class_), fill=(0, 255, 0))
+        elif show_score:
+            draw.text((x0, y1 + 2), "{}".format(score), fill=(0, 255, 0))
+
+    del draw
+    return image
+
+
+def draw_stonesoup_tracks(
+    image, tracks, show_history=True, show_class=True, show_score=True
+):
+    """Draw tracks on an image
+
+    Parameters
+    ----------
+    image: :class:`PIL.Image`
+        Image on which to draw the tracks
+    detections: : set of :class:`~.Tracks`
+        A set of tracks generated by our :class:`~.MultiTargetTracker`
+    show_history: bool
+        Whether to draw the trajectory of the track. Default is ``True``
+    show_class: bool
+        Whether to draw the class of the object. Default is ``True``
+    show_score: bool
+        Whether to draw the score of the object. Default is ``True``
+
+    Returns
+    -------
+    : :class:`PIL.Image`
+        Image with tracks drawn
+
+    """
+    draw = ImageDraw.Draw(image)
+    for track in tracks:
+        bboxes = np.array(
+            [
+                np.array(state.state_vector[[0, 2, 4, 5]]).reshape(4)
+                for state in track.states
+            ]
+        )
+        x0, y0, w, h = bboxes[-1]
+        x1 = x0 + w
+        y1 = y0 + h
+        draw.rectangle([x0, y0, x1, y1], outline=(255, 0, 0), width=2)
+
+        if show_history:
+            pts = [(box[0] + box[2] / 2, box[1] + box[3] / 2) for box in bboxes]
+            draw.line(pts, fill=(255, 0, 0), width=2)
+
+        class_ = track.metadata["class"]["name"]
+        score = round(
+            float(track.metadata["score"]) if track.metadata["score"] else 0.0, 2
+        )
+        if show_class and show_score:
+            draw.text((x0, y1 + 2), "{}:{}".format(class_, score), fill=(255, 0, 0))
+        elif show_class:
+            draw.text((x0, y1 + 2), "{}".format(class_), fill=(255, 0, 0))
+        elif show_score:
+            draw.text((x0, y1 + 2), "{}".format(score), fill=(255, 0, 0))
+    return image