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

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

rocket_welder_sdk/session_id.py

@@ -1,4 +1,4 @@
-"""SessionId parsing utilities for NNG URL generation.
+"""SessionId parsing utilities.
 
 SessionId format: ps-{guid} (e.g., ps-a1b2c3d4-e5f6-7890-abcd-ef1234567890)
 Prefix "ps" = PipelineSession.
@@ -6,33 +6,20 @@ Prefix "ps" = PipelineSession.
 This module provides utilities to:
 1. Parse SessionId from environment variable
 2. Extract the Guid portion
-3. Generate NNG IPC URLs for streaming results
-4. Read explicit NNG URLs from environment variables (preferred)
 
-## URL Configuration Priority
+## URL Configuration
 
-The SDK supports two ways to configure NNG URLs:
-
-1. **Explicit URLs (PREFERRED)** - Set by rocket-welder2:
-   - SEGMENTATION_SINK_URL
-   - KEYPOINTS_SINK_URL
-   - ACTIONS_SINK_URL
-
-2. **Derived from SessionId (FALLBACK)** - For backwards compatibility:
-   - SessionId env var → parse GUID → generate URLs
-
-Use `get_nng_urls_from_env()` for explicit URLs (preferred).
-Use `get_nng_urls(session_id)` for SessionId-derived URLs (fallback).
+Output URLs are configured via environment variables set by rocket-welder2:
+- 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)
+- ACTIONS_SINK_URL: URL for actions output
 """
 
 from __future__ import annotations
 
-import logging
 import os
 import uuid
 
-logger = logging.getLogger(__name__)
-
 SESSION_ID_PREFIX = "ps-"
 SESSION_ID_ENV_VAR = "SessionId"
 
@@ -40,6 +27,7 @@ SESSION_ID_ENV_VAR = "SessionId"
 SEGMENTATION_SINK_URL_ENV = "SEGMENTATION_SINK_URL"
 KEYPOINTS_SINK_URL_ENV = "KEYPOINTS_SINK_URL"
 ACTIONS_SINK_URL_ENV = "ACTIONS_SINK_URL"
+GRAPHICS_SINK_URL_ENV = "GRAPHICS_SINK_URL"
 
 
 def parse_session_id(session_id: str) -> uuid.UUID:
@@ -73,166 +61,3 @@ def get_session_id_from_env() -> str | None:
         SessionId string or None if not set
     """
     return os.environ.get(SESSION_ID_ENV_VAR)
-
-
-def get_nng_urls(session_id: str) -> dict[str, str]:
-    """Generate NNG IPC URLs from SessionId.
-
-    Args:
-        session_id: SessionId string (e.g., "ps-a1b2c3d4-...")
-
-    Returns:
-        Dictionary with 'segmentation', 'keypoints', 'actions' URLs
-
-    Examples:
-        >>> urls = get_nng_urls("ps-a1b2c3d4-e5f6-7890-abcd-ef1234567890")
-        >>> urls["segmentation"]
-        'ipc:///tmp/rw-a1b2c3d4-e5f6-7890-abcd-ef1234567890-seg.sock'
-    """
-    guid = parse_session_id(session_id)
-    return {
-        "segmentation": f"ipc:///tmp/rw-{guid}-seg.sock",
-        "keypoints": f"ipc:///tmp/rw-{guid}-kp.sock",
-        "actions": f"ipc:///tmp/rw-{guid}-actions.sock",
-    }
-
-
-def get_segmentation_url(session_id: str) -> str:
-    """Get NNG URL for segmentation stream.
-
-    Args:
-        session_id: SessionId string (e.g., "ps-a1b2c3d4-...")
-
-    Returns:
-        IPC URL for segmentation stream
-    """
-    guid = parse_session_id(session_id)
-    return f"ipc:///tmp/rw-{guid}-seg.sock"
-
-
-def get_keypoints_url(session_id: str) -> str:
-    """Get NNG URL for keypoints stream.
-
-    Args:
-        session_id: SessionId string (e.g., "ps-a1b2c3d4-...")
-
-    Returns:
-        IPC URL for keypoints stream
-    """
-    guid = parse_session_id(session_id)
-    return f"ipc:///tmp/rw-{guid}-kp.sock"
-
-
-def get_actions_url(session_id: str) -> str:
-    """Get NNG URL for actions stream.
-
-    Args:
-        session_id: SessionId string (e.g., "ps-a1b2c3d4-...")
-
-    Returns:
-        IPC URL for actions stream
-    """
-    guid = parse_session_id(session_id)
-    return f"ipc:///tmp/rw-{guid}-actions.sock"
-
-
-# ============================================================================
-# Explicit URL functions (PREFERRED - URLs set by rocket-welder2)
-# ============================================================================
-
-
-def get_nng_urls_from_env() -> dict[str, str | None]:
-    """Get NNG URLs from explicit environment variables.
-
-    This is the PREFERRED method for getting NNG URLs. rocket-welder2
-    sets these environment variables when starting containers.
-
-    Returns:
-        Dictionary with 'segmentation', 'keypoints', 'actions' URLs.
-        Values are None if not configured.
-
-    Examples:
-        >>> os.environ["SEGMENTATION_SINK_URL"] = "ipc:///tmp/rw-abc-seg.sock"
-        >>> urls = get_nng_urls_from_env()
-        >>> urls["segmentation"]
-        'ipc:///tmp/rw-abc-seg.sock'
-    """
-    return {
-        "segmentation": os.environ.get(SEGMENTATION_SINK_URL_ENV),
-        "keypoints": os.environ.get(KEYPOINTS_SINK_URL_ENV),
-        "actions": os.environ.get(ACTIONS_SINK_URL_ENV),
-    }
-
-
-def get_segmentation_url_from_env() -> str | None:
-    """Get segmentation NNG URL from environment variable.
-
-    Returns:
-        IPC URL for segmentation stream, or None if not configured.
-    """
-    return os.environ.get(SEGMENTATION_SINK_URL_ENV)
-
-
-def get_keypoints_url_from_env() -> str | None:
-    """Get keypoints NNG URL from environment variable.
-
-    Returns:
-        IPC URL for keypoints stream, or None if not configured.
-    """
-    return os.environ.get(KEYPOINTS_SINK_URL_ENV)
-
-
-def get_actions_url_from_env() -> str | None:
-    """Get actions NNG URL from environment variable.
-
-    Returns:
-        IPC URL for actions stream, or None if not configured.
-    """
-    return os.environ.get(ACTIONS_SINK_URL_ENV)
-
-
-def has_explicit_nng_urls() -> bool:
-    """Check if explicit NNG URLs are configured.
-
-    Returns:
-        True if at least segmentation OR keypoints URL is configured.
-    """
-    urls = get_nng_urls_from_env()
-    return bool(urls["segmentation"] or urls["keypoints"])
-
-
-def get_configured_nng_urls() -> dict[str, str]:
-    """Get all configured NNG URLs (explicit or derived from SessionId).
-
-    Priority:
-    1. Explicit URLs from environment (SEGMENTATION_SINK_URL, etc.)
-    2. Derived from SessionId environment variable (fallback)
-
-    Returns:
-        Dictionary with 'segmentation', 'keypoints', 'actions' URLs.
-        Only includes URLs that are actually configured.
-
-    Raises:
-        ValueError: If no NNG URLs are configured (neither explicit nor SessionId).
-    """
-    # Try explicit URLs first (preferred)
-    explicit_urls = get_nng_urls_from_env()
-    result: dict[str, str] = {}
-
-    for name, url in explicit_urls.items():
-        if url:
-            result[name] = url
-
-    # If we have at least one explicit URL, return what we have
-    if result:
-        return result
-
-    # Fallback: derive from SessionId
-    session_id = get_session_id_from_env()
-    if session_id:
-        return get_nng_urls(session_id)
-
-    raise ValueError(
-        "No NNG URLs configured. Set SEGMENTATION_SINK_URL/KEYPOINTS_SINK_URL "
-        "environment variables, or set SessionId for URL derivation."
-    )

rocket_welder_sdk/transport/__init__.py

@@ -6,7 +6,6 @@ Provides transport-agnostic frame sink/source abstractions for protocols.
 
 from .frame_sink import IFrameSink, NullFrameSink
 from .frame_source import IFrameSource
-from .nng_transport import NngFrameSink, NngFrameSource
 from .stream_transport import StreamFrameSink, StreamFrameSource
 from .tcp_transport import TcpFrameSink, TcpFrameSource
 from .unix_socket_transport import (
@@ -14,12 +13,16 @@ from .unix_socket_transport import (
     UnixSocketFrameSource,
     UnixSocketServer,
 )
+from .websocket_transport import (
+    WebSocketFrameSink,
+    WebSocketFrameSource,
+    connect_websocket_sink,
+    connect_websocket_source,
+)
 
 __all__ = [
     "IFrameSink",
     "IFrameSource",
-    "NngFrameSink",
-    "NngFrameSource",
     "NullFrameSink",
     "StreamFrameSink",
     "StreamFrameSource",
@@ -28,4 +31,8 @@ __all__ = [
     "UnixSocketFrameSink",
     "UnixSocketFrameSource",
     "UnixSocketServer",
+    "WebSocketFrameSink",
+    "WebSocketFrameSource",
+    "connect_websocket_sink",
+    "connect_websocket_source",
 ]

rocket_welder_sdk/transport/frame_sink.py

@@ -9,7 +9,7 @@ class IFrameSink(ABC):
 
     Transport-agnostic interface that handles the question: "where do frames go?"
     This abstraction decouples protocol logic (KeyPoints, SegmentationResults) from
-    transport mechanisms (File, TCP, WebSocket, NNG). Each frame is written atomically.
+    transport mechanisms (File, TCP, WebSocket, Unix Socket). Each frame is written atomically.
     """
 
     @abstractmethod
@@ -37,7 +37,7 @@ class IFrameSink(ABC):
         """
         Flush any buffered data to the transport synchronously.
 
-        For message-based transports (NNG, WebSocket), this may be a no-op.
+        For message-based transports (WebSocket), this may be a no-op.
         """
         pass
 
@@ -46,7 +46,7 @@ class IFrameSink(ABC):
         """
         Flush any buffered data to the transport asynchronously.
 
-        For message-based transports (NNG, WebSocket), this may be a no-op.
+        For message-based transports (WebSocket), this may be a no-op.
         """
         pass