rocket-welder-sdk 1.1.42__py3-none-any.whl → 1.1.44__py3-none-any.whl

rocket_welder_sdk/rocket_welder_client.py

@@ -24,12 +24,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 +61,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 +80,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 +127,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
@@ -403,15 +335,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:

rocket_welder_sdk/segmentation_result.py

@@ -23,13 +23,22 @@ import io
 import struct
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from typing import BinaryIO, Iterator, List, Optional, Tuple, Union
+from typing import (
+    AsyncIterator,
+    BinaryIO,
+    Iterator,
+    List,
+    Optional,
+    Sequence,
+    Tuple,
+    Union,
+)
 
 import numpy as np
 import numpy.typing as npt
 from typing_extensions import TypeAlias
 
-from .transport import IFrameSink, StreamFrameSink
+from .transport import IFrameSink, IFrameSource, StreamFrameSink
 
 # Type aliases
 Point = Tuple[int, int]
@@ -122,6 +131,40 @@ class SegmentationInstance:
         return [(int(x), int(y)) for x, y in self.points]
 
 
+@dataclass(frozen=True)
+class SegmentationFrame:
+    """
+    Represents a decoded segmentation frame containing instance segmentation results.
+
+    Matches C# SegmentationFrame record struct.
+    Used for round-trip testing of segmentation protocol encoding/decoding.
+
+    Attributes:
+        frame_id: Frame identifier for temporal ordering.
+        width: Frame width in pixels.
+        height: Frame height in pixels.
+        instances: Segmentation instances detected in this frame.
+    """
+
+    frame_id: int
+    width: int
+    height: int
+    instances: Sequence[SegmentationInstance]
+
+    @property
+    def count(self) -> int:
+        """Number of instances in the frame."""
+        return len(self.instances)
+
+    def find_by_class(self, class_id: int) -> List[SegmentationInstance]:
+        """Find all instances with the specified class ID."""
+        return [inst for inst in self.instances if inst.class_id == class_id]
+
+    def find_by_instance(self, instance_id: int) -> List[SegmentationInstance]:
+        """Find all instances with the specified instance ID."""
+        return [inst for inst in self.instances if inst.instance_id == instance_id]
+
+
 class SegmentationResultWriter:
     """
     Writes segmentation results for a single frame via IFrameSink.
@@ -550,3 +593,345 @@ class SegmentationResultSink(ISegmentationResultSink):
         """Close the sink and release resources."""
         if self._owns_sink:
             self._frame_sink.close()
+
+
+class ISegmentationResultSource(ABC):
+    """
+    Interface for streaming segmentation frames from a source.
+
+    Mirrors the pattern from IKeyPointsSource for consistency.
+    """
+
+    @abstractmethod
+    def read_frames(self) -> Iterator[SegmentationFrame]:
+        """
+        Read frames synchronously as an iterator.
+
+        Yields:
+            SegmentationFrame for each frame in the source.
+        """
+        pass
+
+    def read_frames_async(self) -> AsyncIterator[SegmentationFrame]:
+        """
+        Read frames asynchronously as an async iterator.
+
+        Yields:
+            SegmentationFrame for each frame in the source.
+
+        Raises:
+            NotImplementedError: Subclass must implement for async support.
+        """
+        raise NotImplementedError("Subclass must implement read_frames_async")
+
+    def close(self) -> None:  # noqa: B027
+        """Close the source and release resources."""
+        pass
+
+    async def close_async(self) -> None:  # noqa: B027
+        """Close the source and release resources asynchronously."""
+        pass
+
+    def __enter__(self) -> "ISegmentationResultSource":
+        """Context manager entry."""
+        return self
+
+    def __exit__(self, *args: object) -> None:
+        """Context manager exit."""
+        self.close()
+
+    async def __aenter__(self) -> "ISegmentationResultSource":
+        """Async context manager entry."""
+        return self
+
+    async def __aexit__(self, *args: object) -> None:
+        """Async context manager exit."""
+        await self.close_async()
+
+
+class SegmentationResultSource(ISegmentationResultSource):
+    """
+    High-level segmentation result source that reads from any IFrameSource.
+
+    Wraps IFrameSource to provide iterator-based access to SegmentationFrame objects.
+
+    Thread-safe: No (caller must synchronize)
+    """
+
+    def __init__(self, frame_source: IFrameSource) -> None:
+        """
+        Create a segmentation result source.
+
+        Args:
+            frame_source: Low-level frame source to read from.
+        """
+        self._frame_source = frame_source
+        self._closed = False
+
+    def read_frames(self) -> Iterator[SegmentationFrame]:
+        """
+        Read frames synchronously as an iterator.
+
+        Yields:
+            SegmentationFrame for each frame in the source.
+
+        Raises:
+            RuntimeError: If source is closed.
+        """
+        if self._closed:
+            raise RuntimeError("SegmentationResultSource is closed")
+
+        while True:
+            frame_data = self._frame_source.read_frame()
+            if frame_data is None or len(frame_data) == 0:
+                break
+
+            frame = SegmentationProtocol.read(frame_data)
+            yield frame
+
+    async def read_frames_async(self) -> AsyncIterator[SegmentationFrame]:
+        """
+        Read frames asynchronously as an async iterator.
+
+        Yields:
+            SegmentationFrame for each frame in the source.
+
+        Raises:
+            RuntimeError: If source is closed.
+        """
+        if self._closed:
+            raise RuntimeError("SegmentationResultSource is closed")
+
+        while True:
+            frame_data = await self._frame_source.read_frame_async()
+            if frame_data is None or len(frame_data) == 0:
+                break
+
+            frame = SegmentationProtocol.read(frame_data)
+            yield frame
+
+    def close(self) -> None:
+        """Close the source and release resources."""
+        if self._closed:
+            return
+        self._closed = True
+        self._frame_source.close()
+
+    async def close_async(self) -> None:
+        """Close the source and release resources asynchronously."""
+        if self._closed:
+            return
+        self._closed = True
+        await self._frame_source.close_async()
+
+
+class SegmentationProtocol:
+    """
+    Static helpers for encoding and decoding segmentation protocol data.
+
+    Pure protocol logic with no transport or rendering dependencies.
+    Matches C# SegmentationProtocol static class.
+
+    Frame Format:
+        [FrameId: 8 bytes, little-endian uint64]
+        [Width: varint]
+        [Height: varint]
+        [Instances...]
+
+    Instance Format:
+        [ClassId: 1 byte]
+        [InstanceId: 1 byte]
+        [PointCount: varint]
+        [Point0: X zigzag-varint, Y zigzag-varint]  (absolute)
+        [Point1+: deltaX zigzag-varint, deltaY zigzag-varint]
+    """
+
+    @staticmethod
+    def write(buffer: bytearray, frame: SegmentationFrame) -> int:
+        """
+        Write a complete segmentation frame to a buffer.
+
+        Args:
+            buffer: Pre-allocated buffer to write to.
+            frame: Frame to encode.
+
+        Returns:
+            Number of bytes written.
+        """
+        stream = io.BytesIO()
+
+        # Write header
+        stream.write(struct.pack("<Q", frame.frame_id))
+        _write_varint(stream, frame.width)
+        _write_varint(stream, frame.height)
+
+        # Write instances
+        for instance in frame.instances:
+            SegmentationProtocol._write_instance_core(stream, instance)
+
+        data = stream.getvalue()
+        buffer[: len(data)] = data
+        return len(data)
+
+    @staticmethod
+    def _write_instance_core(stream: BinaryIO, instance: SegmentationInstance) -> None:
+        """Write a single instance to the stream."""
+        stream.write(bytes([instance.class_id, instance.instance_id]))
+        point_count = len(instance.points)
+        _write_varint(stream, point_count)
+
+        if point_count == 0:
+            return
+
+        prev_x, prev_y = 0, 0
+        for i, point in enumerate(instance.points):
+            x, y = int(point[0]), int(point[1])
+            if i == 0:
+                # First point is absolute (but still zigzag encoded)
+                _write_varint(stream, _zigzag_encode(x))
+                _write_varint(stream, _zigzag_encode(y))
+            else:
+                # Subsequent points are deltas
+                _write_varint(stream, _zigzag_encode(x - prev_x))
+                _write_varint(stream, _zigzag_encode(y - prev_y))
+            prev_x, prev_y = x, y
+
+    @staticmethod
+    def write_header(buffer: bytearray, frame_id: int, width: int, height: int) -> int:
+        """
+        Write just the frame header (frameId, width, height).
+
+        Args:
+            buffer: Pre-allocated buffer to write to.
+            frame_id: Frame identifier.
+            width: Frame width.
+            height: Frame height.
+
+        Returns:
+            Number of bytes written.
+        """
+        stream = io.BytesIO()
+        stream.write(struct.pack("<Q", frame_id))
+        _write_varint(stream, width)
+        _write_varint(stream, height)
+        data = stream.getvalue()
+        buffer[: len(data)] = data
+        return len(data)
+
+    @staticmethod
+    def write_instance(
+        buffer: bytearray,
+        class_id: int,
+        instance_id: int,
+        points: Union[List[Point], PointArray],
+    ) -> int:
+        """
+        Write a single segmentation instance.
+
+        Points are delta-encoded for compression.
+
+        Args:
+            buffer: Pre-allocated buffer to write to.
+            class_id: Class identifier (0-255).
+            instance_id: Instance identifier (0-255).
+            points: Polygon points.
+
+        Returns:
+            Number of bytes written.
+        """
+        # Convert to numpy array if needed
+        if not isinstance(points, np.ndarray):
+            points_array = np.array(points, dtype=np.int32)
+        else:
+            points_array = points.astype(np.int32)
+
+        instance = SegmentationInstance(class_id, instance_id, points_array)
+
+        stream = io.BytesIO()
+        SegmentationProtocol._write_instance_core(stream, instance)
+        data = stream.getvalue()
+        buffer[: len(data)] = data
+        return len(data)
+
+    @staticmethod
+    def calculate_instance_size(point_count: int) -> int:
+        """
+        Calculate the maximum buffer size needed for an instance.
+
+        Args:
+            point_count: Number of polygon points.
+
+        Returns:
+            Maximum bytes needed.
+        """
+        # classId(1) + instanceId(1) + pointCount(varint, max 5) + points(max 10 bytes each)
+        return 1 + 1 + 5 + (point_count * 10)
+
+    @staticmethod
+    def read(data: bytes) -> SegmentationFrame:
+        """
+        Read a complete segmentation frame from a buffer.
+
+        Args:
+            data: Raw frame data.
+
+        Returns:
+            Decoded SegmentationFrame.
+        """
+        stream = io.BytesIO(data)
+
+        # Read header
+        frame_id_bytes = stream.read(8)
+        if len(frame_id_bytes) != 8:
+            raise EOFError("Failed to read FrameId")
+        frame_id = struct.unpack("<Q", frame_id_bytes)[0]
+        width = _read_varint(stream)
+        height = _read_varint(stream)
+
+        # Read all instances
+        instances: List[SegmentationInstance] = []
+        while True:
+            header = stream.read(2)
+            if len(header) == 0:
+                break
+            if len(header) != 2:
+                raise EOFError("Unexpected end of stream reading instance header")
+
+            class_id = header[0]
+            instance_id = header[1]
+            point_count = _read_varint(stream)
+
+            if point_count == 0:
+                points = np.empty((0, 2), dtype=np.int32)
+            else:
+                points = np.empty((point_count, 2), dtype=np.int32)
+                prev_x, prev_y = 0, 0
+
+                for i in range(point_count):
+                    x = _zigzag_decode(_read_varint(stream))
+                    y = _zigzag_decode(_read_varint(stream))
+                    if i > 0:
+                        x += prev_x
+                        y += prev_y
+                    points[i] = [x, y]
+                    prev_x, prev_y = x, y
+
+            instances.append(SegmentationInstance(class_id, instance_id, points))
+
+        return SegmentationFrame(frame_id, width, height, instances)
+
+    @staticmethod
+    def try_read(data: bytes) -> Optional[SegmentationFrame]:
+        """
+        Try to read a segmentation frame, returning None if the data is invalid.
+
+        Args:
+            data: Raw frame data.
+
+        Returns:
+            SegmentationFrame if successful, None otherwise.
+        """
+        try:
+            return SegmentationProtocol.read(data)
+        except Exception:
+            return None