Multi-class tracking fix (#247)

* Use binary search to split class bboxes

* Remove fastmath

* Fix pep8

* Use default argument for bisect

* Revert find_split_indices and fix imports

fastmot/mot.py

@@ -2,14 +2,15 @@
 from enum import Enum
 import logging
 import numpy as np
+import numba as nb
 import cv2
 
 from .detector import SSDDetector, YOLODetector, PublicDetector
 from .feature_extractor import FeatureExtractor
 from .tracker import MultiTracker
 from .utils import Profiler
 from .utils.visualization import Visualizer
-from .utils.numba import find_split_indices
+from .utils.numba import bisect_right
 
 
 LOGGER = logging.getLogger(__name__)
@@ -143,7 +144,8 @@ def step(self, frame):
                 detections = self.detector.postprocess()
 
             with Profiler('extract'):
-                cls_bboxes = np.split(detections.tlbr, find_split_indices(detections.label))
+                cls_bboxes = self._split_bboxes_by_cls(detections.tlbr, detections.label,
+                                                       self.class_ids)
                 for extractor, bboxes in zip(self.extractors, cls_bboxes):
                     extractor.extract_async(frame, bboxes)
 
@@ -175,6 +177,17 @@ def print_timing_info():
                      f"{Profiler.get_avg_millis('extract'):>6.3f} ms")
         LOGGER.debug(f"{'association time:':<37}{Profiler.get_avg_millis('assoc'):>6.3f} ms")
 
+    @staticmethod
+    @nb.njit(cache=True)
+    def _split_bboxes_by_cls(bboxes, labels, class_ids):
+        cls_bboxes = []
+        begin = 0
+        for cls_id in class_ids:
+            end = bisect_right(labels, cls_id, begin)
+            cls_bboxes.append(bboxes[begin:end])
+            begin = end
+        return cls_bboxes
+
     def _draw(self, frame, detections):
         visible_tracks = list(self.visible_tracks())
         self.visualizer.render(frame, visible_tracks, detections, self.tracker.klt_bboxes.values(),

fastmot/utils/numba.py

@@ -39,6 +39,19 @@ def mask_area(mask):
     return count
 
 
+@nb.njit(fastmath=True, cache=True)
+def bisect_right(arr, val, left=0):
+    """Utility to search a value in a sorted array."""
+    right = len(arr)
+    while left < right:
+        mid = left + (right - left) // 2
+        if arr[mid] >= val:
+            left = mid + 1
+        else:
+            right = mid
+    return left
+
+
 @nb.njit(fastmath=True, cache=True)
 def find_split_indices(arr):
     """Utility to find indices of unique elements in sorted array."""