Add drawing functions

Author: GeekAlexis

fastmot/mot.py

@@ -6,7 +6,8 @@
 from .detector import SSDDetector, YoloDetector, PublicDetector
 from .feature_extractor import FeatureExtractor
 from .tracker import MultiTracker
-from .utils.visualization import draw_track, draw_detection, draw_bg_flow
+from .utils.visualization import draw_tracks, draw_detections
+from .utils.visualization import draw_flow_bboxes, draw_background_flow
 
 
 LOGGER = logging.getLogger(__name__)
@@ -116,11 +117,10 @@ def step(self, frame):
         self.frame_count += 1
 
     def _draw(self, frame, detections):
-        for track in self.visible_tracks:
-            draw_track(frame, track, draw_flow=self.verbose)
+        draw_tracks(frame, self.visible_tracks, draw_flow=self.verbose)
         if self.verbose:
-            for det in detections:
-                draw_detection(frame, det)
-            draw_bg_flow(frame, self.tracker)
+            draw_detections(frame, detections)
+            draw_flow_bboxes(frame, self.tracker)
+            draw_background_flow(frame, self.tracker)
         cv2.putText(frame, f'visible: {len(self.visible_tracks)}', (30, 30),
                     cv2.FONT_HERSHEY_SIMPLEX, 1, 0, 2, cv2.LINE_AA)

fastmot/tracker.py

@@ -113,11 +113,11 @@ def step_kalman_filter(self):
             mean, cov = self.kf.warp(mean, cov, self.homography)
             mean, cov = self.kf.predict(mean, cov)
             if trk_id in self.flow_bboxes:
-                flow_bbox = self.flow_bboxes[trk_id]
+                flow_tlbr = self.flow_bboxes[trk_id]
                 # give large flow uncertainty for occluded tracks
                 # usually these with high age and low inlier ratio
                 std_multiplier = max(self.age_factor * track.age, 1) / track.inlier_ratio
-                mean, cov = self.kf.update(mean, cov, flow_bbox, MeasType.FLOW, std_multiplier)
+                mean, cov = self.kf.update(mean, cov, flow_tlbr, MeasType.FLOW, std_multiplier)
             next_tlbr = as_rect(mean[:4])
             track.state = (mean, cov)
             track.tlbr = next_tlbr

fastmot/utils/rect.py

@@ -33,12 +33,6 @@ def get_center(tlbr):
     return np.array([(xmin + xmax) / 2, (ymin + ymax) / 2])
 
 
-@nb.njit(cache=True)
-def get_corners(tlbr):
-    xmin, ymin, xmax, ymax = tlbr
-    return np.array([[xmin, ymin], [xmax, ymin], [xmax, ymax], [xmin, ymax]])
-
-
 @nb.njit(cache=True)
 def to_tlwh(tlbr):
     return np.append(tlbr[:2], get_size(tlbr))
@@ -83,9 +77,9 @@ def crop(img, tlbr):
 
 @nb.njit(cache=True)
 def multi_crop(img, tlbrs):
-    _tlbrs = tlbrs.astype(np.int_)
-    return [img[_tlbrs[i][1]:_tlbrs[i][3] + 1, _tlbrs[i][0]:_tlbrs[i][2] + 1]
-            for i in range(len(_tlbrs))]
+    tlbrs_ = tlbrs.astype(np.int_)
+    return [img[tlbrs_[i][1]:tlbrs_[i][3] + 1, tlbrs_[i][0]:tlbrs_[i][2] + 1]
+            for i in range(len(tlbrs_))]
 
 
 @nb.njit(fastmath=True, cache=True)

fastmot/utils/visualization.py

@@ -6,37 +6,43 @@
 GOLDEN_RATIO = 0.618033988749895
 
 
-def draw_track(frame, trk, draw_flow=False):
-    _draw_bbox(frame, trk.tlbr, _get_color(trk.trk_id), 2, str(trk.trk_id))
-    if draw_flow:
-        _draw_feature_match(frame, trk.keypoints, trk.prev_keypoints, (0, 255, 255))
+def draw_tracks(frame, tracks, draw_flow=False):
+    for track in tracks:
+        draw_bbox(frame, track.tlbr, get_color(track.trk_id), 2, str(track.trk_id))
+        if draw_flow:
+            draw_feature_match(frame, track.prev_keypoints, track.keypoints, (0, 255, 255))
 
 
-def draw_detection(frame, det):
-    _draw_bbox(frame, det.tlbr, (255, 255, 255), 1)
+def draw_detections(frame, detections):
+    for det in detections:
+        draw_bbox(frame, det.tlbr, (255, 255, 255), 1)
 
 
-def draw_tile(frame, detector):
+def draw_flow_bboxes(frame, tracker):
+    for tlbr in tracker.flow_bboxes.values():
+        draw_bbox(frame, tlbr, 0, 1)
+
+
+def draw_tiles(frame, detector):
     assert hasattr(detector, 'tiles')
     for tile in detector.tiles:
-        tl = np.rint(tile[:2] * detector.scale_factor).astype(int)
-        br = np.rint(tile[2:] * detector.scale_factor).astype(int)
-        cv2.rectangle(frame, tuple(tl), tuple(br), 0, 1)
+        tlbr = np.rint(tile * np.tile(detector.scale_factor, 2))
+        draw_bbox(frame, tlbr, 0, 1)
 
 
-def draw_bg_flow(frame, tracker):
-    _draw_feature_match(frame, tracker.flow.bg_keypoints,
-                        tracker.flow.prev_bg_keypoints, (0, 0, 255))
+def draw_background_flow(frame, tracker):
+    draw_feature_match(frame, tracker.flow.prev_bg_keypoints,
+                       tracker.flow.bg_keypoints, (0, 0, 255))
 
 
-def _get_color(idx, s=0.8, vmin=0.7):
+def get_color(idx, s=0.8, vmin=0.7):
     h = np.fmod(idx * GOLDEN_RATIO, 1.)
     v = 1. - np.fmod(idx * GOLDEN_RATIO, 1. - vmin)
     r, g, b = colorsys.hsv_to_rgb(h, s, v)
     return int(255 * b), int(255 * g), int(255 * r)
 
 
-def _draw_bbox(frame, tlbr, color, thickness, text=None):
+def draw_bbox(frame, tlbr, color, thickness, text=None):
     tlbr = tlbr.astype(int)
     tl, br = tuple(tlbr[:2]), tuple(tlbr[2:])
     cv2.rectangle(frame, tl, br, color, thickness)
@@ -48,7 +54,7 @@ def _draw_bbox(frame, tlbr, color, thickness, text=None):
                     0.5, 0, 1, cv2.LINE_AA)
 
 
-def _draw_feature_match(frame, cur_pts, prev_pts, color):
+def draw_feature_match(frame, prev_pts, cur_pts, color):
     if len(cur_pts) > 0:
         cur_pts = np.rint(cur_pts).astype(np.int32)
         for pt in cur_pts:
@@ -57,3 +63,22 @@ def _draw_feature_match(frame, cur_pts, prev_pts, color):
             prev_pts = np.rint(prev_pts).astype(np.int32)
             for pt1, pt2 in zip(prev_pts, cur_pts):
                 cv2.line(frame, tuple(pt1), tuple(pt2), color, 1, cv2.LINE_AA)
+
+
+def draw_covariance(frame, tlbr, covariance):
+    tlbr = tlbr.astype(int)
+    tl, br = tuple(tlbr[:2]), tuple(tlbr[2:])
+
+    def ellipse(cov):
+        vals, vecs = np.linalg.eigh(cov)
+        order = vals.argsort()[::-1]
+        # 95% confidence ellipse
+        vals, vecs = np.sqrt(vals[order] * 5.9915), vecs[:, order]
+        axes = int(vals[0] + 0.5), int(vals[1] + 0.5)
+        angle = np.degrees(np.arctan2(vecs[1, 0], vecs[0, 0]))
+        return axes, angle
+
+    axes, angle = ellipse(covariance[:2, :2])
+    cv2.ellipse(frame, tl, axes, angle, 0, 360, (255, 255, 255), 1)
+    axes, angle = ellipse(covariance[2:4, 2:4])
+    cv2.ellipse(frame, br, axes, angle, 0, 360, (255, 255, 255), 1)