File size: 3,074 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 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 | from typing import Dict, List, Optional
class DefectLifecycleManager:
"""
Tracks lifecycle of defects across frames.
A defect is:
- CREATED when first detected
- UPDATED when seen again
- FINALIZED when it disappears for a threshold
"""
def __init__(self, max_missing_frames: int = 5, match_threshold: int = 50):
self.active_defects: Dict[int, Dict] = {}
self.next_id: int = 0
self.max_missing_frames = max_missing_frames
self.match_threshold = match_threshold
def _is_match(self, bbox1, bbox2) -> bool:
"""Check if two bounding boxes belong to same defect."""
x1, y1, _, _ = bbox1
x2, y2, _, _ = bbox2
return (
abs(x1 - x2) < self.match_threshold and
abs(y1 - y2) < self.match_threshold
)
def _find_match(self, detection: Dict) -> Optional[int]:
"""Find matching defect ID for a detection."""
for defect_id, data in self.active_defects.items():
if self._is_match(detection["bbox"], data["bbox"]):
return defect_id
return None
def update(self, detections: List[Dict], frame_id: int) -> List[Dict]:
"""
Update lifecycle with new detections.
Args:
detections: List of processed defects
frame_id: Current frame number
Returns:
List of finalized defects
"""
finalized: List[Dict] = []
# Step 1: Mark all as unseen
for defect in self.active_defects.values():
defect["seen"] = False
# Step 2: Match detections
for d in detections:
matched_id = self._find_match(d)
if matched_id is not None:
data = self.active_defects[matched_id]
data["bbox"] = d["bbox"]
data["last_seen"] = frame_id
data["frames_seen"] += 1
data["seen"] = True
# Update max severity
if d["severity_score"] > data["max_severity"]:
data["max_severity"] = d["severity_score"]
data["final_defect"] = d
else:
# Create new defect
self.active_defects[self.next_id] = {
"bbox": d["bbox"],
"last_seen": frame_id,
"frames_seen": 1,
"max_severity": d.get("severity_score", 0),
"final_defect": d,
"seen": True
}
self.next_id += 1
# Step 3: Finalize disappeared defects
to_delete = []
for defect_id, data in self.active_defects.items():
if not data["seen"]:
if frame_id - data["last_seen"] > self.max_missing_frames:
finalized.append(data["final_defect"])
to_delete.append(defect_id)
# Step 4: Cleanup
for defect_id in to_delete:
del self.active_defects[defect_id]
return finalized |