feature: added support for pickable trackers

Makefile

@@ -5,7 +5,7 @@ install-develop:
 	python setup.py develop
 
 test:
-	python -m pytest ./tests
+	python -m pytest -vvs ./tests
 
 env-create:
 	conda create --name env_motpy python=3.7 -y

motpy/tracker.py

@@ -41,8 +41,18 @@ def get_kalman_object_tracker(model: Model, x0: Optional[Vector] = None) -> Kalm
 DEFAULT_MODEL_SPEC = ModelPreset.constant_velocity_and_static_box_size_2d.value
 
 
-def exponential_moving_average_fn(gamma: float) -> Callable:
-    def fn(old, new):
+class EMA:
+    """Pickable interface for callable Exponential Moving Average"""
+    def __init__(self,gamma: float):
+        self.gamma = gamma
+
+    def exponential_moving_average_fn(self, *args, **kwargs) -> float:
+        if len(args) > 0:
+            old, new = args
+        else:
+            old = kwargs.get("old", None)
+            new = kwargs.get("new", None)
+
         if new is None:
             return old
 
@@ -55,9 +65,7 @@ def fn(old, new):
         if isinstance(old, Iterable):
             old = np.array(old)
 
-        return gamma * old + (1 - gamma) * new
-
-    return fn
+        return self.gamma * old + (1 - self.gamma) * new   
 
 
 class SingleObjectTracker:
@@ -73,8 +81,8 @@ def __init__(self,
         self.staleness: float = 0.0
         self.max_staleness: float = max_staleness
 
-        self.update_score_fn: Callable = exponential_moving_average_fn(smooth_score_gamma)
-        self.update_feature_fn: Callable = exponential_moving_average_fn(smooth_feature_gamma)
+        self.update_score_fn: Callable = EMA(smooth_score_gamma).exponential_moving_average_fn
+        self.update_feature_fn: Callable = EMA(smooth_feature_gamma).exponential_moving_average_fn
 
         self.score: Optional[float] = score0
         self.feature: Optional[Vector] = None
@@ -188,7 +196,7 @@ def __init__(self,
         super(SimpleTracker, self).__init__(**kwargs)
         self._box: Box = box0
 
-        self.update_box_fn: Callable = exponential_moving_average_fn(box_update_gamma)
+        self.update_box_fn: Callable = EMA(box_update_gamma).exponential_moving_average_fn
 
     def _predict(self) -> None:
         pass
@@ -363,18 +371,32 @@ def __init__(self, dt: float,
     def active_tracks(self,
                       max_staleness_to_positive_ratio: float = 3.0,
                       max_staleness: float = 999,
-                      min_steps_alive: int = -1) -> List[Track]:
-        """ returns all active tracks after optional filtering by tracker steps count and staleness """
+                      min_steps_alive: int = -1,
+                      return_indices: bool = False) -> Union[List[Track], Tuple[List[Track], List[int]]]:
+        """ returns all active tracks after optional filtering by tracker steps count and staleness 
+            returns indices array from last step, -1 is a body not in passed detected boxes
+        """
 
         tracks: List[Track] = []
+        if return_indices:
+            tracks_indices: List[int] = []
+
         for tracker in self.trackers:
             cond1 = tracker.staleness / tracker.steps_positive < max_staleness_to_positive_ratio  # early stage
             cond2 = tracker.staleness < max_staleness
             cond3 = tracker.steps_alive >= min_steps_alive
             if cond1 and cond2 and cond3:
                 tracks.append(Track(id=tracker.id, box=tracker.box(), score=tracker.score, class_id=tracker.class_id))
+                if return_indices:
+                    try:
+                        tracks_indices.append(self.detections_matched_ids.index(tracker))
+                    except ValueError:
+                        tracks_indices.append(-1)
 
         logger.debug('active/all tracks: %d/%d' % (len(self.trackers), len(tracks)))
+        if return_indices:
+            return tracks, tracks_indices
+
         return tracks
 
     def cleanup_trackers(self) -> None:
@@ -407,7 +429,7 @@ def step(self, detections: Sequence[Detection]) -> List[Track]:
         for match in matches:
             track_idx, det_idx = match[0], match[1]
             self.trackers[track_idx].update(detection=detections[det_idx])
-            self.detections_matched_ids[det_idx] = self.trackers[track_idx].id
+            self.detections_matched_ids[det_idx] = self.trackers[track_idx]
 
         # not assigned detections: create new trackers POF
         assigned_det_idxs = set(matches[:, 1]) if len(matches) > 0 else []
@@ -417,7 +439,7 @@ def step(self, detections: Sequence[Detection]) -> List[Track]:
                                         score0=det.score,
                                         class_id0=det.class_id,
                                         **self.tracker_kwargs)
-            self.detections_matched_ids[det_idx] = tracker.id
+            self.detections_matched_ids[det_idx] = tracker
             self.trackers.append(tracker)
 
         # unassigned trackers

tests/test_tracker.py

@@ -1,11 +1,13 @@
 from collections import Counter
 
 import numpy as np
+import pickle
 import pytest
+from motpy import ModelPreset
 from motpy.core import Detection, setup_logger
 from motpy.testing import data_generator
 from motpy.tracker import (IOUAndFeatureMatchingFunction, MultiObjectTracker,
-                           exponential_moving_average_fn, match_by_cost_matrix)
+                           EMA, match_by_cost_matrix)
 from numpy.testing import assert_almost_equal, assert_array_equal
 
 logger = setup_logger(__name__)
@@ -94,6 +96,23 @@ def test_tracker_diverges():
     assert len(mot.trackers) == 1
     assert mot.active_tracks()[0].id != first_track_id
 
+def test_tracker_det_indices():
+    mot = MultiObjectTracker(dt=5)
+    box0 = np.array([0, 0, 10, 10])
+    box1 = np.array([20, 20, 30, 30])
+    mot.step([Detection(box=box) for box in [box0, box1]])
+    track_ids = [t.id for t in mot.active_tracks()]
+    assert len(track_ids) == 2
+    _, indices = mot.active_tracks(return_indices=True)
+    assert indices == [0, 1]
+    mot.step([Detection(box=box) for box in [box1, box0]])
+    assert track_ids == [t.id for t in mot.active_tracks()]
+    track_ids_idx, indices = mot.active_tracks(return_indices=True)
+    assert track_ids == [t.id for t in track_ids_idx]
+    assert indices == [1, 0]
+    mot.step([Detection(box=box) for box in []])
+    _, indices = mot.active_tracks(return_indices=True)
+    assert indices == [-1, -1]
 
 def test_class_smoothing():
     box = np.array([0, 0, 10, 10])
@@ -106,9 +125,26 @@ def test_class_smoothing():
     mot.step([Detection(box=box, class_id=1)])
     assert mot.trackers[0].class_id == 1
 
+def test_pickable():
+    box = np.array([0, 0, 10, 10])
+    mot = MultiObjectTracker(dt=0.1)
+    dumped_empty = pickle.dumps(mot)
+    assert len(dumped_empty) 
+    mot.step([Detection(box=box, class_id=1)])
+    tracks = mot.active_tracks()
+    dumped_nonempty = pickle.dumps(mot)
+    assert len(dumped_nonempty)
+    mot2 = pickle.loads(dumped_nonempty)
+    tracks2 = mot2.active_tracks()
+    assert len(tracks) == len(tracks2)
+    assert tracks[0].id == tracks2[0].id
+    assert np.array_equal(tracks[0].box, tracks2[0].box)
+    assert tracks[0].score == tracks2[0].score
+    assert tracks[0].class_id == tracks2[0].class_id
+
 
 def test_exponential_moving_average():
-    update_fn = exponential_moving_average_fn(0.5)
+    update_fn = EMA(0.5).exponential_moving_average_fn
 
     # scalars
     assert update_fn(None, 100.) == 100.