tracker.py

# Simple centroid tracker (frame-to-frame assignment by nearest centroid with distance gate)
import numpy as np

class SimpleTracker:
    def __init__(self, max_lost=30, dist_thresh=80):
        self.next_id = 1
        self.tracks = {}         # id -> {'bbox':(x1,y1,x2,y2), 'centroid':(x,y), 'lost':0, 'history':[]}
        self.max_lost = max_lost
        self.dist_thresh = dist_thresh

    def _centroid(self, bbox):
        x1,y1,x2,y2 = bbox
        return ((x1+x2)/2.0, (y1+y2)/2.0)

    def update(self, person_bboxes):
        # person_bboxes: list of (x1,y1,x2,y2)
        centroids = [self._centroid(b) for b in person_bboxes]

        # Build cost matrix
        track_ids = list(self.tracks.keys())
        T = len(track_ids); D = len(centroids)
        cost = np.full((T, D), fill_value=1e9, dtype=float)
        for i, tid in enumerate(track_ids):
            tx, ty = self.tracks[tid]['centroid']
            for j, c in enumerate(centroids):
                cost[i,j] = np.linalg.norm(np.array([tx,ty]) - np.array(c))

        # Greedy assignment by nearest
        assigned_tracks = set()
        assigned_dets = set()

        while True:
            i,j = np.unravel_index(cost.argmin(), cost.shape) if cost.size else (None,None)
            if cost.size == 0: break
            if cost[i,j] > self.dist_thresh:
                break
            # assign
            tid = track_ids[i]
            self.tracks[tid]['bbox'] = person_bboxes[j]
            self.tracks[tid]['centroid'] = centroids[j]
            self.tracks[tid]['lost'] = 0
            self.tracks[tid]['history'].append(centroids[j])
            assigned_tracks.add(i); assigned_dets.add(j)
            cost[i,:] = 1e9
            cost[:,j] = 1e9

        # Increase lost for unassigned tracks
        for idx, tid in enumerate(track_ids):
            if idx not in assigned_tracks:
                self.tracks[tid]['lost'] += 1

        # Remove lost tracks
        for tid in list(self.tracks.keys()):
            if self.tracks[tid]['lost'] > self.max_lost:
                del self.tracks[tid]

        # Add new tracks for unassigned detections
        for j, b in enumerate(person_bboxes):
            if j not in assigned_dets:
                c = self._centroid(b)
                self.tracks[self.next_id] = {'bbox': b, 'centroid': c, 'lost':0, 'history':[c]}
                self.next_id += 1

        return self.tracks