Upload engine/tracker.py with huggingface_hub

engine/tracker.py

@@ -0,0 +1,318 @@
+"""
+ByteTrack-inspired Face Tracker for temporal stability in video.
+
+ByteTrack (Zhang et al., 2022) key insight: use ALL detection boxes
+(high + low confidence) for association, not just high-confidence ones.
+Low-confidence detections are valuable for tracking occluded/blurred faces.
+
+Flow:
+1. High-confidence detections → match to existing tracks (IoU + Kalman)
+2. Unmatched tracks + low-confidence detections → second matching round
+3. Remaining unmatched high-confidence → initialize new tracks
+4. Unmatched tracks → mark lost → delete after max_lost frames
+
+Kalman state: [x_center, y_center, aspect_ratio, height, vx, vy, va, vh]
+"""
+
+import numpy as np
+from typing import List, Tuple, Optional, Dict
+from dataclasses import dataclass, field
+
+
+class KalmanBoxTracker:
+    """
+    Kalman filter for bounding box tracking.
+
+    State vector: [cx, cy, s, r, vcx, vcy, vs, vr]
+    where s = area, r = aspect ratio (w/h)
+
+    Measurement: [cx, cy, s, r]
+    """
+
+    _count = 0
+
+    def __init__(self, bbox: np.ndarray):
+        """Initialize tracker with bounding box [x1, y1, x2, y2]."""
+        # State: [cx, cy, s, r, vcx, vcy, vs, vr]
+        self.dim_x = 8
+        self.dim_z = 4
+
+        # State vector
+        self.x = np.zeros(self.dim_x)
+        cx = (bbox[0] + bbox[2]) / 2
+        cy = (bbox[1] + bbox[3]) / 2
+        w = bbox[2] - bbox[0]
+        h = bbox[3] - bbox[1]
+        self.x[0] = cx
+        self.x[1] = cy
+        self.x[2] = w * h  # area
+        self.x[3] = w / max(h, 1e-6)  # aspect ratio
+
+        # State covariance
+        self.P = np.eye(self.dim_x)
+        self.P[4:, 4:] *= 10  # High uncertainty on velocities
+        self.P *= 10
+
+        # Transition matrix (constant velocity)
+        self.F = np.eye(self.dim_x)
+        self.F[0, 4] = 1  # cx += vcx
+        self.F[1, 5] = 1  # cy += vcy
+        self.F[2, 6] = 1  # s += vs
+        self.F[3, 7] = 1  # r += vr
+
+        # Measurement matrix
+        self.H = np.zeros((self.dim_z, self.dim_x))
+        self.H[:4, :4] = np.eye(4)
+
+        # Process noise
+        self.Q = np.eye(self.dim_x) * 0.01
+        self.Q[4:, 4:] *= 0.01
+
+        # Measurement noise
+        self.R = np.eye(self.dim_z) * 1.0
+
+        KalmanBoxTracker._count += 1
+        self.id = KalmanBoxTracker._count
+        self.age = 0
+        self.hits = 0
+        self.time_since_update = 0
+
+    def predict(self) -> np.ndarray:
+        """Predict next state. Returns predicted bbox [x1, y1, x2, y2]."""
+        # Prevent negative area
+        if self.x[2] + self.x[6] <= 0:
+            self.x[6] = 0
+
+        # Kalman predict
+        self.x = self.F @ self.x
+        self.P = self.F @ self.P @ self.F.T + self.Q
+        self.age += 1
+        self.time_since_update += 1
+
+        return self._state_to_bbox()
+
+    def update(self, bbox: np.ndarray):
+        """Update state with measurement [x1, y1, x2, y2]."""
+        cx = (bbox[0] + bbox[2]) / 2
+        cy = (bbox[1] + bbox[3]) / 2
+        w = bbox[2] - bbox[0]
+        h = bbox[3] - bbox[1]
+        z = np.array([cx, cy, w * h, w / max(h, 1e-6)])
+
+        # Kalman update
+        y = z - self.H @ self.x
+        S = self.H @ self.P @ self.H.T + self.R
+        K = self.P @ self.H.T @ np.linalg.inv(S)
+        self.x = self.x + K @ y
+        self.P = (np.eye(self.dim_x) - K @ self.H) @ self.P
+
+        self.hits += 1
+        self.time_since_update = 0
+
+    def _state_to_bbox(self) -> np.ndarray:
+        """Convert state [cx, cy, s, r] to bbox [x1, y1, x2, y2]."""
+        cx, cy, s, r = self.x[:4]
+        s = max(s, 1)
+        w = np.sqrt(s * r)
+        h = s / max(w, 1e-6)
+        return np.array([cx - w/2, cy - h/2, cx + w/2, cy + h/2])
+
+    def get_state(self) -> np.ndarray:
+        """Get current bbox estimate."""
+        return self._state_to_bbox()
+
+
+@dataclass
+class Track:
+    """Single face track."""
+    track_id: int
+    bbox: np.ndarray        # Current bounding box [x1, y1, x2, y2]
+    score: float             # Detection confidence
+    age: int = 0             # Frames since track creation
+    hits: int = 0            # Total detection associations
+    time_since_update: int = 0
+    is_confirmed: bool = False
+    landmarks: Optional[np.ndarray] = None
+
+
+class ByteTracker:
+    """
+    ByteTrack face tracker for video temporal stability.
+
+    Features:
+    - Two-stage association (high + low confidence)
+    - Kalman filter prediction for smooth trajectories
+    - Track lifecycle management (init, confirm, lose, delete)
+    - IoU-based association (no appearance features needed for faces)
+
+    Args:
+        high_thresh: High detection confidence threshold (default: 0.5)
+        low_thresh: Low detection confidence threshold (default: 0.1)
+        match_thresh: IoU threshold for association (default: 0.3)
+        max_lost: Frames before deleting lost tracks (default: 30)
+        min_hits: Detections needed to confirm a track (default: 3)
+    """
+
+    def __init__(self,
+                 high_thresh: float = 0.5,
+                 low_thresh: float = 0.1,
+                 match_thresh: float = 0.3,
+                 max_lost: int = 30,
+                 min_hits: int = 3):
+        self.high_thresh = high_thresh
+        self.low_thresh = low_thresh
+        self.match_thresh = match_thresh
+        self.max_lost = max_lost
+        self.min_hits = min_hits
+
+        self.tracks: List[KalmanBoxTracker] = []
+        self.track_scores: Dict[int, float] = {}
+        self.frame_count = 0
+
+    def update(self, detections: np.ndarray, scores: np.ndarray,
+               landmarks: Optional[np.ndarray] = None) -> List[Track]:
+        """
+        Update tracker with new detections.
+
+        Args:
+            detections: [N, 4] bounding boxes (x1, y1, x2, y2)
+            scores: [N] confidence scores
+            landmarks: [N, 10] optional landmarks
+
+        Returns:
+            List of active Track objects with stable IDs
+        """
+        self.frame_count += 1
+
+        # Split into high and low confidence
+        high_mask = scores >= self.high_thresh
+        low_mask = (scores >= self.low_thresh) & (~high_mask)
+
+        high_dets = detections[high_mask]
+        high_scores = scores[high_mask]
+        low_dets = detections[low_mask]
+        low_scores = scores[low_mask]
+
+        high_lmk = landmarks[high_mask] if landmarks is not None else None
+        low_lmk = landmarks[low_mask] if landmarks is not None else None
+
+        # Predict existing tracks
+        predicted_boxes = []
+        for t in self.tracks:
+            pred = t.predict()
+            predicted_boxes.append(pred)
+        predicted_boxes = np.array(predicted_boxes) if predicted_boxes else np.empty((0, 4))
+
+        # === First association: high-confidence detections ===
+        if len(self.tracks) > 0 and len(high_dets) > 0:
+            iou_matrix = self._iou_batch(predicted_boxes, high_dets)
+            matches_h, unmatched_tracks_h, unmatched_dets_h = \
+                self._hungarian_match(iou_matrix, self.match_thresh)
+        else:
+            matches_h = np.empty((0, 2), dtype=int)
+            unmatched_tracks_h = list(range(len(self.tracks)))
+            unmatched_dets_h = list(range(len(high_dets)))
+
+        # Update matched tracks
+        for t_idx, d_idx in matches_h:
+            self.tracks[t_idx].update(high_dets[d_idx])
+            self.track_scores[self.tracks[t_idx].id] = high_scores[d_idx]
+
+        # === Second association: low-confidence detections with remaining tracks ===
+        remaining_tracks = [self.tracks[i] for i in unmatched_tracks_h]
+        if len(remaining_tracks) > 0 and len(low_dets) > 0:
+            remaining_preds = np.array([t.get_state() for t in remaining_tracks])
+            iou_matrix_l = self._iou_batch(remaining_preds, low_dets)
+            matches_l, unmatched_tracks_l, _ = \
+                self._hungarian_match(iou_matrix_l, self.match_thresh)
+
+            for t_local, d_idx in matches_l:
+                remaining_tracks[t_local].update(low_dets[d_idx])
+                self.track_scores[remaining_tracks[t_local].id] = low_scores[d_idx]
+        else:
+            unmatched_tracks_l = list(range(len(remaining_tracks)))
+
+        # === Initialize new tracks from unmatched high-confidence detections ===
+        for d_idx in unmatched_dets_h:
+            new_tracker = KalmanBoxTracker(high_dets[d_idx])
+            self.tracks.append(new_tracker)
+            self.track_scores[new_tracker.id] = high_scores[d_idx]
+
+        # === Remove lost tracks ===
+        active_tracks = []
+        for t in self.tracks:
+            if t.time_since_update <= self.max_lost:
+                active_tracks.append(t)
+        self.tracks = active_tracks
+
+        # === Build output ===
+        results = []
+        for t in self.tracks:
+            if t.hits >= self.min_hits or self.frame_count <= self.min_hits:
+                bbox = t.get_state()
+                score = self.track_scores.get(t.id, 0.5)
+                track = Track(
+                    track_id=t.id,
+                    bbox=bbox,
+                    score=score,
+                    age=t.age,
+                    hits=t.hits,
+                    time_since_update=t.time_since_update,
+                    is_confirmed=(t.hits >= self.min_hits),
+                )
+                results.append(track)
+
+        return results
+
+    @staticmethod
+    def _iou_batch(boxes1: np.ndarray, boxes2: np.ndarray) -> np.ndarray:
+        """Compute IoU matrix between two sets of boxes."""
+        x1 = np.maximum(boxes1[:, 0:1], boxes2[:, 0:1].T)
+        y1 = np.maximum(boxes1[:, 1:2], boxes2[:, 1:2].T)
+        x2 = np.minimum(boxes1[:, 2:3], boxes2[:, 2:3].T)
+        y2 = np.minimum(boxes1[:, 3:4], boxes2[:, 3:4].T)
+
+        inter = np.maximum(0, x2 - x1) * np.maximum(0, y2 - y1)
+
+        area1 = (boxes1[:, 2] - boxes1[:, 0]) * (boxes1[:, 3] - boxes1[:, 1])
+        area2 = (boxes2[:, 2] - boxes2[:, 0]) * (boxes2[:, 3] - boxes2[:, 1])
+
+        union = area1[:, None] + area2[None, :] - inter
+        return inter / (union + 1e-6)
+
+    @staticmethod
+    def _hungarian_match(iou_matrix: np.ndarray, threshold: float):
+        """Greedy matching by IoU (fast approximation of Hungarian algorithm)."""
+        matches = []
+        unmatched_rows = list(range(iou_matrix.shape[0]))
+        unmatched_cols = list(range(iou_matrix.shape[1]))
+
+        if iou_matrix.size == 0:
+            return np.empty((0, 2), dtype=int), unmatched_rows, unmatched_cols
+
+        # Greedy: take highest IoU pairs iteratively
+        while True:
+            if iou_matrix.size == 0:
+                break
+            max_idx = np.unravel_index(iou_matrix.argmax(), iou_matrix.shape)
+            if iou_matrix[max_idx] < threshold:
+                break
+
+            row, col = max_idx
+            matches.append([unmatched_rows[row], unmatched_cols[col]])
+
+            # Remove matched row and col
+            iou_matrix = np.delete(iou_matrix, row, axis=0)
+            iou_matrix = np.delete(iou_matrix, col, axis=1)
+            unmatched_rows.pop(row)
+            unmatched_cols.pop(col)
+
+        return (np.array(matches) if matches else np.empty((0, 2), dtype=int),
+                unmatched_rows, unmatched_cols)
+
+    def reset(self):
+        """Reset tracker state."""
+        self.tracks.clear()
+        self.track_scores.clear()
+        self.frame_count = 0
+        KalmanBoxTracker._count = 0