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..b22638d 100644 --- a/motpy/tracker.py +++ b/motpy/tracker.py @@ -363,18 +363,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 +421,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 +431,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..4871780 100644 --- a/tests/test_tracker.py +++ b/tests/test_tracker.py @@ -2,6 +2,7 @@ import numpy as np 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, @@ -94,6 +95,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])