inspection/tracker.py

import numpy as np
from typing import List, Dict, Tuple


class DefectTracker:
    """
    Tracks detected defects across frames using centroid distance.

    Attributes:
        max_distance: Maximum distance to consider a match
        confirm_frames: Frames required to confirm a detection
    """

    def __init__(self, max_distance: int = 60, confirm_frames: int = 3):
        self.max_distance = max_distance
        self.confirm_frames = confirm_frames

        self.tracks: List[Dict] = []

    def _compute_center(self, bbox: Tuple[int, int, int, int]) -> Tuple[float, float]:
        """Compute center of bounding box."""
        x, y, w, h = bbox
        return x + w / 2, y + h / 2

    def update(self, detections: List[Dict]) -> List[Dict]:
        """
        Update tracker with new detections.

        Args:
            detections: List of detection dictionaries

        Returns:
            List of confirmed detections
        """

        confirmed: List[Dict] = []

        # Mark all tracks as not updated
        for track in self.tracks:
            track["updated"] = False

        for det in detections:

            center = self._compute_center(det["bbox"])
            matched = False

            for track in self.tracks:

                tx, ty = track["center"]
                dist = np.sqrt((center[0] - tx) ** 2 + (center[1] - ty) ** 2)

                if dist < self.max_distance:

                    # Update track
                    track["center"] = center
                    track["count"] += 1
                    track["data"] = det
                    track["updated"] = True

                    matched = True

                    # Confirm detection after enough frames
                    if track["count"] >= self.confirm_frames:
                        confirmed.append(det)

                    break

            if not matched:
                self.tracks.append({
                    "center": center,
                    "count": 1,
                    "data": det,
                    "updated": True
                })

        # 🧹 Cleanup stale tracks (IMPORTANT for long runs)
        self.tracks = [
            t for t in self.tracks if t["updated"] or t["count"] < self.confirm_frames
        ]

        return confirmed