visiongraph.tracker.FlateTracker

  1from argparse import ArgumentParser
  2from collections import defaultdict
  3from dataclasses import dataclass
  4from typing import List, Optional, Dict
  5
  6from visiongraph.result.ResultList import ResultList
  7from visiongraph.result.spatial.ObjectDetectionResult import ObjectDetectionResult
  8from visiongraph.tracker.BaseObjectDetectionTracker import BaseObjectDetectionTracker
  9from visiongraph.tracker.ObjectAssignmentSolver import ObjectAssignmentSolver, CostFunctionType
 10
 11
 12@dataclass
 13class _FlateTrack:
 14    id: int
 15    reference: ObjectDetectionResult
 16    age: int = 0
 17    stale: int = 0
 18
 19    def update_reference(self):
 20        self.reference.tracking_id = self.id
 21
 22
 23class FlateTracker(BaseObjectDetectionTracker):
 24    """
 25    Fast localization and tracking engine.
 26    """
 27
 28    def __init__(
 29        self,
 30        max_cost: float = 0.5,
 31        min_alive: int = 0,
 32        max_lost: int = 5,
 33        class_aware: bool = False,
 34        cost_function: Optional[CostFunctionType] = None,
 35    ):
 36        """
 37        Initializes the FlateTracker with specified parameters.
 38
 39        :param max_cost: Maximum cost for a trackable match.
 40        :param min_alive: Minimum number of frames a track must be visible to be considered alive.
 41        :param max_lost: Maximum number of frames a track can be lost before it is removed.
 42        :param class_aware: If True, run class-aware matching.
 43        :param cost_function: Cost function to use for matching. Defaults to L2 distance.
 44        """
 45        self.max_cost: float = max_cost
 46
 47        self.min_alive: int = min_alive
 48        self.max_lost: int = max_lost
 49
 50        self.include_stale: bool = False
 51
 52        self.class_aware: bool = class_aware
 53        self.cost_function: CostFunctionType = (
 54            cost_function if cost_function is not None else ObjectAssignmentSolver.l2_cost_function
 55        )
 56
 57        self._tracks: List[_FlateTrack] = []
 58        self._unique_id: int = 0
 59
 60    def setup(self):
 61        """
 62        Prepares the tracker for a new tracking session by clearing existing tracks and resetting the unique ID.
 63        """
 64        self._tracks.clear()
 65        self._unique_id = 0
 66
 67    def _new_id(self) -> int:
 68        """
 69        Generates a new unique track ID.
 70
 71        :return: A new unique track ID.
 72        """
 73        track_id = self._unique_id
 74        self._unique_id += 1
 75        return track_id
 76
 77    def process(self, detections: List[ObjectDetectionResult]) -> ResultList[ObjectDetectionResult]:
 78        """
 79        Processes the given detections to update tracks and create new ones if necessary.
 80
 81        :param detections: A list of detected objects to process.
 82
 83        :return: A list of tracked objects.
 84        """
 85        if not self._tracks and not detections:
 86            return ResultList([])
 87
 88        if not self._tracks:
 89            for det in detections:
 90                tr = _FlateTrack(self._new_id(), det)
 91                tr.update_reference()
 92                self._tracks.append(tr)
 93            return ResultList([t.reference for t in self._tracks])
 94
 95        if not detections:
 96            for t in self._tracks:
 97                t.stale += 1
 98                t.reference.staleness = t.stale
 99            self._tracks = [t for t in self._tracks if t.stale <= self.max_lost]
100            return ResultList(
101                [t.reference for t in self._tracks if t.age >= self.min_alive and (self.include_stale or t.stale == 0)]
102            )
103
104        tracks_ref = [t.reference for t in self._tracks]
105
106        solver = ObjectAssignmentSolver(self.cost_function, self.max_cost)
107
108        assignments: Dict[ObjectDetectionResult, Optional[ObjectDetectionResult]] = {}
109        unassigned_destinations: List[ObjectDetectionResult] = []
110
111        if not self.class_aware:
112            result = solver.solve(tracks_ref, detections)
113            assignments = result.assignments
114            unassigned_destinations = result.unassigned_destinations
115        else:
116            # Class-aware matching
117            track_classes = [r.class_id for r in tracks_ref]
118            det_classes = [d.class_id for d in detections]
119
120            # Group by class
121            tracks_by_class = defaultdict(list)
122            for i, c in enumerate(track_classes):
123                if c is not None:
124                    tracks_by_class[c].append(tracks_ref[i])
125
126            dets_by_class = defaultdict(list)
127            for i, c in enumerate(det_classes):
128                if c is not None:
129                    dets_by_class[c].append(detections[i])
130
131            known_classes = sorted(set(tracks_by_class.keys()) | set(dets_by_class.keys()))
132
133            matched_tracks = set()
134            matched_dets = set()
135
136            # Step 1: per-class matching for known classes
137            for cls_id in known_classes:
138                t_subset = tracks_by_class.get(cls_id, [])
139                d_subset = dets_by_class.get(cls_id, [])
140
141                if not t_subset or not d_subset:
142                    continue
143
144                res = solver.solve(t_subset, d_subset)
145
146                for src, dst in res.assignments.items():
147                    assignments[src] = dst
148                    matched_tracks.add(src)
149                    if dst:
150                        matched_dets.add(dst)
151
152            # Step 2: handle unknown-class tracks against any remaining detections
153            unknown_class_tracks = [
154                t for i, t in enumerate(tracks_ref) if track_classes[i] is None and t not in matched_tracks
155            ]
156            rem_dets = [d for d in detections if d not in matched_dets]
157
158            if unknown_class_tracks and rem_dets:
159                res = solver.solve(unknown_class_tracks, rem_dets)
160                for src, dst in res.assignments.items():
161                    assignments[src] = dst
162                    matched_tracks.add(src)
163                    if dst:
164                        matched_dets.add(dst)
165
166            # Fill unassigned destinations
167            unassigned_destinations = [d for d in detections if d not in matched_dets]
168
169            # Ensure all tracks are in assignments
170            for t in tracks_ref:
171                if t not in assignments:
172                    assignments[t] = None
173
174        # Update tracks
175        for track in self._tracks:
176            dest = assignments.get(track.reference)
177
178            if dest is not None:
179                track.age += 1
180                track.stale = 0
181                track.reference = dest
182                track.update_reference()
183            else:
184                track.stale += 1
185
186            track.reference.staleness = track.stale
187
188        # Create new tracks
189        for det in unassigned_destinations:
190            tr = _FlateTrack(self._new_id(), det)
191            tr.update_reference()
192            tr.reference.staleness = 0
193            self._tracks.append(tr)
194
195        # Clean up stale tracks
196        self._tracks = [t for t in self._tracks if t.stale <= self.max_lost]
197
198        return ResultList(
199            [t.reference for t in self._tracks if t.age >= self.min_alive and (self.include_stale or t.stale == 0)]
200        )
201
202    def release(self):
203        """
204        Releases resources and clears all tracks.
205        """
206        self._tracks.clear()
207
208    def configure(self, args):
209        """
210        Configures the tracker with parameters from the provided argument parser.
211        """
212        self.max_cost = self._get_param(args, "tracker_max_cost", self.max_cost)
213        self.min_alive = self._get_param(args, "tracker_min_alive", self.min_alive)
214        self.max_lost = self._get_param(args, "tracker_max_lost", self.max_lost)
215
216    @staticmethod
217    def add_params(parser: ArgumentParser):
218        """
219        Adds command line parameters for configuring the tracker.
220        """
221        parser.add_argument("--tracker-max-cost", type=float, default=0.5, help="Max cost for trackable match.")
222        parser.add_argument("--tracker-min-alive", type=int, default=0, help="Min frames trackable visible.")
223        parser.add_argument("--tracker-max-lost", type=int, default=5, help="Max frames trackable not visible.")