inspection/lifecycle.py

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