rocket-welder-sdk 1.1.27__py3-none-any.whl → 1.1.29__py3-none-any.whl

rocket_welder_sdk/rocket_welder_client.py

@@ -6,19 +6,25 @@ Main entry point for the RocketWelder SDK.
 from __future__ import annotations
 
 import logging
+import queue
 import threading
-from typing import TYPE_CHECKING, Any, Callable
+from typing import TYPE_CHECKING, Any, Callable, List, Optional, Union
 
 import numpy as np
 
 from .connection_string import ConnectionMode, ConnectionString, Protocol
 from .controllers import DuplexShmController, IController, OneWayShmController
+from .opencv_controller import OpenCvController
 
 if TYPE_CHECKING:
+    import numpy.typing as npt
+
     from .gst_metadata import GstMetadata
 
-# Type alias for OpenCV Mat
-Mat = np.ndarray[Any, Any]
+    # Use numpy array type for Mat - OpenCV Mat is essentially a numpy array
+    Mat = npt.NDArray[np.uint8]
+else:
+    Mat = np.ndarray  # type: ignore[misc]
 
 # Module logger
 logger = logging.getLogger(__name__)
@@ -31,7 +37,7 @@ class RocketWelderClient:
     Provides a unified interface for different connection types and protocols.
     """
 
-    def __init__(self, connection: str | ConnectionString):
+    def __init__(self, connection: Union[str, ConnectionString]):
         """
         Initialize the RocketWelder client.
 
@@ -43,9 +49,17 @@ class RocketWelderClient:
         else:
             self._connection = connection
 
-        self._controller: IController | None = None
+        self._controller: Optional[IController] = None
         self._lock = threading.Lock()
 
+        # Preview support
+        self._preview_enabled = (
+            self._connection.parameters.get("preview", "false").lower() == "true"
+        )
+        self._preview_queue: queue.Queue[Optional[Mat]] = queue.Queue(maxsize=2)  # type: ignore[valid-type]  # Small buffer
+        self._preview_window_name = "RocketWelder Preview"
+        self._original_callback: Any = None
+
     @property
     def connection(self) -> ConnectionString:
         """Get the connection configuration."""
@@ -57,7 +71,7 @@ class RocketWelderClient:
         with self._lock:
             return self._controller is not None and self._controller.is_running
 
-    def get_metadata(self) -> GstMetadata | None:
+    def get_metadata(self) -> Optional[GstMetadata]:
         """
         Get the current GStreamer metadata.
 
@@ -71,8 +85,8 @@ class RocketWelderClient:
 
     def start(
         self,
-        on_frame: Callable[[Mat], None] | Callable[[Mat, Mat], None],
-        cancellation_token: threading.Event | None = None,
+        on_frame: Union[Callable[[Mat], None], Callable[[Mat, Mat], None]],  # type: ignore[valid-type]
+        cancellation_token: Optional[threading.Event] = None,
     ) -> None:
         """
         Start receiving/processing video frames.
@@ -97,11 +111,55 @@ class RocketWelderClient:
                     self._controller = DuplexShmController(self._connection)
                 else:
                     self._controller = OneWayShmController(self._connection)
+            elif self._connection.protocol in (Protocol.FILE, Protocol.MJPEG):
+                self._controller = OpenCvController(self._connection)
             else:
                 raise ValueError(f"Unsupported protocol: {self._connection.protocol}")
 
+            # If preview is enabled, wrap the callback to capture frames
+            if self._preview_enabled:
+                self._original_callback = on_frame
+
+                # Determine if duplex or one-way
+                if self._connection.connection_mode == ConnectionMode.DUPLEX:
+
+                    def preview_wrapper_duplex(input_frame: Mat, output_frame: Mat) -> None:  # type: ignore[valid-type]
+                        # Call original callback
+                        on_frame(input_frame, output_frame)  # type: ignore[call-arg]
+                        # Queue the OUTPUT frame for preview
+                        try:
+                            self._preview_queue.put_nowait(output_frame.copy())  # type: ignore[attr-defined]
+                        except queue.Full:
+                            # Drop oldest frame if queue is full
+                            try:
+                                self._preview_queue.get_nowait()
+                                self._preview_queue.put_nowait(output_frame.copy())  # type: ignore[attr-defined]
+                            except queue.Empty:
+                                pass
+
+                    actual_callback = preview_wrapper_duplex
+                else:
+
+                    def preview_wrapper_oneway(frame: Mat) -> None:  # type: ignore[valid-type]
+                        # Call original callback
+                        on_frame(frame)  # type: ignore[call-arg]
+                        # Queue frame for preview
+                        try:
+                            self._preview_queue.put_nowait(frame.copy())  # type: ignore[attr-defined]
+                        except queue.Full:
+                            # Drop oldest frame if queue is full
+                            try:
+                                self._preview_queue.get_nowait()
+                                self._preview_queue.put_nowait(frame.copy())  # type: ignore[attr-defined]
+                            except queue.Empty:
+                                pass
+
+                    actual_callback = preview_wrapper_oneway  # type: ignore[assignment]
+            else:
+                actual_callback = on_frame  # type: ignore[assignment]
+
             # Start the controller
-            self._controller.start(on_frame, cancellation_token)  # type: ignore[arg-type]
+            self._controller.start(actual_callback, cancellation_token)  # type: ignore[arg-type]
             logger.info("RocketWelder client started with %s", self._connection)
 
     def stop(self) -> None:
@@ -110,8 +168,83 @@ class RocketWelderClient:
             if self._controller:
                 self._controller.stop()
                 self._controller = None
+
+                # Signal preview to stop if enabled
+                if self._preview_enabled:
+                    self._preview_queue.put(None)  # Sentinel value
+
                 logger.info("RocketWelder client stopped")
 
+    def show(self, cancellation_token: Optional[threading.Event] = None) -> None:
+        """
+        Display preview frames in a window (main thread only).
+
+        This method should be called from the main thread after start().
+        - If preview=true: blocks and displays frames until stopped or 'q' pressed
+        - If preview=false or not set: returns immediately
+
+        Args:
+            cancellation_token: Optional cancellation token to stop preview
+
+        Example:
+            client = RocketWelderClient("file:///video.mp4?preview=true")
+            client.start(process_frame)
+            client.show()  # Blocks and shows preview
+            client.stop()
+        """
+        if not self._preview_enabled:
+            # No preview requested, return immediately
+            return
+
+        try:
+            import cv2
+        except ImportError:
+            logger.warning("OpenCV not available, cannot show preview")
+            return
+
+        logger.info("Starting preview display in main thread")
+
+        # Create window
+        cv2.namedWindow(self._preview_window_name, cv2.WINDOW_NORMAL)
+
+        try:
+            while True:
+                # Check for cancellation
+                if cancellation_token and cancellation_token.is_set():
+                    break
+
+                try:
+                    # Get frame with timeout
+                    frame = self._preview_queue.get(timeout=0.1)
+
+                    # Check for stop sentinel
+                    if frame is None:
+                        break
+
+                    # Display frame
+                    cv2.imshow(self._preview_window_name, frame)
+
+                    # Process window events and check for 'q' key
+                    key = cv2.waitKey(1) & 0xFF
+                    if key == ord("q"):
+                        logger.info("User pressed 'q', stopping preview")
+                        break
+
+                except queue.Empty:
+                    # No frame available, check if still running
+                    if not self.is_running:
+                        break
+                    # Process window events even without new frame
+                    if cv2.waitKey(1) & 0xFF == ord("q"):
+                        logger.info("User pressed 'q', stopping preview")
+                        break
+
+        finally:
+            # Clean up window
+            cv2.destroyWindow(self._preview_window_name)
+            cv2.waitKey(1)  # Process pending events
+            logger.info("Preview display stopped")
+
     def __enter__(self) -> RocketWelderClient:
         """Context manager entry."""
         return self
@@ -120,6 +253,110 @@ class RocketWelderClient:
         """Context manager exit."""
         self.stop()
 
+    @classmethod
+    def from_connection_string(cls, connection_string: str) -> RocketWelderClient:
+        """
+        Create a client from a connection string.
+
+        Args:
+            connection_string: Connection string (e.g., 'shm://buffer?mode=Duplex')
+
+        Returns:
+            Configured RocketWelderClient instance
+        """
+        return cls(connection_string)
+
+    @classmethod
+    def from_args(cls, args: List[str]) -> RocketWelderClient:
+        """
+        Create a client from command line arguments.
+
+        Checks in order:
+        1. First positional argument from args
+        2. CONNECTION_STRING environment variable
+
+        Args:
+            args: Command line arguments (typically sys.argv)
+
+        Returns:
+            Configured RocketWelderClient instance
+
+        Raises:
+            ValueError: If no connection string is found
+        """
+        import os
+
+        # Check for positional argument (skip script name if present)
+        connection_string = None
+        for arg in args[1:] if len(args) > 0 and args[0].endswith(".py") else args:
+            if not arg.startswith("-"):
+                connection_string = arg
+                break
+
+        # Fall back to environment variable
+        if not connection_string:
+            connection_string = os.environ.get("CONNECTION_STRING")
+
+        if not connection_string:
+            raise ValueError(
+                "No connection string provided. "
+                "Provide as argument or set CONNECTION_STRING environment variable"
+            )
+
+        return cls(connection_string)
+
+    @classmethod
+    def from_(cls, *args: Any, **kwargs: Any) -> RocketWelderClient:
+        """
+        Create a client with automatic configuration detection.
+
+        This is the most convenient factory method that:
+        1. Checks kwargs for 'args' parameter (command line arguments)
+        2. Checks args for command line arguments
+        3. Falls back to CONNECTION_STRING environment variable
+
+        Examples:
+            client = RocketWelderClient.from_()  # Uses env var
+            client = RocketWelderClient.from_(sys.argv)  # Uses command line
+            client = RocketWelderClient.from_(args=sys.argv)  # Named param
+
+        Returns:
+            Configured RocketWelderClient instance
+
+        Raises:
+            ValueError: If no connection string is found
+        """
+        import os
+
+        # Check kwargs first
+        argv = kwargs.get("args")
+
+        # Then check positional args
+        if not argv and args:
+            # If first arg looks like sys.argv (list), use it
+            if isinstance(args[0], list):
+                argv = args[0]
+            # If first arg is a string, treat it as connection string
+            elif isinstance(args[0], str):
+                return cls(args[0])
+
+        # Try to get from command line args if provided
+        if argv:
+            try:
+                return cls.from_args(argv)
+            except ValueError:
+                pass  # Fall through to env var check
+
+        # Fall back to environment variable
+        connection_string = os.environ.get("CONNECTION_STRING")
+        if connection_string:
+            return cls(connection_string)
+
+        raise ValueError(
+            "No connection string provided. "
+            "Provide as argument or set CONNECTION_STRING environment variable"
+        )
+
     @classmethod
     def create_oneway_shm(
         cls,