Refactor: PHASE 2: Core Module Extraction

detector_worker/detection/tracking_manager.py

@@ -0,0 +1,481 @@
+"""
+Object tracking management and state handling.
+
+This module provides tracking state management, session handling,
+and occupancy detection functionality for the detection pipeline.
+"""
+
+import time
+import logging
+from typing import Dict, List, Any, Optional, Tuple, Set
+from dataclasses import dataclass, field
+from datetime import datetime
+
+from ..core.constants import SESSION_TIMEOUT_SECONDS
+from ..core.exceptions import TrackingError, create_detection_error
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class SessionState:
+    """Session state for tracking management."""
+    active: bool = True
+    waiting_for_backend_session: bool = False
+    wait_start_time: float = 0.0
+    reset_tracker_on_resume: bool = False
+    occupancy_mode: bool = False
+    occupancy_enabled_at: Optional[float] = None
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary format."""
+        return {
+            "active": self.active,
+            "waiting_for_backend_session": self.waiting_for_backend_session,
+            "wait_start_time": self.wait_start_time,
+            "reset_tracker_on_resume": self.reset_tracker_on_resume,
+            "occupancy_mode": self.occupancy_mode,
+            "occupancy_enabled_at": self.occupancy_enabled_at
+        }
+    
+    def from_dict(self, data: Dict[str, Any]) -> None:
+        """Update from dictionary data."""
+        self.active = data.get("active", True)
+        self.waiting_for_backend_session = data.get("waiting_for_backend_session", False)
+        self.wait_start_time = data.get("wait_start_time", 0.0)
+        self.reset_tracker_on_resume = data.get("reset_tracker_on_resume", False)
+        self.occupancy_mode = data.get("occupancy_mode", False)
+        self.occupancy_enabled_at = data.get("occupancy_enabled_at")
+
+
+@dataclass
+class TrackingData:
+    """Complete tracking data for a camera and model."""
+    track_stability_counters: Dict[int, int] = field(default_factory=dict)
+    stable_tracks: Set[int] = field(default_factory=set)
+    session_state: SessionState = field(default_factory=SessionState)
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary format."""
+        return {
+            "track_stability_counters": dict(self.track_stability_counters),
+            "stable_tracks": list(self.stable_tracks),
+            "session_state": self.session_state.to_dict()
+        }
+
+
+class TrackingManager:
+    """
+    Manages object tracking state and session handling across cameras and models.
+    
+    This class provides centralized tracking state management including:
+    - Track stability counters
+    - Session state management
+    - Backend session timeout handling
+    - Occupancy detection mode
+    """
+    
+    def __init__(self, session_timeout_seconds: int = SESSION_TIMEOUT_SECONDS):
+        """
+        Initialize tracking manager.
+        
+        Args:
+            session_timeout_seconds: Timeout for backend session waiting
+        """
+        self.session_timeout_seconds = session_timeout_seconds
+        self._camera_tracking_data: Dict[str, Dict[str, TrackingData]] = {}
+        self._lock = None
+    
+    def _ensure_thread_safety(self):
+        """Initialize thread safety if not already done."""
+        if self._lock is None:
+            import threading
+            self._lock = threading.RLock()
+    
+    def get_camera_tracking_data(self, camera_id: str, model_id: str) -> TrackingData:
+        """
+        Get or create tracking data for a specific camera and model.
+        
+        Args:
+            camera_id: Unique camera identifier
+            model_id: Model identifier
+            
+        Returns:
+            Tracking data for the camera and model
+        """
+        self._ensure_thread_safety()
+        
+        with self._lock:
+            if camera_id not in self._camera_tracking_data:
+                self._camera_tracking_data[camera_id] = {}
+            
+            if model_id not in self._camera_tracking_data[camera_id]:
+                logger.warning(f"🔄 Camera {camera_id}: Creating NEW tracking data for {model_id} - this will reset any cooldown!")
+                self._camera_tracking_data[camera_id][model_id] = TrackingData()
+            
+            return self._camera_tracking_data[camera_id][model_id]
+    
+    def check_stable_tracks(self, 
+                          camera_id: str, 
+                          model_id: str, 
+                          regions_dict: Dict[str, Any]) -> Tuple[bool, List[Dict[str, Any]]]:
+        """
+        Check if any stable tracks match the detected classes for a specific camera.
+        
+        Args:
+            camera_id: Unique camera identifier
+            model_id: Model identifier
+            regions_dict: Dictionary mapping class names to detection regions
+            
+        Returns:
+            Tuple of (has_stable_tracks, stable_detections)
+        """
+        self._ensure_thread_safety()
+        
+        with self._lock:
+            tracking_data = self.get_camera_tracking_data(camera_id, model_id)
+            stable_tracks = tracking_data.stable_tracks
+            
+            if not stable_tracks:
+                return False, []
+            
+            # Check for track-based stability
+            stable_detections = []
+            
+            for class_name, region_data in regions_dict.items():
+                detection = region_data.get("detection", {})
+                track_id = detection.get("id")
+                
+                if track_id in stable_tracks:
+                    stable_detections.append(detection)
+            
+            has_stable = len(stable_detections) > 0
+            logger.debug(f"Camera {camera_id}: Stable track check - stable_tracks: {list(stable_tracks)}, found: {has_stable}")
+            
+            return has_stable, stable_detections
+    
+    def reset_tracking_state(self, 
+                           camera_id: str, 
+                           model_id: str, 
+                           reason: str = "session ended") -> None:
+        """
+        Reset tracking state after session completion or timeout.
+        
+        Args:
+            camera_id: Unique camera identifier
+            model_id: Model identifier
+            reason: Reason for reset (for logging)
+        """
+        self._ensure_thread_safety()
+        
+        with self._lock:
+            tracking_data = self.get_camera_tracking_data(camera_id, model_id)
+            session_state = tracking_data.session_state
+            
+            # Clear all tracking data for fresh start
+            tracking_data.track_stability_counters.clear()
+            tracking_data.stable_tracks.clear()
+            session_state.active = True
+            session_state.waiting_for_backend_session = False
+            session_state.wait_start_time = 0.0
+            session_state.reset_tracker_on_resume = True
+            
+            logger.info(f"Camera {camera_id}: 🔄 Reset tracking state - {reason}")
+            logger.info(f"Camera {camera_id}: 🧹 Cleared stability counters and stable tracks for fresh session")
+    
+    def is_camera_active(self, camera_id: str, model_id: str) -> bool:
+        """
+        Check if camera should be processing detections.
+        
+        Args:
+            camera_id: Unique camera identifier
+            model_id: Model identifier
+            
+        Returns:
+            True if camera should process detections, False otherwise
+        """
+        self._ensure_thread_safety()
+        
+        with self._lock:
+            tracking_data = self.get_camera_tracking_data(camera_id, model_id)
+            session_state = tracking_data.session_state
+            
+            # Check if waiting for backend sessionId has timed out
+            if session_state.waiting_for_backend_session:
+                current_time = time.time()
+                wait_start_time = session_state.wait_start_time
+                elapsed_time = current_time - wait_start_time
+                
+                if elapsed_time >= self.session_timeout_seconds:
+                    logger.warning(f"Camera {camera_id}: Backend sessionId timeout ({self.session_timeout_seconds}s) - resetting tracking")
+                    self.reset_tracking_state(camera_id, model_id, "backend sessionId timeout")
+                    return True
+                else:
+                    remaining_time = self.session_timeout_seconds - elapsed_time
+                    logger.debug(f"Camera {camera_id}: Still waiting for backend sessionId ({remaining_time:.1f}s remaining)")
+                    return False
+            
+            return session_state.active
+    
+    def set_waiting_for_backend_session(self, 
+                                      camera_id: str, 
+                                      model_id: str, 
+                                      waiting: bool = True) -> None:
+        """
+        Set whether camera is waiting for backend session ID.
+        
+        Args:
+            camera_id: Unique camera identifier
+            model_id: Model identifier
+            waiting: Whether to wait for backend session
+        """
+        self._ensure_thread_safety()
+        
+        with self._lock:
+            tracking_data = self.get_camera_tracking_data(camera_id, model_id)
+            session_state = tracking_data.session_state
+            
+            session_state.waiting_for_backend_session = waiting
+            if waiting:
+                session_state.wait_start_time = time.time()
+                logger.debug(f"Camera {camera_id}: Started waiting for backend sessionId")
+            else:
+                session_state.wait_start_time = 0.0
+                logger.debug(f"Camera {camera_id}: Stopped waiting for backend sessionId")
+    
+    def cleanup_camera_tracking(self, camera_id: str) -> None:
+        """
+        Clean up tracking data when a camera is disconnected.
+        
+        Args:
+            camera_id: Unique camera identifier
+        """
+        self._ensure_thread_safety()
+        
+        with self._lock:
+            if camera_id in self._camera_tracking_data:
+                del self._camera_tracking_data[camera_id]
+                logger.info(f"Cleaned up tracking data for camera {camera_id}")
+    
+    def set_occupancy_mode(self, 
+                         camera_id: str, 
+                         model_id: str, 
+                         enable: bool = True) -> bool:
+        """
+        Enable or disable occupancy detection mode.
+        
+        Occupancy mode stops model inference after pipeline completion
+        while backend session handling continues in background.
+        
+        Args:
+            camera_id: Unique camera identifier
+            model_id: Model identifier
+            enable: True to enable occupancy mode, False to disable
+            
+        Returns:
+            Current occupancy mode state
+        """
+        self._ensure_thread_safety()
+        
+        with self._lock:
+            tracking_data = self.get_camera_tracking_data(camera_id, model_id)
+            session_state = tracking_data.session_state
+            
+            if enable:
+                session_state.occupancy_mode = True
+                session_state.occupancy_enabled_at = time.time()
+                logger.debug(f"Camera {camera_id}: Occupancy mode ENABLED - model will stop after pipeline completion")
+            else:
+                session_state.occupancy_mode = False
+                session_state.occupancy_enabled_at = None
+                logger.debug(f"Camera {camera_id}: Occupancy mode DISABLED - model will continue running")
+            
+            return session_state.occupancy_mode
+    
+    def is_occupancy_mode_enabled(self, camera_id: str, model_id: str) -> bool:
+        """
+        Check if occupancy mode is enabled for a camera.
+        
+        Args:
+            camera_id: Unique camera identifier
+            model_id: Model identifier
+            
+        Returns:
+            True if occupancy mode is enabled
+        """
+        self._ensure_thread_safety()
+        
+        with self._lock:
+            tracking_data = self.get_camera_tracking_data(camera_id, model_id)
+            return tracking_data.session_state.occupancy_mode
+    
+    def get_occupancy_duration(self, camera_id: str, model_id: str) -> Optional[float]:
+        """
+        Get duration since occupancy mode was enabled.
+        
+        Args:
+            camera_id: Unique camera identifier
+            model_id: Model identifier
+            
+        Returns:
+            Duration in seconds since occupancy mode was enabled, or None if not enabled
+        """
+        self._ensure_thread_safety()
+        
+        with self._lock:
+            tracking_data = self.get_camera_tracking_data(camera_id, model_id)
+            session_state = tracking_data.session_state
+            
+            if not session_state.occupancy_mode or not session_state.occupancy_enabled_at:
+                return None
+            
+            return time.time() - session_state.occupancy_enabled_at
+    
+    def get_session_state(self, camera_id: str, model_id: str) -> SessionState:
+        """
+        Get session state for a camera and model.
+        
+        Args:
+            camera_id: Unique camera identifier
+            model_id: Model identifier
+            
+        Returns:
+            Session state object
+        """
+        self._ensure_thread_safety()
+        
+        with self._lock:
+            tracking_data = self.get_camera_tracking_data(camera_id, model_id)
+            return tracking_data.session_state
+    
+    def update_session_state(self, 
+                           camera_id: str, 
+                           model_id: str, 
+                           **kwargs) -> None:
+        """
+        Update session state properties.
+        
+        Args:
+            camera_id: Unique camera identifier
+            model_id: Model identifier
+            **kwargs: Session state properties to update
+        """
+        self._ensure_thread_safety()
+        
+        with self._lock:
+            tracking_data = self.get_camera_tracking_data(camera_id, model_id)
+            session_state = tracking_data.session_state
+            
+            for key, value in kwargs.items():
+                if hasattr(session_state, key):
+                    setattr(session_state, key, value)
+                    logger.debug(f"Camera {camera_id}: Updated session state {key} = {value}")
+                else:
+                    logger.warning(f"Camera {camera_id}: Unknown session state property: {key}")
+    
+    def get_tracking_statistics(self, camera_id: str, model_id: str) -> Dict[str, Any]:
+        """
+        Get comprehensive tracking statistics for a camera and model.
+        
+        Args:
+            camera_id: Unique camera identifier
+            model_id: Model identifier
+            
+        Returns:
+            Dictionary with tracking statistics
+        """
+        self._ensure_thread_safety()
+        
+        with self._lock:
+            tracking_data = self.get_camera_tracking_data(camera_id, model_id)
+            current_time = time.time()
+            
+            stats = tracking_data.to_dict()
+            
+            # Add computed statistics
+            stats["total_tracked_objects"] = len(tracking_data.track_stability_counters)
+            stats["stable_track_count"] = len(tracking_data.stable_tracks)
+            
+            session_state = tracking_data.session_state
+            if session_state.waiting_for_backend_session and session_state.wait_start_time > 0:
+                stats["backend_wait_duration"] = current_time - session_state.wait_start_time
+                stats["backend_wait_remaining"] = max(0, self.session_timeout_seconds - stats["backend_wait_duration"])
+            
+            if session_state.occupancy_mode and session_state.occupancy_enabled_at:
+                stats["occupancy_duration"] = current_time - session_state.occupancy_enabled_at
+            
+            return stats
+    
+    def get_all_camera_stats(self) -> Dict[str, Dict[str, Dict[str, Any]]]:
+        """
+        Get tracking statistics for all monitored cameras.
+        
+        Returns:
+            Nested dictionary: {camera_id: {model_id: stats}}
+        """
+        self._ensure_thread_safety()
+        
+        with self._lock:
+            all_stats = {}
+            for camera_id, models in self._camera_tracking_data.items():
+                all_stats[camera_id] = {}
+                for model_id in models.keys():
+                    all_stats[camera_id][model_id] = self.get_tracking_statistics(camera_id, model_id)
+            
+            return all_stats
+
+
+# Global tracking manager instance
+tracking_manager = TrackingManager()
+
+
+# ===== CONVENIENCE FUNCTIONS =====
+# These provide the same interface as the original functions in pympta.py
+
+def check_stable_tracks(camera_id: str, model_id: str, regions_dict: Dict[str, Any]) -> Tuple[bool, List[Dict[str, Any]]]:
+    """Check if any stable tracks match the detected classes for a specific camera."""
+    return tracking_manager.check_stable_tracks(camera_id, model_id, regions_dict)
+
+
+def reset_tracking_state(camera_id: str, model_id: str, reason: str = "session ended") -> None:
+    """Reset tracking state after session completion or timeout."""
+    tracking_manager.reset_tracking_state(camera_id, model_id, reason)
+
+
+def is_camera_active(camera_id: str, model_id: str) -> bool:
+    """Check if camera should be processing detections."""
+    return tracking_manager.is_camera_active(camera_id, model_id)
+
+
+def cleanup_camera_stability(camera_id: str) -> None:
+    """Clean up tracking data when a camera is disconnected."""
+    tracking_manager.cleanup_camera_tracking(camera_id)
+
+
+def occupancy_detector(camera_id: str, model_id: str, enable: bool = True) -> bool:
+    """
+    Enable or disable occupancy detection mode.
+    
+    When enabled:
+    - Model stops inference after completing full pipeline
+    - Backend sessionId handling continues in background
+    
+    Note: This is a temporary function that will be changed in the future.
+    """
+    return tracking_manager.set_occupancy_mode(camera_id, model_id, enable)
+
+
+def get_camera_tracking_data(camera_id: str, model_id: str) -> Dict[str, Any]:
+    """Get tracking data for a specific camera and model."""
+    tracking_data = tracking_manager.get_camera_tracking_data(camera_id, model_id)
+    return tracking_data.to_dict()
+
+
+def set_waiting_for_backend_session(camera_id: str, model_id: str, waiting: bool = True) -> None:
+    """Set whether camera is waiting for backend session ID."""
+    tracking_manager.set_waiting_for_backend_session(camera_id, model_id, waiting)
+
+
+def get_tracking_statistics(camera_id: str, model_id: str) -> Dict[str, Any]:
+    """Get comprehensive tracking statistics for a camera and model."""
+    return tracking_manager.get_tracking_statistics(camera_id, model_id)