Spaces:
Configuration error
Configuration error
File size: 2,418 Bytes
98b1db7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 | # 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
|