tracker.py

"""
Simple centroid-based multi-object tracker for video person/car counting.
Solves the problem of counting the same person multiple times across frames.
"""
import numpy as np


class CentroidTracker:
    """Track objects across frames by nearest-centroid matching."""

    def __init__(self, max_disappeared: int = 3, min_appeared: int = 2,
                 distance_threshold: float = 120):
        """
        Args:
            max_disappeared: frames a track can be missing before removal
            min_appeared: minimum frames a track must appear to be "confirmed"
            distance_threshold: max pixel distance for matching centroids
        """
        self.next_id = 0
        self.tracks = {}          # id -> track_info dict
        self.max_disappeared = max_disappeared
        self.min_appeared = min_appeared
        self.distance_threshold = distance_threshold

    def update(self, detections: list, frame_idx: int) -> dict:
        """
        Update tracker with new detections.
        Args:
            detections: [{bbox: (x,y,w,h), class_name, confidence}, ...]
            frame_idx: current frame number
        Returns:
            {frame_idx, active_tracks: [{track_id, bbox, class_name, confidence}, ...],
             unique_counts: {person: N, car: M}}
        """
        if len(detections) == 0:
            # Increment disappeared for all tracks
            for tid in self.tracks:
                self.tracks[tid]["disappeared"] += 1
            self._remove_stale()
            return self._result(frame_idx)

        # Compute centroids from detections
        input_centroids = []
        for d in detections:
            x, y, w, h = d["bbox"]
            cx = x + w / 2
            cy = y + h / 2
            input_centroids.append((cx, cy))

        if len(self.tracks) == 0:
            for i, d in enumerate(detections):
                self.tracks[self.next_id] = {
                    "centroid": input_centroids[i],
                    "class_name": d["class_name"],
                    "confidence": d.get("confidence", 0),
                    "disappeared": 0,
                    "appeared": 1,
                    "bbox": d["bbox"],
                    "confidences": [d.get("confidence", 0)],
                }
                self.next_id += 1
            return self._result(frame_idx)

        # Match existing tracks to new detections
        track_ids = list(self.tracks.keys())
        track_centroids = [self.tracks[tid]["centroid"] for tid in track_ids]

        # Distance matrix: [n_tracks, n_detections]
        D = np.zeros((len(track_ids), len(input_centroids)))
        for i, tc in enumerate(track_centroids):
            for j, ic in enumerate(input_centroids):
                D[i, j] = np.sqrt((tc[0] - ic[0]) ** 2 + (tc[1] - ic[1]) ** 2)

        used_tracks = set()
        used_dets = set()
        matches = []

        # Greedy matching: iterate rows sorted by min distance
        for i in D.min(axis=1).argsort():
            j = D[i].argmin()
            if D[i, j] > self.distance_threshold:
                continue
            if i in used_tracks or j in used_dets:
                continue
            matches.append((i, j))
            used_tracks.add(i)
            used_dets.add(j)

        # Update matched tracks
        for i, j in matches:
            tid = track_ids[i]
            self.tracks[tid]["centroid"] = input_centroids[j]
            self.tracks[tid]["disappeared"] = 0
            self.tracks[tid]["appeared"] += 1
            self.tracks[tid]["bbox"] = detections[j]["bbox"]
            self.tracks[tid]["confidence"] = max(
                self.tracks[tid]["confidence"], detections[j].get("confidence", 0)
            )
            self.tracks[tid]["class_name"] = detections[j]["class_name"]
            self.tracks[tid]["confidences"].append(detections[j].get("confidence", 0))

        # Mark unmatched tracks as disappeared
        for i in range(len(track_ids)):
            if i not in used_tracks:
                self.tracks[track_ids[i]]["disappeared"] += 1

        # Create new tracks for unmatched detections
        for j in range(len(detections)):
            if j not in used_dets:
                self.tracks[self.next_id] = {
                    "centroid": input_centroids[j],
                    "class_name": detections[j]["class_name"],
                    "confidence": detections[j].get("confidence", 0),
                    "disappeared": 0,
                    "appeared": 1,
                    "bbox": detections[j]["bbox"],
                    "confidences": [detections[j].get("confidence", 0)],
                }
                self.next_id += 1

        self._remove_stale()
        return self._result(frame_idx)

    def _remove_stale(self):
        to_remove = [tid for tid, t in self.tracks.items()
                     if t["disappeared"] > self.max_disappeared]
        for tid in to_remove:
            del self.tracks[tid]

    def _result(self, frame_idx: int) -> dict:
        active = []
        unique_counts = {"person": 0, "car": 0}
        seen = {"person": set(), "car": set()}

        for tid, t in self.tracks.items():
            if t["disappeared"] == 0:
                active.append({
                    "track_id": tid,
                    "bbox": t["bbox"],
                    "class_name": t["class_name"],
                    "confidence": round(t["confidence"], 3),
                })

        # Count confirmed tracks as unique individuals
        for tid, t in self.tracks.items():
            if t["appeared"] >= self.min_appeared:
                cls = t.get("class_name", "person")
                unique_counts[cls] = unique_counts.get(cls, 0) + 1

        return {
            "frame_idx": frame_idx,
            "active_tracks": active,
            "unique_counts": unique_counts,
        }

    def final_summary(self) -> dict:
        """Get final tracking summary with unique counts by class."""
        summary = {
            "person": {"count": 0, "tracks": [], "avg_confidence": 0},
            "car": {"count": 0, "tracks": [], "avg_confidence": 0},
        }
        for tid, t in self.tracks.items():
            if t["appeared"] >= self.min_appeared:
                cls = t.get("class_name", "person")
                if cls not in summary:
                    summary[cls] = {"count": 0, "tracks": [], "avg_confidence": 0}
                summary[cls]["count"] += 1
                summary[cls]["tracks"].append({
                    "track_id": tid,
                    "class_name": cls,
                    "confidence": round(t["confidence"], 3),
                    "frames_seen": t["appeared"],
                    "bbox": t["bbox"],
                })

        for cls in summary:
            tracks = summary[cls]["tracks"]
            if tracks:
                summary[cls]["avg_confidence"] = round(
                    sum(tr["confidence"] for tr in tracks) / len(tracks), 3
                )

        return summary