dexcontrol 0.2.1__py3-none-any.whl

dexcontrol/utils/subscribers/base.py

@@ -0,0 +1,260 @@
+# Copyright (c) 2025 Dexmate CORPORATION & AFFILIATES. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 with Commons Clause License
+# Condition v1.0 [see LICENSE for details].
+
+"""Base Zenoh subscriber utilities.
+
+This module provides the abstract base class for all Zenoh subscribers
+and common utilities used across different subscriber implementations.
+"""
+
+import threading
+import time
+from abc import ABC, abstractmethod
+from collections.abc import Callable
+from typing import Any, Final, TypeVar
+
+import zenoh
+from google.protobuf.message import Message
+from loguru import logger
+
+from dexcontrol.utils.os_utils import resolve_key_name
+
+# Type variable for Message subclasses
+M = TypeVar("M", bound=Message)
+
+# Type alias for custom data handler functions
+CustomDataHandler = Callable[[zenoh.Sample], None]
+
+
+class BaseZenohSubscriber(ABC):
+    """Base class for Zenoh subscribers with configurable data handling.
+
+    This class provides a common interface for subscribing to Zenoh topics
+    and processing incoming data through configurable decoder functions.
+    It handles the Zenoh communication setup and provides thread-safe access
+    to the latest data.
+
+    Attributes:
+        _active: Whether subscriber is receiving data updates.
+        _data_lock: Lock for thread-safe data access.
+        _zenoh_session: Active Zenoh session for communication.
+        _subscriber: Zenoh subscriber for data.
+        _topic: The resolved Zenoh topic name.
+        _enable_fps_tracking: Whether to track and log FPS metrics.
+        _frame_count: Counter for frames processed since last FPS calculation.
+        _fps: Most recently calculated frames per second value.
+        _last_fps_time: Timestamp of last FPS calculation.
+        _fps_log_interval: Number of frames between FPS calculations.
+        _name: Name for logging purposes.
+        _custom_data_handler: Optional custom data handler function.
+        _last_data_time: Timestamp of last data received.
+    """
+
+    def __init__(
+        self,
+        topic: str,
+        zenoh_session: zenoh.Session,
+        name: str = "subscriber",
+        enable_fps_tracking: bool = False,
+        fps_log_interval: int = 100,
+        custom_data_handler: CustomDataHandler | None = None,
+    ) -> None:
+        """Initialize the base Zenoh subscriber.
+
+        Args:
+            topic: Zenoh topic to subscribe to for data.
+            zenoh_session: Active Zenoh session for communication.
+            name: Name for logging purposes.
+            enable_fps_tracking: Whether to track and log FPS metrics.
+            fps_log_interval: Number of frames between FPS calculations.
+            custom_data_handler: Optional custom function to handle incoming data.
+                                If provided, this will replace the default data
+                                handling logic entirely.
+        """
+        self._active: bool = False
+        self._data_lock: Final[threading.RLock] = threading.RLock()
+        self._zenoh_session: Final[zenoh.Session] = zenoh_session
+        self._name = name
+        self._custom_data_handler = custom_data_handler
+
+        # Data freshness tracking
+        self._last_data_time: float | None = None
+
+        # FPS tracking
+        self._enable_fps_tracking = enable_fps_tracking
+        self._fps_log_interval = fps_log_interval
+        self._frame_count = 0
+        self._fps = 0.0
+        self._last_fps_time = time.time()
+
+        # Setup Zenoh subscriber
+        self._topic: Final[str] = resolve_key_name(topic)
+        self._subscriber: Final[zenoh.Subscriber] = zenoh_session.declare_subscriber(
+            self._topic, self._data_handler_wrapper
+        )
+
+        logger.info(f"Initialized {self._name} subscriber for topic: {self._topic}")
+
+    def _data_handler_wrapper(self, sample: zenoh.Sample) -> None:
+        """Wrapper for data handling that calls either custom or default handler.
+
+        Args:
+            sample: Zenoh sample containing data.
+        """
+        # Update data freshness timestamp
+        with self._data_lock:
+            self._last_data_time = time.monotonic()
+
+        # Call custom data handler if provided, otherwise call default handler
+        if self._custom_data_handler is not None:
+            try:
+                self._custom_data_handler(sample)
+            except Exception as e:
+                logger.error(f"Custom data handler failed for {self._name}: {e}")
+        else:
+            # Call the default data handler
+            self._data_handler(sample)
+
+    @abstractmethod
+    def _data_handler(self, sample: zenoh.Sample) -> None:
+        """Handle incoming data.
+
+        This method must be implemented by subclasses to process the specific
+        type of data they handle.
+
+        Args:
+            sample: Zenoh sample containing data.
+        """
+        pass
+
+    @abstractmethod
+    def get_latest_data(self) -> Any | None:
+        """Get the latest data.
+
+        Returns:
+            Latest data if available, None otherwise.
+        """
+        pass
+
+    def _update_fps_metrics(self) -> None:
+        """Update FPS tracking metrics.
+
+        Increments frame counter and recalculates FPS at specified intervals.
+        Only has an effect if fps_tracking was enabled during initialization.
+        """
+        if not self._enable_fps_tracking:
+            return
+
+        self._frame_count += 1
+        if self._frame_count >= self._fps_log_interval:
+            current_time = time.time()
+            elapsed = current_time - self._last_fps_time
+            self._fps = self._frame_count / elapsed
+            logger.info(f"{self._name} {self._topic} frequency: {self._fps:.2f} Hz")
+            self._frame_count = 0
+            self._last_fps_time = current_time
+
+    def wait_for_active(self, timeout: float = 5.0) -> bool:
+        """Wait for the subscriber to start receiving data.
+
+        Args:
+            timeout: Maximum time to wait in seconds.
+
+        Returns:
+            True if subscriber becomes active, False if timeout is reached.
+        """
+        start_time = time.monotonic()
+        check_interval = min(0.05, timeout / 10)  # Check every 10ms or less
+
+        while True:
+            with self._data_lock:
+                if self._active:
+                    return True
+
+            elapsed = time.monotonic() - start_time
+            if elapsed >= timeout:
+                logger.error(f"No data received from {self._topic} after {timeout}s")
+                return False
+
+            # Sleep for the shorter of: remaining time or check interval
+            sleep_time = min(check_interval, timeout - elapsed)
+            time.sleep(sleep_time)
+
+    def is_active(self) -> bool:
+        """Check if the subscriber is actively receiving data.
+
+        Returns:
+            True if subscriber is active, False otherwise.
+        """
+        with self._data_lock:
+            return self._active
+
+    def shutdown(self) -> None:
+        """Stop the subscriber and release resources."""
+        if hasattr(self, "_subscriber") and self._subscriber:
+            self._subscriber.undeclare()
+
+    @property
+    def topic(self) -> str:
+        """Get the Zenoh topic name.
+
+        Returns:
+            The resolved Zenoh topic name.
+        """
+        return self._topic
+
+    @property
+    def fps(self) -> float:
+        """Get the current FPS measurement.
+
+        Returns:
+            Current frames per second measurement.
+        """
+        return self._fps
+
+    @property
+    def name(self) -> str:
+        """Get the subscriber name.
+
+        Returns:
+            The subscriber name.
+        """
+        return self._name
+
+    def is_data_fresh(self, max_age_seconds: float) -> bool:
+        """Check if the most recent data is fresh (received within the specified time).
+
+        This method checks if data has been received within the specified time window,
+        regardless of whether the data content has changed. This is useful for
+        detecting communication failures or stale data streams.
+
+        Args:
+            max_age_seconds: Maximum age in seconds for data to be considered fresh.
+
+        Returns:
+            True if fresh data was received within the time window, False otherwise.
+            Returns False if no data has ever been received.
+        """
+        with self._data_lock:
+            if self._last_data_time is None:
+                return False
+
+            current_time = time.monotonic()
+            age = current_time - self._last_data_time
+            return age <= max_age_seconds
+
+    def get_time_since_last_data(self) -> float | None:
+        """Get the time elapsed since the last data was received.
+
+        Returns:
+            Time in seconds since last data was received, or None if no data
+            has ever been received.
+        """
+        with self._data_lock:
+            if self._last_data_time is None:
+                return None
+
+            current_time = time.monotonic()
+            return current_time - self._last_data_time

dexcontrol/utils/subscribers/camera.py

@@ -0,0 +1,328 @@
+# Copyright (c) 2025 Dexmate CORPORATION & AFFILIATES. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 with Commons Clause License
+# Condition v1.0 [see LICENSE for details].
+
+"""Camera Zenoh subscribers for RGB and depth data.
+
+This module provides specialized subscribers for camera data including RGB images
+and depth images, using the serialization formats from dexsensor.
+"""
+
+import numpy as np
+import zenoh
+from loguru import logger
+
+from .base import BaseZenohSubscriber, CustomDataHandler
+
+# Import camera serialization functions from dexsensor
+try:
+    from dexsensor.serialization.camera import decode_depth, decode_image
+except ImportError:
+    logger.error(
+        "Failed to import dexsensor camera serialization functions. Please install dexsensor."
+    )
+    decode_image = None
+    decode_depth = None
+
+
+class RGBCameraSubscriber(BaseZenohSubscriber):
+    """Zenoh subscriber for RGB camera data.
+
+    This subscriber handles RGB image data encoded using the dexsensor
+    camera serialization format with JPEG compression.
+    Uses lazy decoding - data is only decoded when requested.
+    """
+
+    def __init__(
+        self,
+        topic: str,
+        zenoh_session: zenoh.Session,
+        name: str = "rgb_camera_subscriber",
+        enable_fps_tracking: bool = True,
+        fps_log_interval: int = 30,
+        custom_data_handler: CustomDataHandler | None = None,
+    ) -> None:
+        """Initialize the RGB camera subscriber.
+
+        Args:
+            topic: Zenoh topic to subscribe to for RGB data.
+            zenoh_session: Active Zenoh session for communication.
+            name: Name for logging purposes.
+            enable_fps_tracking: Whether to track and log FPS metrics.
+            fps_log_interval: Number of frames between FPS calculations.
+            custom_data_handler: Optional custom function to handle incoming data.
+                                If provided, this will replace the default data
+                                handling logic entirely.
+        """
+        super().__init__(
+            topic,
+            zenoh_session,
+            name,
+            enable_fps_tracking,
+            fps_log_interval,
+            custom_data_handler,
+        )
+        self._latest_raw_data: bytes | None = None
+
+    def _data_handler(self, sample: zenoh.Sample) -> None:
+        """Handle incoming RGB image data.
+
+        Args:
+            sample: Zenoh sample containing encoded RGB image data.
+        """
+        with self._data_lock:
+            self._latest_raw_data = sample.payload.to_bytes()
+            self._active = True
+
+        self._update_fps_metrics()
+
+    def get_latest_data(self) -> np.ndarray | None:
+        """Get the latest RGB image.
+
+        Returns:
+            Latest RGB image as numpy array (HxWxC) if available, None otherwise.
+        """
+        with self._data_lock:
+            if self._latest_raw_data is None:
+                return None
+
+            if decode_image is None:
+                logger.error(
+                    f"Cannot decode RGB image for {self._name}: dexsensor not available"
+                )
+                return None
+
+            try:
+                # Decode the RGB image
+                image = decode_image(self._latest_raw_data)
+                return image
+            except Exception as e:
+                logger.error(f"Failed to decode RGB image for {self._name}: {e}")
+                return None
+
+    def get_latest_image(self) -> np.ndarray | None:
+        """Get the latest RGB image.
+
+        Alias for get_latest_data() for clarity.
+
+        Returns:
+            Latest RGB image as numpy array (HxWxC) if available, None otherwise.
+        """
+        return self.get_latest_data()
+
+
+class DepthCameraSubscriber(BaseZenohSubscriber):
+    """Zenoh subscriber for depth camera data.
+
+    This subscriber handles depth image data encoded using the dexsensor
+    camera serialization format with compression.
+    Uses lazy decoding - data is only decoded when requested.
+    """
+
+    def __init__(
+        self,
+        topic: str,
+        zenoh_session: zenoh.Session,
+        name: str = "depth_camera_subscriber",
+        enable_fps_tracking: bool = True,
+        fps_log_interval: int = 30,
+        custom_data_handler: CustomDataHandler | None = None,
+    ) -> None:
+        """Initialize the depth camera subscriber.
+
+        Args:
+            topic: Zenoh topic to subscribe to for depth data.
+            zenoh_session: Active Zenoh session for communication.
+            name: Name for logging purposes.
+            enable_fps_tracking: Whether to track and log FPS metrics.
+            fps_log_interval: Number of frames between FPS calculations.
+            custom_data_handler: Optional custom function to handle incoming data.
+                                If provided, this will replace the default data
+                                handling logic entirely.
+        """
+        super().__init__(
+            topic,
+            zenoh_session,
+            name,
+            enable_fps_tracking,
+            fps_log_interval,
+            custom_data_handler,
+        )
+        self._latest_raw_data: bytes | None = None
+
+    def _data_handler(self, sample: zenoh.Sample) -> None:
+        """Handle incoming depth image data.
+
+        Args:
+            sample: Zenoh sample containing encoded depth image data.
+        """
+        with self._data_lock:
+            self._latest_raw_data = sample.payload.to_bytes()
+            self._active = True
+
+        self._update_fps_metrics()
+
+    def get_latest_data(self) -> np.ndarray | None:
+        """Get the latest depth data.
+
+        Returns:
+            Tuple of (depth_image, depth_min, depth_max) if available, None otherwise.
+            depth_image: Depth values in meters as numpy array (HxW)
+        """
+        with self._data_lock:
+            if self._latest_raw_data is None:
+                return None
+
+            if decode_depth is None:
+                logger.error(
+                    f"Cannot decode depth image for {self._name}: dexsensor not available"
+                )
+                return None
+
+            try:
+                # Decode the depth image
+                depth, depth_min, depth_max = decode_depth(self._latest_raw_data)
+                return depth
+            except Exception as e:
+                logger.error(f"Failed to decode depth image for {self._name}: {e}")
+                return None
+
+
+class RGBDCameraSubscriber(BaseZenohSubscriber):
+    """Zenoh subscriber for RGBD camera data.
+
+    This subscriber handles both RGB and depth data from an RGBD camera,
+    subscribing to separate topics for RGB and depth streams.
+    """
+
+    def __init__(
+        self,
+        rgb_topic: str,
+        depth_topic: str,
+        zenoh_session: zenoh.Session,
+        name: str = "rgbd_camera_subscriber",
+        enable_fps_tracking: bool = True,
+        fps_log_interval: int = 30,
+        custom_data_handler: CustomDataHandler | None = None,
+    ) -> None:
+        """Initialize the RGBD camera subscriber.
+
+        Args:
+            rgb_topic: Zenoh topic for RGB data.
+            depth_topic: Zenoh topic for depth data.
+            zenoh_session: Active Zenoh session for communication.
+            name: Name for logging purposes.
+            enable_fps_tracking: Whether to track and log FPS metrics.
+            fps_log_interval: Number of frames between FPS calculations.
+            custom_data_handler: Optional custom function to handle incoming data.
+                                If provided, this will replace the default data
+                                handling logic entirely.
+        """
+        # Initialize with RGB topic as primary
+        super().__init__(
+            rgb_topic,
+            zenoh_session,
+            name,
+            enable_fps_tracking,
+            fps_log_interval,
+            custom_data_handler,
+        )
+
+        # Create separate subscribers for RGB and depth
+        self._rgb_subscriber = RGBCameraSubscriber(
+            rgb_topic,
+            zenoh_session,
+            f"{name}_rgb",
+            enable_fps_tracking,
+            fps_log_interval,
+            custom_data_handler,
+        )
+        self._depth_subscriber = DepthCameraSubscriber(
+            depth_topic,
+            zenoh_session,
+            f"{name}_depth",
+            enable_fps_tracking,
+            fps_log_interval,
+            custom_data_handler,
+        )
+
+    def _data_handler(self, sample: zenoh.Sample) -> None:
+        """Handle incoming data - not used as we use separate subscribers."""
+        pass
+
+    def get_latest_data(self) -> tuple[np.ndarray, np.ndarray] | None:
+        """Get the latest RGBD data.
+
+        Returns:
+            Tuple of (rgb_image, depth_image, depth_min, depth_max) if both available, None otherwise.
+            rgb_image: RGB image as numpy array (HxWxC)
+            depth_image: Depth values in meters as numpy array (HxW)
+        """
+        rgb_image = self._rgb_subscriber.get_latest_data()
+        depth_image = self._depth_subscriber.get_latest_data()
+
+        if rgb_image is not None and depth_image is not None:
+            return rgb_image, depth_image
+        return None
+
+    def get_latest_rgb(self) -> np.ndarray | None:
+        """Get the latest RGB image.
+
+        Returns:
+            Latest RGB image as numpy array (HxWxC) if available, None otherwise.
+        """
+        return self._rgb_subscriber.get_latest_image()
+
+    def get_latest_depth(self) -> np.ndarray | None:
+        """Get the latest depth image.
+
+        Returns:
+            Latest depth image as numpy array (HxW) with values in meters if available, None otherwise.
+        """
+        return self._depth_subscriber.get_latest_data()
+
+    def wait_for_active(self, timeout: float = 5.0) -> bool:
+        """Wait for both RGB and depth subscribers to start receiving data.
+
+        Args:
+            timeout: Maximum time to wait in seconds.
+
+        Returns:
+            True if both subscribers become active, False if timeout is reached.
+        """
+        rgb_active = self._rgb_subscriber.wait_for_active(timeout)
+        depth_active = self._depth_subscriber.wait_for_active(timeout)
+        return rgb_active and depth_active
+
+    def is_active(self) -> bool:
+        """Check if both RGB and depth subscribers are actively receiving data.
+
+        Returns:
+            True if both subscribers are active, False otherwise.
+        """
+        return self._rgb_subscriber.is_active() and self._depth_subscriber.is_active()
+
+    def shutdown(self) -> None:
+        """Stop both subscribers and release resources."""
+        self._rgb_subscriber.shutdown()
+        self._depth_subscriber.shutdown()
+        super().shutdown()
+
+    @property
+    def rgb_fps(self) -> float:
+        """Get the RGB stream FPS measurement.
+
+        Returns:
+            Current RGB frames per second measurement.
+        """
+        return self._rgb_subscriber.fps
+
+    @property
+    def depth_fps(self) -> float:
+        """Get the depth stream FPS measurement.
+
+        Returns:
+            Current depth frames per second measurement.
+        """
+        return self._depth_subscriber.fps

dexcontrol/utils/subscribers/decoders.py

@@ -0,0 +1,83 @@
+# Copyright (c) 2025 Dexmate CORPORATION & AFFILIATES. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 with Commons Clause License
+# Condition v1.0 [see LICENSE for details].
+
+"""Decoder functions for Zenoh subscribers.
+
+This module provides common decoder functions that can be used with the
+GenericZenohSubscriber to transform raw Zenoh bytes into specific data formats.
+"""
+
+import json
+from collections.abc import Callable
+from typing import Any, Type, TypeVar
+
+import zenoh
+from google.protobuf.message import Message
+
+M = TypeVar("M", bound=Message)
+
+# Type alias for decoder functions
+DecoderFunction = Callable[[zenoh.ZBytes], Any]
+
+
+def protobuf_decoder(message_type: Type[M]) -> DecoderFunction:
+    """Create a decoder function for a specific protobuf message type.
+
+    Args:
+        message_type: Protobuf message class to decode to.
+
+    Returns:
+        Decoder function that parses bytes into the specified protobuf message.
+    """
+
+    def decode(data: zenoh.ZBytes) -> M:
+        message = message_type()
+        message.ParseFromString(data.to_bytes())
+        return message
+
+    return decode
+
+
+def raw_bytes_decoder(data: zenoh.ZBytes) -> bytes:
+    """Decoder that returns raw bytes.
+
+    Args:
+        data: Zenoh bytes data.
+
+    Returns:
+        Raw bytes data.
+    """
+    return data.to_bytes()
+
+
+def json_decoder(data: zenoh.ZBytes) -> Any:
+    """Decoder that parses JSON from bytes.
+
+    Args:
+        data: Zenoh bytes data containing JSON.
+
+    Returns:
+        Parsed JSON data.
+
+    Raises:
+        json.JSONDecodeError: If the data is not valid JSON.
+    """
+    return json.loads(data.to_bytes().decode("utf-8"))
+
+
+def string_decoder(data: zenoh.ZBytes, encoding: str = "utf-8") -> str:
+    """Decoder that converts bytes to string.
+
+    Args:
+        data: Zenoh bytes data.
+        encoding: Text encoding to use for decoding.
+
+    Returns:
+        Decoded string.
+
+    Raises:
+        UnicodeDecodeError: If the data cannot be decoded with the specified encoding.
+    """
+    return data.to_bytes().decode(encoding)