nedo-vision-worker-core 0.2.0__py3-none-any.whl → 0.3.0__py3-none-any.whl

nedo_vision_worker_core/streams/VideoStream.py

@@ -3,322 +3,284 @@ import cv2
 import threading
 import time
 import logging
-from typing import Optional
+from typing import Optional, Union
+from enum import Enum
+
+
+class StreamState(Enum):
+    DISCONNECTED = "disconnected"
+    CONNECTING = "connecting"
+    CONNECTED = "connected"
+    RECONNECTING = "reconnecting"
+    STOPPED = "stopped"
 
 
 class VideoStream(threading.Thread):
-    """Threaded class for capturing video from a source, with automatic reconnection.
+    """Thread-safe video capture with automatic reconnection and error handling."""
     
-    This class provides a thread-safe way to capture video frames from various sources
-    (cameras, files, network streams) with automatic reconnection on failure.
-    
-    Attributes:
-        source: The video source (int for webcams, string for files/URLs)
-        reconnect_interval: Time in seconds to wait before reconnecting on failure
-        retry_limit: Maximum number of consecutive retries before giving up
-    """
     def __init__(
         self, 
-        source: str, 
-        reconnect_interval: int = 5,
-        retry_limit: int = 5,
+        source: Union[str, int], 
+        reconnect_interval: float = 5.0,
+        max_failures: int = 5,
         max_reconnect_attempts: int = 10,
-        reconnect_backoff_factor: float = 1.5
+        backoff_factor: float = 1.5,
+        target_fps: Optional[float] = None
     ):
-        super().__init__()
+        super().__init__(daemon=True)
         
-        # Stream configuration
         self.source = source
         self.reconnect_interval = reconnect_interval
-        self.retry_limit = retry_limit
+        self.max_failures = max_failures
+        self.max_reconnect_attempts = max_reconnect_attempts
+        self.backoff_factor = backoff_factor
+        self.target_fps = target_fps
         
-        # Stream state
         self.capture = None
-        self.running = True
-        self.connected = False
-        self.start_time = time.time()
-        self.is_file = isinstance(source, (int, str, bytes, os.PathLike)) and os.path.isfile(source)
-        self.fps = 30  # Default FPS until determined
+        self.state = StreamState.DISCONNECTED
+        self.fps = 30.0
         self.frame_count = 0
+        self.start_time = time.time()
         
-        # Reconnection control
-        self.max_reconnect_attempts = max_reconnect_attempts
-        self.reconnect_attempts = 0
-        self.reconnect_backoff_factor = reconnect_backoff_factor
-        self.current_reconnect_interval = reconnect_interval
-        
-        # Thread synchronization
-        self.lock = threading.Lock()
+        self._running = True
+        self._lock = threading.Lock()
         self._latest_frame = None
+        self._reconnect_attempts = 0
+        self._current_interval = reconnect_interval
         
-        # Start the capture thread
-        self.start()
-
-    def _initialize_capture(self) -> bool:
-        """Initialize or reinitialize the capture device.
+        self.is_file = self._is_file_source()
         
-        Returns:
-            True if successful, False otherwise
-        """
+    def _is_file_source(self) -> bool:
+        """Check if source is a file path."""
+        if isinstance(self.source, int):
+            return False
+        return isinstance(self.source, (str, bytes, os.PathLike)) and os.path.isfile(str(self.source))
+    
+    def _get_source_for_cv2(self) -> Union[str, int]:
+        """Convert source to format suitable for cv2.VideoCapture."""
+        if isinstance(self.source, str) and self.source.isdigit():
+            return int(self.source)
+        return self.source
+    
+    def _initialize_capture(self) -> bool:
+        """Initialize video capture device."""
         try:
-            logging.info(f"🔄 Attempting to connect to stream: {self.source} (attempt {self.reconnect_attempts + 1}/{self.max_reconnect_attempts})")
+            self.state = StreamState.CONNECTING
+            logging.info(f"Connecting to {self.source} (attempt {self._reconnect_attempts + 1})")
             
-            # Clean up existing capture if needed
             if self.capture:
                 self.capture.release()
-                
-            # Create new capture object
-            self.capture = cv2.VideoCapture(self.source)
+            
+            self.capture = cv2.VideoCapture(self._get_source_for_cv2())
             
             if not self.capture.isOpened():
-                logging.error(f"❌ Failed to open video source: {self.source}")
+                logging.error(f"Failed to open video source: {self.source}")
                 return False
-                
-            # Get FPS if available or estimate it
-            self.fps = self.capture.get(cv2.CAP_PROP_FPS)
-            
-            # Check FPS validity and estimate if needed
-            if not self.fps or self.fps <= 0 or self.fps > 240:
-                if self.reconnect_attempts:
-                    return False
-
-                logging.warning(f"⚠️ Invalid FPS reported ({self.fps}). Using default 30 FPS.")
-                self.fps = 30
             
-            if self.is_file:
-                total_frames = self.capture.get(cv2.CAP_PROP_FRAME_COUNT)
-                duration = (total_frames / self.fps) if self.fps > 0 and total_frames > 0 else 0
-                logging.info(f"✅ Connected to video file {self.source} at {self.fps:.2f} FPS, {total_frames} frames, {duration:.2f}s duration")
-            else:
-                logging.info(f"✅ Connected to stream {self.source} at {self.fps:.2f} FPS")
-                
-            self.connected = True
+            self._configure_capture()
+            self.state = StreamState.CONNECTED
             return True
             
         except Exception as e:
-            logging.error(f"❌ Error initializing capture for {self.source}: {e}")
-            self.connected = False
-            if self.capture:
-                try:
-                    self.capture.release()
-                except:
-                    pass
+            logging.error(f"Error initializing capture: {e}")
+            self._cleanup_capture()
+            return False
+    
+    def _configure_capture(self):
+        """Configure capture properties and determine FPS."""
+        detected_fps = self.capture.get(cv2.CAP_PROP_FPS)
+        
+        if self.target_fps:
+            self.fps = self.target_fps
+            if hasattr(cv2, 'CAP_PROP_FPS'):
+                self.capture.set(cv2.CAP_PROP_FPS, self.target_fps)
+        elif detected_fps and 0 < detected_fps <= 240:
+            self.fps = detected_fps
+        else:
+            self.fps = 30.0
+            logging.warning(f"Invalid FPS detected ({detected_fps}), using {self.fps}")
+        
+        if self.is_file:
+            total_frames = self.capture.get(cv2.CAP_PROP_FRAME_COUNT)
+            duration = total_frames / self.fps if self.fps > 0 else 0
+            logging.info(f"Video file: {self.fps:.1f} FPS, {int(total_frames)} frames, {duration:.1f}s")
+        else:
+            logging.info(f"Stream connected: {self.fps:.1f} FPS")
+    
+    def _cleanup_capture(self):
+        """Clean up capture resources."""
+        if self.capture:
+            try:
+                self.capture.release()
+            except Exception as e:
+                logging.error(f"Error releasing capture: {e}")
+            finally:
                 self.capture = None
+        self.state = StreamState.DISCONNECTED
+    
+    def _handle_reconnection(self) -> bool:
+        """Handle reconnection logic with backoff."""
+        if self._reconnect_attempts >= self.max_reconnect_attempts:
+            logging.error(f"Max reconnection attempts reached for {self.source}")
             return False
-
+        
+        self._reconnect_attempts += 1
+        self.state = StreamState.RECONNECTING
+        self._current_interval = min(self._current_interval * self.backoff_factor, 60.0)
+        
+        logging.warning(f"Reconnecting in {self._current_interval:.1f}s...")
+        return self._sleep_interruptible(self._current_interval)
+    
+    def _sleep_interruptible(self, duration: float) -> bool:
+        """Sleep with ability to interrupt on stop."""
+        end_time = time.time() + duration
+        while time.time() < end_time and self._running:
+            time.sleep(0.1)
+        return self._running
+    
+    def _handle_file_end(self) -> bool:
+        """Handle video file reaching end."""
+        if not self.is_file:
+            return False
+        
+        try:
+            current_pos = self.capture.get(cv2.CAP_PROP_POS_FRAMES)
+            total_frames = self.capture.get(cv2.CAP_PROP_FRAME_COUNT)
+            
+            if current_pos >= total_frames - 1:
+                logging.info(f"Video file ended, restarting: {self.source}")
+                self.capture.set(cv2.CAP_PROP_POS_FRAMES, 0)
+                return True
+        except Exception as e:
+            logging.error(f"Error handling file end: {e}")
+        
+        return False
+    
     def run(self):
-        """Main thread loop that continuously captures frames with automatic reconnection."""
-        retry_count = 0
-        frame_interval = 0  # Will be calculated once we know the FPS
+        """Main capture loop."""
+        failure_count = 0
+        frame_interval = 1.0 / self.fps
         
-        while self.running:
+        while self._running:
             try:
-                # Check if we should exit
-                if not self.running:
-                    break
-                    
-                # (Re)connect if needed
-                if self.capture is None or not self.capture.isOpened():
-                    # Check if we've exceeded the maximum reconnection attempts
-                    if self.reconnect_attempts >= self.max_reconnect_attempts:
-                        logging.error(f"❌ Exceeded maximum reconnection attempts ({self.max_reconnect_attempts}) for {self.source}. Giving up.")
-                        self.running = False
-                        break
-                    
-                    if retry_count:
-                        self.reconnect_attempts += 1
-
+                if not self.capture or not self.capture.isOpened():
                     if not self._initialize_capture():
-                        # Apply exponential backoff to reconnection interval
-                        self.current_reconnect_interval *= self.reconnect_backoff_factor
-                        
-                        logging.warning(f"⚠️ Reconnection attempt {self.reconnect_attempts}/{self.max_reconnect_attempts} "
-                                      f"failed. Next attempt in {self.current_reconnect_interval:.1f}s...")
-                        
-                        # Check for thread stop before sleeping
-                        if not self.running:
+                        if not self._handle_reconnection():
                             break
-                            
-                        time.sleep(self.current_reconnect_interval)
                         continue
                     
-                    # Reset counters on successful connection
-                    retry_count = 0
-                    # Only reset reconnect_attempts if we've been connected for a meaningful duration
-                    # This prevents rapid success/fail cycles from resetting the counter
-                    if self.connected:
-                        self.reconnect_attempts = 0
-                        self.current_reconnect_interval = self.reconnect_interval
-                    
-                    frame_interval = 1.0 / self.fps if self.fps > 0 else 0.033
-                
-                # Check if we should exit before reading frame
-                if not self.running:
-                    break
+                    failure_count = 0
+                    self._reconnect_attempts = 0
+                    self._current_interval = self.reconnect_interval
+                    frame_interval = 1.0 / self.fps
                 
-                # Read the next frame
-                read_start = time.time()
+                start_time = time.time()
                 ret, frame = self.capture.read()
                 
-                # Handle frame read failure
                 if not ret or frame is None or frame.size == 0:
-                    if self.is_file:
-                        # For video files, check if we've reached the end
-                        current_pos = self.capture.get(cv2.CAP_PROP_POS_FRAMES)
-                        total_frames = self.capture.get(cv2.CAP_PROP_FRAME_COUNT)
-                        
-                        # If we're at or beyond the end of the video, restart from beginning
-                        if current_pos >= total_frames - 1:
-                            logging.info(f"🔄 Video file {self.source} reached end. Restarting from beginning...")
-                            self.capture.set(cv2.CAP_PROP_POS_FRAMES, 0)
-                            # Reset retry count for video files when restarting
-                            retry_count = 0
-                            continue
-                        else:
-                            # We're not at the end, this might be a real error
-                            retry_count += 1
-                            logging.warning(f"⚠️ Failed to read frame from {self.source} at position {current_pos}/{total_frames} (attempt {retry_count}/{self.retry_limit})")
-                    else:
-                        # For non-file sources (cameras, streams), increment retry count
-                        retry_count += 1
-                        logging.warning(f"⚠️ Failed to read frame from {self.source} (attempt {retry_count}/{self.retry_limit})")
+                    if self._handle_file_end():
+                        continue
                     
-                    if retry_count > self.retry_limit:
-                        logging.error(f"❌ Too many consecutive frame failures. Reconnecting...")
-                        self.connected = False
-                        if self.capture and self.running:  # Only release if we're still running
-                            try:
-                                self.capture.release()
-                                self.capture = None
-                            except Exception as e:
-                                logging.error(f"Error releasing capture during failure: {e}")
+                    failure_count += 1
+                    if failure_count > self.max_failures:
+                        logging.error("Too many consecutive failures, reconnecting...")
+                        self._cleanup_capture()
+                        failure_count = 0
                         continue
                     
-                    if not self.running:
+                    if not self._sleep_interruptible(0.1):
                         break
-
-                    time.sleep(0.1)
                     continue
                 
-                # Reset retry count on successful frame
-                retry_count = 0
+                failure_count = 0
                 self.frame_count += 1
                 
-                # Store the frame - but make sure we're still running first
-                if not self.running:
-                    break
-                    
-                with self.lock:
-                    self._latest_frame = frame.copy()
+                with self._lock:
+                    if self._running:
+                        self._latest_frame = frame.copy()
                 
-                # Regulate frame rate to match source FPS for efficiency
-                # This helps prevent CPU overuse when reading from fast sources
-                elapsed = time.time() - read_start
+                elapsed = time.time() - start_time
                 sleep_time = max(0, frame_interval - elapsed)
-                if sleep_time > 0 and self.running:  # Check running before sleep
-                    time.sleep(sleep_time)
+                if sleep_time > 0 and not self._sleep_interruptible(sleep_time):
+                    break
                     
-            except cv2.error as cv_err:
-                logging.error(f"❌ OpenCV error in frame processing: {cv_err}")
-                self.connected = False
-                if not self.running:
+            except cv2.error as e:
+                logging.error(f"OpenCV error: {e}")
+                self._cleanup_capture()
+                if not self._sleep_interruptible(1.0):
                     break
-                time.sleep(1)  # Brief pause before retry
-                
+                    
             except Exception as e:
-                logging.error(f"❌ Error in VideoStream {self.source}: {e}", exc_info=True)
-                self.connected = False
-                if not self.running:
+                logging.error(f"Unexpected error: {e}", exc_info=True)
+                if not self._sleep_interruptible(self.reconnect_interval):
                     break
-                time.sleep(self.reconnect_interval)
         
-        # Final cleanup
-        self._cleanup()
-
-    def _cleanup(self):
-        """Release resources and perform cleanup."""
-        try:
-            if self.capture:
-                self.capture.release()
-                self.capture = None
-                
-            with self.lock:
-                self._latest_frame = None
-                        
-            self.connected = False
-            logging.info(f"🔴 Stream {self.source} stopped and cleaned up.")
-            
-        except Exception as e:
-            logging.error(f"Error during final VideoStream cleanup: {e}")
-
-    def get_frame(self) -> Optional[cv2.Mat]:
-        """Returns the latest available frame (thread-safe).
+        self._final_cleanup()
+    
+    def _final_cleanup(self):
+        """Final resource cleanup."""
+        self.state = StreamState.STOPPED
+        self._cleanup_capture()
         
-        Returns:
-            A copy of the latest frame or None if no frames are available
-        """
-        if not self.running:
+        with self._lock:
+            self._latest_frame = None
+        
+        logging.info(f"VideoStream stopped: {self.source}")
+    
+    def get_frame(self) -> Optional[cv2.Mat]:
+        """Get the latest frame (thread-safe)."""
+        if not self._running or self.state not in (StreamState.CONNECTED, StreamState.RECONNECTING):
             return None
         
-        with self.lock:
-            if self._latest_frame is not None:
-                return self._latest_frame.copy()
-        return None
-
+        with self._lock:
+            return self._latest_frame.copy() if self._latest_frame is not None else None
+    
+    def is_connected(self) -> bool:
+        """Check if stream is currently connected."""
+        return self.state == StreamState.CONNECTED
+    
+    @property
+    def running(self) -> bool:
+        """Check if stream is currently running."""
+        return self._running and self.state != StreamState.STOPPED
+    
+    def get_state(self) -> StreamState:
+        """Get current stream state."""
+        return self.state
+    
     def is_video_ended(self) -> bool:
-        """Check if a video file has reached its end (only for video files).
-        
-        Returns:
-            True if video file has ended, False otherwise
-        """
-        if not self.is_file or not self.capture or not self.capture.isOpened():
+        """Check if video file has ended."""
+        if not self.is_file or not self.capture:
             return False
-            
+        
         try:
             current_pos = self.capture.get(cv2.CAP_PROP_POS_FRAMES)
             total_frames = self.capture.get(cv2.CAP_PROP_FRAME_COUNT)
             return current_pos >= total_frames - 1
         except Exception:
             return False
-
-    def stop(self):
-        """Stops the video stream and releases resources safely."""
-        if not self.running:  # Prevent multiple stops
-            return
     
-        # Step 1: Signal the thread to stop
-        self.running = False
+    def stop(self, timeout: float = 5.0):
+        """Stop the video stream gracefully."""
+        if not self._running:
+            return
         
-        # Step 2: Wait for any ongoing OpenCV operations to complete
-        time.sleep(0.2)  # Short delay to let operations finish
+        logging.info(f"Stopping VideoStream: {self.source}")
+        self._running = False
         
-        # Step 3: Explicitly release capture device before joining thread
-        # This is critical to avoid segmentation faults in some OpenCV implementations
-        if self.capture:
-            try:
-                logging.debug(f"Releasing capture device for {self.source}")
-                with self.lock:  # Use lock to ensure thread isn't accessing capture during release
-                    if self.capture:
-                        self.capture.release()
-                        self.capture = None
-            except Exception as e:
-                logging.error(f"Error releasing capture device: {e}")
-        
-        # Step 4: Clear any references to frames to help garbage collection
-        with self.lock:
+        with self._lock:
             self._latest_frame = None
-                
-        # Step 5: Wait for thread to exit
-        try:
-            if self.is_alive():
-                logging.debug(f"Waiting for thread to exit for stream {self.source}")
-                self.join(timeout=30)  # Reduced timeout since we already released resources
-                
-                # Check if thread is still alive
-                if self.is_alive():
-                    logging.warning(f"⚠️ Stream {self.source} thread did not exit cleanly within timeout")
-        except Exception as e:
-            logging.error(f"Error joining thread: {e}")
         
-        # Final status update
-        self.connected = False
+        if self.is_alive():
+            self.join(timeout=timeout)
+            if self.is_alive():
+                logging.warning(f"Stream thread did not exit within {timeout}s")
+    
+    def __enter__(self):
+        """Context manager entry."""
+        self.start()
+        return self
+    
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Context manager exit."""
+        self.stop()