rocket-welder-sdk 1.1.43__py3-none-any.whl → 1.1.45__py3-none-any.whl

rocket_welder_sdk/rocket_welder_client.py

@@ -15,7 +15,12 @@ import numpy as np
 from .connection_string import ConnectionMode, ConnectionString, Protocol
 from .controllers import DuplexShmController, IController, OneWayShmController
 from .frame_metadata import FrameMetadata  # noqa: TC001 - used at runtime in callbacks
-from .high_level.connection_strings import KeyPointsConnectionString, SegmentationConnectionString
+from .graphics import ILayerCanvas, IStageSink, IStageWriter, RgbColor, StageSink
+from .high_level.connection_strings import (
+    GraphicsConnectionString,
+    KeyPointsConnectionString,
+    SegmentationConnectionString,
+)
 from .high_level.frame_sink_factory import FrameSinkFactory
 from .keypoints_protocol import IKeyPointsSink, IKeyPointsWriter, KeyPointsSink
 from .opencv_controller import OpenCvController
@@ -24,12 +29,6 @@ from .segmentation_result import (
     ISegmentationResultWriter,
     SegmentationResultSink,
 )
-from .session_id import (
-    get_configured_nng_urls,
-    get_nng_urls_from_env,
-    has_explicit_nng_urls,
-)
-from .transport.nng_transport import NngFrameSink
 
 if TYPE_CHECKING:
     import numpy.typing as npt
@@ -67,9 +66,6 @@ class RocketWelderClient:
         self._controller: Optional[IController] = None
         self._lock = threading.Lock()
 
-        # NNG publishers for streaming results (auto-created if SessionId env var is set)
-        self._nng_publishers: dict[str, NngFrameSink] = {}
-
         # Preview support
         self._preview_enabled = (
             self._connection.parameters.get("preview", "false").lower() == "true"
@@ -89,50 +85,6 @@ class RocketWelderClient:
         with self._lock:
             return self._controller is not None and self._controller.is_running
 
-    @property
-    def nng_publishers(self) -> dict[str, NngFrameSink]:
-        """Get NNG publishers for streaming results.
-
-        Returns:
-            Dictionary with 'segmentation', 'keypoints', 'actions' keys.
-            Empty if SessionId env var was not set at startup.
-
-        Example:
-            client.nng_publishers["segmentation"].write_frame(seg_data)
-        """
-        return self._nng_publishers
-
-    def _create_nng_publishers(self) -> None:
-        """Create NNG publishers for result streaming.
-
-        URLs are read from environment variables (preferred) or derived from SessionId (fallback).
-
-        Priority:
-        1. Explicit URLs: SEGMENTATION_SINK_URL, KEYPOINTS_SINK_URL, ACTIONS_SINK_URL
-        2. Derived from SessionId environment variable (backwards compatibility)
-        """
-        try:
-            urls = get_configured_nng_urls()
-
-            for name, url in urls.items():
-                sink = NngFrameSink.create_publisher(url)
-                self._nng_publishers[name] = sink
-                logger.info("NNG publisher ready: %s at %s", name, url)
-
-            # Log configuration summary
-            logger.info(
-                "NNG publishers configured: seg=%s, kp=%s, actions=%s",
-                urls.get("segmentation", "(not configured)"),
-                urls.get("keypoints", "(not configured)"),
-                urls.get("actions", "(not configured)"),
-            )
-        except ValueError as ex:
-            # No URLs configured - this is expected for containers that don't publish results
-            logger.debug("NNG publishers not configured: %s", ex)
-        except Exception as ex:
-            logger.warning("Failed to create NNG publishers: %s", ex)
-            # Don't fail start() - NNG is optional for backwards compatibility
-
     def get_metadata(self) -> Optional[GstMetadata]:
         """
         Get the current GStreamer metadata.
@@ -180,21 +132,6 @@ class RocketWelderClient:
             else:
                 raise ValueError(f"Unsupported protocol: {self._connection.protocol}")
 
-            # Auto-create NNG publishers if URLs are configured
-            # (explicit URLs via SEGMENTATION_SINK_URL etc., or derived from SessionId)
-            if has_explicit_nng_urls():
-                self._create_nng_publishers()
-            else:
-                # Log that NNG is not configured (informational)
-                urls = get_nng_urls_from_env()
-                logger.info(
-                    "NNG sink URLs not configured (this is normal if not publishing AI results). "
-                    "seg=%s, kp=%s, actions=%s",
-                    urls.get("segmentation") or "(not set)",
-                    urls.get("keypoints") or "(not set)",
-                    urls.get("actions") or "(not set)",
-                )
-
             # If preview is enabled, wrap the callback to capture frames
             if self._preview_enabled:
                 self._original_callback = on_frame
@@ -256,25 +193,26 @@ class RocketWelderClient:
 
     def start_with_writers(
         self,
-        on_frame: Callable[[Mat, ISegmentationResultWriter, IKeyPointsWriter, Mat], None],  # type: ignore[valid-type]
+        on_frame: Callable[[Mat, ISegmentationResultWriter, IKeyPointsWriter, IStageWriter, Mat], None],  # type: ignore[valid-type]
         cancellation_token: Optional[threading.Event] = None,
     ) -> None:
         """
-        Start receiving frames with segmentation and keypoints output support.
+        Start receiving frames with segmentation, keypoints, and graphics output support.
 
         Creates sinks for streaming AI results to rocket-welder2.
 
         Configuration via environment variables:
         - SEGMENTATION_SINK_URL: URL for segmentation output (e.g., socket:///tmp/seg.sock)
         - KEYPOINTS_SINK_URL: URL for keypoints output (e.g., socket:///tmp/kp.sock)
+        - STAGE_SINK_URL: URL for graphics/stage output (e.g., socket:///tmp/stage.sock)
 
         Args:
-            on_frame: Callback receiving (input_mat, seg_writer, kp_writer, output_mat).
+            on_frame: Callback receiving (input_mat, seg_writer, kp_writer, stage_writer, output_mat).
                      The writers are created per-frame and auto-flush on context exit.
             cancellation_token: Optional cancellation token
 
         Example:
-            def process_frame(input_mat, seg_writer, kp_writer, output_mat):
+            def process_frame(input_mat, seg_writer, kp_writer, stage_writer, output_mat):
                 # Run AI inference
                 result = ai_model.infer(input_mat)
 
@@ -286,6 +224,11 @@ class RocketWelderClient:
                 for kp in result.keypoints:
                     kp_writer.append(kp.id, kp.x, kp.y, kp.confidence)
 
+                # Draw graphics overlay
+                layer = stage_writer[0]
+                layer.set_font_size(24)
+                layer.draw_text("Detection count: 5", 10, 30)
+
                 # Copy/draw to output
                 output_mat[:] = input_mat
 
@@ -316,11 +259,13 @@ class RocketWelderClient:
             # Create sinks from environment
             seg_sink = self._get_or_create_segmentation_sink()
             kp_sink = self._get_or_create_keypoints_sink()
+            stage_sink = self._get_or_create_stage_sink()
 
             logger.info(
-                "Starting RocketWelder client with AI output support: seg=%s, kp=%s",
+                "Starting RocketWelder client with AI output support: seg=%s, kp=%s, stage=%s",
                 "configured" if seg_sink else "null",
                 "configured" if kp_sink else "null",
+                "configured" if stage_sink else "null",
             )
 
             # Wrapper callback that creates per-frame writers
@@ -337,17 +282,26 @@ class RocketWelderClient:
                         frame_metadata.frame_number,
                     )
                     # Use no-op writers
-                    on_frame(
-                        input_mat, _NoOpSegmentationWriter(), _NoOpKeyPointsWriter(), output_mat
-                    )
+                    with stage_sink.create_writer(frame_metadata.frame_number) as stage_writer:
+                        on_frame(
+                            input_mat,
+                            _NoOpSegmentationWriter(),
+                            _NoOpKeyPointsWriter(),
+                            stage_writer,
+                            output_mat,
+                        )
                     return
 
-                # Create per-frame writers from sinks
+                # Create per-frame writers from sinks (all auto-flush on context exit)
                 with seg_sink.create_writer(
                     frame_metadata.frame_number, caps.width, caps.height
-                ) as seg_writer, kp_sink.create_writer(frame_metadata.frame_number) as kp_writer:
+                ) as seg_writer, kp_sink.create_writer(
+                    frame_metadata.frame_number
+                ) as kp_writer, stage_sink.create_writer(
+                    frame_metadata.frame_number
+                ) as stage_writer:
                     # Call user callback with writers
-                    on_frame(input_mat, seg_writer, kp_writer, output_mat)
+                    on_frame(input_mat, seg_writer, kp_writer, stage_writer, output_mat)
                     # Writers auto-flush on context exit
 
             # Start the controller with our wrapper
@@ -392,6 +346,23 @@ class RocketWelderClient:
             logger.warning("Failed to create keypoints sink from %s: %s", url, ex)
             return _NullKeyPointsSink()
 
+    def _get_or_create_stage_sink(self) -> IStageSink:
+        """Get or create graphics stage sink from environment."""
+        import os
+
+        url = os.environ.get("GRAPHICS_SINK_URL")
+        if not url:
+            logger.debug("GRAPHICS_SINK_URL not set, using null sink")
+            return _NullStageSink()
+
+        try:
+            cs = GraphicsConnectionString.parse(url)
+            frame_sink = FrameSinkFactory.create(cs.protocol, cs.address)
+            return StageSink(frame_sink=frame_sink, owns_sink=True)
+        except Exception as ex:
+            logger.warning("Failed to create graphics stage sink from %s: %s", url, ex)
+            return _NullStageSink()
+
     def stop(self) -> None:
         """Stop the client and clean up resources."""
         with self._lock:
@@ -403,15 +374,6 @@ class RocketWelderClient:
                 if self._preview_enabled:
                     self._preview_queue.put(None)  # Sentinel value
 
-                # Clean up NNG publishers
-                for name, sink in self._nng_publishers.items():
-                    try:
-                        sink.close()
-                        logger.debug("Closed NNG publisher: %s", name)
-                    except Exception as ex:
-                        logger.warning("Failed to close NNG publisher %s: %s", name, ex)
-                self._nng_publishers.clear()
-
                 logger.info("RocketWelder client stopped")
 
     def show(self, cancellation_token: Optional[threading.Event] = None) -> None:
@@ -697,3 +659,162 @@ class _NullSegmentationSink(ISegmentationResultSink):
     def close(self) -> None:
         """No-op close."""
         pass
+
+
+class _NoOpLayerCanvas(ILayerCanvas):
+    """No-op layer canvas that discards all drawing operations."""
+
+    @property
+    def layer_id(self) -> int:
+        """The layer ID."""
+        return 0
+
+    # Frame type
+    def master(self) -> None:
+        """No-op."""
+        pass
+
+    def remain(self) -> None:
+        """No-op."""
+        pass
+
+    def clear(self) -> None:
+        """No-op."""
+        pass
+
+    # Context state - Styling
+    def set_stroke(self, color: RgbColor) -> None:
+        """No-op."""
+        pass
+
+    def set_fill(self, color: RgbColor) -> None:
+        """No-op."""
+        pass
+
+    def set_thickness(self, width: int) -> None:
+        """No-op."""
+        pass
+
+    def set_font_size(self, size: int) -> None:
+        """No-op."""
+        pass
+
+    def set_font_color(self, color: RgbColor) -> None:
+        """No-op."""
+        pass
+
+    # Context state - Transforms
+    def translate(self, dx: float, dy: float) -> None:
+        """No-op."""
+        pass
+
+    def rotate(self, degrees: float) -> None:
+        """No-op."""
+        pass
+
+    def scale(self, sx: float, sy: float) -> None:
+        """No-op."""
+        pass
+
+    def skew(self, kx: float, ky: float) -> None:
+        """No-op."""
+        pass
+
+    def set_matrix(
+        self,
+        scale_x: float,
+        skew_x: float,
+        trans_x: float,
+        skew_y: float,
+        scale_y: float,
+        trans_y: float,
+    ) -> None:
+        """No-op."""
+        pass
+
+    # Context stack
+    def save(self) -> None:
+        """No-op."""
+        pass
+
+    def restore(self) -> None:
+        """No-op."""
+        pass
+
+    def reset_context(self) -> None:
+        """No-op."""
+        pass
+
+    # Draw operations
+    def draw_polygon(self, points: Any) -> None:
+        """No-op."""
+        pass
+
+    def draw_text(self, text: str, x: int, y: int) -> None:
+        """No-op."""
+        pass
+
+    def draw_circle(self, center_x: int, center_y: int, radius: int) -> None:
+        """No-op."""
+        pass
+
+    def draw_rectangle(self, x: int, y: int, width: int, height: int) -> None:
+        """No-op."""
+        pass
+
+    def draw_line(self, x1: int, y1: int, x2: int, y2: int) -> None:
+        """No-op."""
+        pass
+
+    def draw_jpeg(self, jpeg_data: bytes, x: int, y: int, width: int, height: int) -> None:
+        """No-op."""
+        pass
+
+
+# Singleton instance
+_NO_OP_LAYER_CANVAS = _NoOpLayerCanvas()
+
+
+class _NoOpStageWriter(IStageWriter):
+    """No-op stage writer that discards all graphics operations."""
+
+    @property
+    def frame_id(self) -> int:
+        """The frame ID."""
+        return 0
+
+    def __getitem__(self, layer_id: int) -> ILayerCanvas:
+        """Returns no-op layer canvas."""
+        return _NO_OP_LAYER_CANVAS
+
+    def layer(self, layer_id: int) -> ILayerCanvas:
+        """Returns no-op layer canvas."""
+        return _NO_OP_LAYER_CANVAS
+
+    def close(self) -> None:
+        """No-op close."""
+        pass
+
+    def __enter__(self) -> _NoOpStageWriter:
+        """Context manager entry."""
+        return self
+
+    def __exit__(self, exc_type: object, exc_val: object, exc_tb: object) -> None:
+        """Context manager exit."""
+        pass
+
+
+# Singleton instance
+_NO_OP_STAGE_WRITER = _NoOpStageWriter()
+
+
+class _NullStageSink(IStageSink):
+    """Null stage sink that creates no-op writers."""
+
+    def create_writer(self, frame_id: int) -> IStageWriter:
+        """Create a no-op writer."""
+        return _NO_OP_STAGE_WRITER
+
+    def close(self) -> None:
+        """No-op close."""
+        pass