rocket-welder-sdk 1.1.32__py3-none-any.whl → 1.1.34__py3-none-any.whl

rocket_welder_sdk/high_level/client.py

@@ -0,0 +1,262 @@
+"""
+RocketWelderClient - High-level API matching C# RocketWelder.SDK.
+
+Usage:
+    with RocketWelderClient.from_environment() as client:
+        # Define schema
+        nose = client.keypoints.define_point("nose")
+        person = client.segmentation.define_class(1, "person")
+
+        # Start processing
+        client.start(process_frame)
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Callable, Optional
+
+import numpy as np
+import numpy.typing as npt
+from typing_extensions import TypeAlias
+
+from .connection_strings import (
+    KeyPointsConnectionString,
+    SegmentationConnectionString,
+    VideoSourceConnectionString,
+)
+from .data_context import (
+    IKeyPointsDataContext,
+    ISegmentationDataContext,
+    KeyPointsDataContext,
+    SegmentationDataContext,
+)
+from .schema import (
+    IKeyPointsSchema,
+    ISegmentationSchema,
+    KeyPointsSchema,
+    SegmentationSchema,
+)
+from .transport_protocol import TransportKind
+
+if TYPE_CHECKING:
+    from rocket_welder_sdk.keypoints_protocol import KeyPointsSink
+    from rocket_welder_sdk.transport.frame_sink import IFrameSink
+
+# Type alias for OpenCV Mat (numpy array)
+Mat: TypeAlias = npt.NDArray[np.uint8]
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class RocketWelderClientOptions:
+    """Configuration options for RocketWelderClient."""
+
+    video_source: VideoSourceConnectionString = field(
+        default_factory=VideoSourceConnectionString.default
+    )
+    keypoints: KeyPointsConnectionString = field(default_factory=KeyPointsConnectionString.default)
+    segmentation: SegmentationConnectionString = field(
+        default_factory=SegmentationConnectionString.default
+    )
+
+    @classmethod
+    def from_environment(cls) -> RocketWelderClientOptions:
+        """Create from environment variables."""
+        return cls(
+            video_source=VideoSourceConnectionString.from_environment(),
+            keypoints=KeyPointsConnectionString.from_environment(),
+            segmentation=SegmentationConnectionString.from_environment(),
+        )
+
+
+class RocketWelderClient:
+    """
+    High-level client for RocketWelder SDK.
+
+    Mirrors C# RocketWelder.SDK.IRocketWelderClient interface.
+    """
+
+    def __init__(self, options: RocketWelderClientOptions) -> None:
+        self._options = options
+        self._keypoints_schema = KeyPointsSchema()
+        self._segmentation_schema = SegmentationSchema()
+        self._keypoints_sink: Optional[KeyPointsSink] = None
+        self._keypoints_frame_sink: Optional[IFrameSink] = None
+        self._segmentation_frame_sink: Optional[IFrameSink] = None
+        self._closed = False
+        logger.debug("RocketWelderClient created with options: %s", options)
+
+    @classmethod
+    def from_environment(cls) -> RocketWelderClient:
+        """Create client from environment variables."""
+        logger.info("Creating RocketWelderClient from environment variables")
+        return cls(RocketWelderClientOptions.from_environment())
+
+    @classmethod
+    def create(cls, options: Optional[RocketWelderClientOptions] = None) -> RocketWelderClient:
+        """Create client with explicit options."""
+        return cls(options or RocketWelderClientOptions())
+
+    @property
+    def keypoints(self) -> IKeyPointsSchema:
+        """Schema for defining keypoints."""
+        return self._keypoints_schema
+
+    @property
+    def segmentation(self) -> ISegmentationSchema:
+        """Schema for defining segmentation classes."""
+        return self._segmentation_schema
+
+    def start(
+        self,
+        process_frame: Callable[[Mat, ISegmentationDataContext, IKeyPointsDataContext, Mat], None],
+    ) -> None:
+        """Start with both keypoints and segmentation."""
+        self._run_loop(process_frame, use_keypoints=True, use_segmentation=True)
+
+    def start_keypoints(
+        self,
+        process_frame: Callable[[Mat, IKeyPointsDataContext, Mat], None],
+    ) -> None:
+        """Start with keypoints only."""
+        self._run_loop(process_frame, use_keypoints=True, use_segmentation=False)
+
+    def start_segmentation(
+        self,
+        process_frame: Callable[[Mat, ISegmentationDataContext, Mat], None],
+    ) -> None:
+        """Start with segmentation only."""
+        self._run_loop(process_frame, use_keypoints=False, use_segmentation=True)
+
+    def _run_loop(
+        self,
+        process_frame: Callable[..., None],
+        use_keypoints: bool,
+        use_segmentation: bool,
+    ) -> None:
+        """Run processing loop."""
+        from rocket_welder_sdk.keypoints_protocol import KeyPointsSink
+
+        logger.info(
+            "Starting processing loop (keypoints=%s, segmentation=%s)",
+            use_keypoints,
+            use_segmentation,
+        )
+
+        # Initialize sinks
+        if use_keypoints:
+            cs = self._options.keypoints
+            logger.info("Initializing keypoints sink: %s -> %s", cs.protocol, cs.address)
+            self._keypoints_frame_sink = self._create_frame_sink(cs.protocol, cs.address)
+            self._keypoints_sink = KeyPointsSink(
+                frame_sink=self._keypoints_frame_sink,
+                master_frame_interval=cs.master_frame_interval,
+                owns_sink=False,  # We manage frame sink lifecycle in close()
+            )
+            logger.debug(
+                "KeyPointsSink created with master_frame_interval=%d", cs.master_frame_interval
+            )
+
+        if use_segmentation:
+            seg_cs = self._options.segmentation
+            logger.info("Initializing segmentation sink: %s -> %s", seg_cs.protocol, seg_cs.address)
+            self._segmentation_frame_sink = self._create_frame_sink(seg_cs.protocol, seg_cs.address)
+            logger.debug("Segmentation frame sink created")
+
+        # TODO: Video capture loop - for now raise NotImplementedError
+        raise NotImplementedError(
+            "Video capture not implemented. Use process_frame_sync() or low-level API."
+        )
+
+    def process_frame_sync(
+        self,
+        frame_id: int,
+        input_frame: Mat,
+        output_frame: Mat,
+        width: int,
+        height: int,
+    ) -> tuple[Optional[IKeyPointsDataContext], Optional[ISegmentationDataContext]]:
+        """
+        Process a single frame synchronously.
+
+        Returns (keypoints_context, segmentation_context) for the caller to use.
+        Caller must call commit() on contexts when done.
+        """
+        from rocket_welder_sdk.segmentation_result import SegmentationResultWriter
+
+        kp_ctx: Optional[IKeyPointsDataContext] = None
+        seg_ctx: Optional[ISegmentationDataContext] = None
+
+        if self._keypoints_sink is not None:
+            kp_writer = self._keypoints_sink.create_writer(frame_id)
+            kp_ctx = KeyPointsDataContext(frame_id, kp_writer)
+
+        if self._segmentation_frame_sink is not None:
+            seg_writer = SegmentationResultWriter(
+                frame_id, width, height, frame_sink=self._segmentation_frame_sink
+            )
+            seg_ctx = SegmentationDataContext(frame_id, seg_writer)
+
+        return kp_ctx, seg_ctx
+
+    def _create_frame_sink(self, protocol: Any, address: str) -> IFrameSink:
+        """Create frame sink from protocol."""
+        from rocket_welder_sdk.transport import NngFrameSink
+        from rocket_welder_sdk.transport.stream_transport import StreamFrameSink
+        from rocket_welder_sdk.transport.unix_socket_transport import UnixSocketFrameSink
+
+        from .transport_protocol import TransportProtocol
+
+        if not isinstance(protocol, TransportProtocol):
+            raise TypeError(f"Expected TransportProtocol, got {type(protocol)}")
+
+        if protocol.kind == TransportKind.FILE:
+            logger.debug("Creating file sink: %s", address)
+            file_handle = open(address, "wb")
+            try:
+                return StreamFrameSink(file_handle)
+            except Exception:
+                file_handle.close()
+                raise
+        elif protocol.kind == TransportKind.SOCKET:
+            logger.debug("Creating Unix socket sink: %s", address)
+            return UnixSocketFrameSink.connect(address)
+        elif protocol.kind in (TransportKind.NNG_PUSH_IPC, TransportKind.NNG_PUSH_TCP):
+            logger.debug("Creating NNG pusher: %s", address)
+            return NngFrameSink.create_pusher(address)
+        elif protocol.kind in (TransportKind.NNG_PUB_IPC, TransportKind.NNG_PUB_TCP):
+            logger.debug("Creating NNG publisher: %s", address)
+            return NngFrameSink.create_publisher(address)
+        else:
+            raise ValueError(f"Unsupported protocol: {protocol}")
+
+    def close(self) -> None:
+        """Release resources."""
+        if self._closed:
+            return
+
+        logger.info("Closing RocketWelderClient")
+
+        # Close frame sinks (KeyPointsSink has owns_sink=False, so we manage lifecycle)
+        self._keypoints_sink = None
+        if self._keypoints_frame_sink is not None:
+            logger.debug("Closing keypoints frame sink")
+            self._keypoints_frame_sink.close()
+            self._keypoints_frame_sink = None
+
+        if self._segmentation_frame_sink is not None:
+            logger.debug("Closing segmentation frame sink")
+            self._segmentation_frame_sink.close()
+            self._segmentation_frame_sink = None
+
+        self._closed = True
+        logger.info("RocketWelderClient closed")
+
+    def __enter__(self) -> RocketWelderClient:
+        return self
+
+    def __exit__(self, *args: object) -> None:
+        self.close()

rocket_welder_sdk/high_level/connection_strings.py

@@ -0,0 +1,331 @@
+"""
+Strongly-typed connection strings with parsing support.
+
+Connection string format: protocol://path?param1=value1&param2=value2
+
+Examples:
+    nng+push+ipc://tmp/keypoints?masterFrameInterval=300
+    nng+pub+tcp://localhost:5555
+    file:///path/to/output.bin
+    socket:///tmp/my.sock
+"""
+
+from __future__ import annotations
+
+import contextlib
+import os
+from dataclasses import dataclass, field
+from enum import Enum, auto
+from typing import Dict, Optional
+from urllib.parse import parse_qs
+
+from .transport_protocol import TransportProtocol
+
+
+class VideoSourceType(Enum):
+    """Type of video source."""
+
+    CAMERA = auto()
+    FILE = auto()
+    SHARED_MEMORY = auto()
+    RTSP = auto()
+    HTTP = auto()
+
+
+@dataclass(frozen=True)
+class VideoSourceConnectionString:
+    """
+    Strongly-typed connection string for video source input.
+
+    Supported formats:
+    - "0", "1", etc. - Camera device index
+    - file://path/to/video.mp4 - Video file
+    - shm://buffer_name - Shared memory buffer
+    - rtsp://host/stream - RTSP stream
+    """
+
+    value: str
+    source_type: VideoSourceType
+    camera_index: Optional[int] = None
+    path: Optional[str] = None
+    parameters: Dict[str, str] = field(default_factory=dict)
+
+    @classmethod
+    def default(cls) -> VideoSourceConnectionString:
+        """Default video source (camera 0)."""
+        return cls.parse("0")
+
+    @classmethod
+    def from_environment(cls, variable_name: str = "VIDEO_SOURCE") -> VideoSourceConnectionString:
+        """Create from environment variable or use default."""
+        value = os.environ.get(variable_name) or os.environ.get("CONNECTION_STRING")
+        return cls.parse(value) if value else cls.default()
+
+    @classmethod
+    def parse(cls, s: str) -> VideoSourceConnectionString:
+        """Parse a connection string."""
+        result = cls.try_parse(s)
+        if result is None:
+            raise ValueError(f"Invalid video source connection string: {s}")
+        return result
+
+    @classmethod
+    def try_parse(cls, s: str) -> Optional[VideoSourceConnectionString]:
+        """Try to parse a connection string."""
+        if not s or not s.strip():
+            return None
+
+        s = s.strip()
+        parameters: Dict[str, str] = {}
+
+        # Extract query parameters
+        if "?" in s:
+            base, query = s.split("?", 1)
+            for key, values in parse_qs(query).items():
+                parameters[key.lower()] = values[0] if values else ""
+            s = base
+
+        # Check for camera index first
+        if s.isdigit():
+            return cls(
+                value=s,
+                source_type=VideoSourceType.CAMERA,
+                camera_index=int(s),
+                parameters=parameters,
+            )
+
+        # Parse protocol
+        if s.startswith("file://"):
+            path = "/" + s[7:]  # Restore absolute path
+            return cls(
+                value=s,
+                source_type=VideoSourceType.FILE,
+                path=path,
+                parameters=parameters,
+            )
+        elif s.startswith("shm://"):
+            path = s[6:]
+            return cls(
+                value=s,
+                source_type=VideoSourceType.SHARED_MEMORY,
+                path=path,
+                parameters=parameters,
+            )
+        elif s.startswith("rtsp://"):
+            return cls(
+                value=s,
+                source_type=VideoSourceType.RTSP,
+                path=s,
+                parameters=parameters,
+            )
+        elif s.startswith("http://") or s.startswith("https://"):
+            return cls(
+                value=s,
+                source_type=VideoSourceType.HTTP,
+                path=s,
+                parameters=parameters,
+            )
+        elif "://" not in s:
+            # Assume file path
+            return cls(
+                value=s,
+                source_type=VideoSourceType.FILE,
+                path=s,
+                parameters=parameters,
+            )
+
+        return None
+
+    def __str__(self) -> str:
+        return self.value
+
+
+@dataclass(frozen=True)
+class KeyPointsConnectionString:
+    """
+    Strongly-typed connection string for KeyPoints output.
+
+    Supported protocols:
+    - file:///path/to/file.bin - File output (absolute path)
+    - socket:///tmp/socket.sock - Unix domain socket
+    - nng+push+ipc://tmp/keypoints - NNG Push over IPC
+    - nng+push+tcp://host:port - NNG Push over TCP
+
+    Supported parameters:
+    - masterFrameInterval: Interval between master frames (default: 300)
+    """
+
+    value: str
+    protocol: TransportProtocol
+    address: str
+    master_frame_interval: int = 300
+    parameters: Dict[str, str] = field(default_factory=dict)
+
+    @classmethod
+    def default(cls) -> KeyPointsConnectionString:
+        """Default connection string for KeyPoints."""
+        return cls.parse("nng+push+ipc://tmp/rocket-welder-keypoints?masterFrameInterval=300")
+
+    @classmethod
+    def from_environment(
+        cls, variable_name: str = "KEYPOINTS_CONNECTION_STRING"
+    ) -> KeyPointsConnectionString:
+        """Create from environment variable or use default."""
+        value = os.environ.get(variable_name)
+        return cls.parse(value) if value else cls.default()
+
+    @classmethod
+    def parse(cls, s: str) -> KeyPointsConnectionString:
+        """Parse a connection string."""
+        result = cls.try_parse(s)
+        if result is None:
+            raise ValueError(f"Invalid KeyPoints connection string: {s}")
+        return result
+
+    @classmethod
+    def try_parse(cls, s: str) -> Optional[KeyPointsConnectionString]:
+        """Try to parse a connection string."""
+        if not s or not s.strip():
+            return None
+
+        s = s.strip()
+        parameters: Dict[str, str] = {}
+
+        # Extract query parameters
+        endpoint_part = s
+        if "?" in s:
+            endpoint_part, query = s.split("?", 1)
+            for key, values in parse_qs(query).items():
+                parameters[key.lower()] = values[0] if values else ""
+
+        # Parse protocol and address
+        scheme_end = endpoint_part.find("://")
+        if scheme_end <= 0:
+            return None
+
+        schema_str = endpoint_part[:scheme_end]
+        path_part = endpoint_part[scheme_end + 3 :]  # skip "://"
+
+        protocol = TransportProtocol.try_parse(schema_str)
+        if protocol is None:
+            return None
+
+        # Build address based on protocol type
+        if protocol.is_file:
+            # file:///absolute/path -> /absolute/path
+            address = path_part if path_part.startswith("/") else "/" + path_part
+        elif protocol.is_socket:
+            # socket:///tmp/sock -> /tmp/sock
+            address = path_part if path_part.startswith("/") else "/" + path_part
+        elif protocol.is_nng:
+            # NNG protocols need proper address format
+            address = protocol.create_nng_address(path_part)
+        else:
+            return None
+
+        # Parse masterFrameInterval
+        master_frame_interval = 300  # default
+        if "masterframeinterval" in parameters:
+            with contextlib.suppress(ValueError):
+                master_frame_interval = int(parameters["masterframeinterval"])
+
+        return cls(
+            value=s,
+            protocol=protocol,
+            address=address,
+            master_frame_interval=master_frame_interval,
+            parameters=parameters,
+        )
+
+    def __str__(self) -> str:
+        return self.value
+
+
+@dataclass(frozen=True)
+class SegmentationConnectionString:
+    """
+    Strongly-typed connection string for Segmentation output.
+
+    Supported protocols:
+    - file:///path/to/file.bin - File output (absolute path)
+    - socket:///tmp/socket.sock - Unix domain socket
+    - nng+push+ipc://tmp/segmentation - NNG Push over IPC
+    - nng+push+tcp://host:port - NNG Push over TCP
+    """
+
+    value: str
+    protocol: TransportProtocol
+    address: str
+    parameters: Dict[str, str] = field(default_factory=dict)
+
+    @classmethod
+    def default(cls) -> SegmentationConnectionString:
+        """Default connection string for Segmentation."""
+        return cls.parse("nng+push+ipc://tmp/rocket-welder-segmentation")
+
+    @classmethod
+    def from_environment(
+        cls, variable_name: str = "SEGMENTATION_CONNECTION_STRING"
+    ) -> SegmentationConnectionString:
+        """Create from environment variable or use default."""
+        value = os.environ.get(variable_name)
+        return cls.parse(value) if value else cls.default()
+
+    @classmethod
+    def parse(cls, s: str) -> SegmentationConnectionString:
+        """Parse a connection string."""
+        result = cls.try_parse(s)
+        if result is None:
+            raise ValueError(f"Invalid Segmentation connection string: {s}")
+        return result
+
+    @classmethod
+    def try_parse(cls, s: str) -> Optional[SegmentationConnectionString]:
+        """Try to parse a connection string."""
+        if not s or not s.strip():
+            return None
+
+        s = s.strip()
+        parameters: Dict[str, str] = {}
+
+        # Extract query parameters
+        endpoint_part = s
+        if "?" in s:
+            endpoint_part, query = s.split("?", 1)
+            for key, values in parse_qs(query).items():
+                parameters[key.lower()] = values[0] if values else ""
+
+        # Parse protocol and address
+        scheme_end = endpoint_part.find("://")
+        if scheme_end <= 0:
+            return None
+
+        schema_str = endpoint_part[:scheme_end]
+        path_part = endpoint_part[scheme_end + 3 :]  # skip "://"
+
+        protocol = TransportProtocol.try_parse(schema_str)
+        if protocol is None:
+            return None
+
+        # Build address based on protocol type
+        if protocol.is_file:
+            # file:///absolute/path -> /absolute/path
+            address = path_part if path_part.startswith("/") else "/" + path_part
+        elif protocol.is_socket:
+            # socket:///tmp/sock -> /tmp/sock
+            address = path_part if path_part.startswith("/") else "/" + path_part
+        elif protocol.is_nng:
+            # NNG protocols need proper address format
+            address = protocol.create_nng_address(path_part)
+        else:
+            return None
+
+        return cls(
+            value=s,
+            protocol=protocol,
+            address=address,
+            parameters=parameters,
+        )
+
+    def __str__(self) -> str:
+        return self.value