diff --git a/Makefile b/Makefile index c465be8..f9efe35 100644 --- a/Makefile +++ b/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 diff --git a/motpy/tracker.py b/motpy/tracker.py index d5eb9fd..788c231 100644 --- a/motpy/tracker.py +++ b/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 diff --git a/tests/test_tracker.py b/tests/test_tracker.py index ca6ac15..0941309 100644 --- a/tests/test_tracker.py +++ b/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.