File size: 2,293 Bytes
201b13c | 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 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 | 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 |