tmp

Author: juelg

pyproject.toml

@@ -11,7 +11,6 @@ dependencies = ["websockets>=11.0",
     "numpy",
     "typer~=0.9",
     "gymnasium~=0.29.1",
-    "pydantic_yaml~=1.3",
     "absl-py~=2.1",
     "etils[epath]>=1.7.0",
     "glfw~=2.7",

python/rcs/camera/digit_cam.py

@@ -1,40 +1,29 @@
 from digit_interface.digit import Digit
-from rcs.camera.hw import BaseHardwareCameraSet, HWCameraSetConfig
+from rcs._core.common import BaseCameraConfig
+from rcs.camera.hw import HardwareCamera
 from rcs.camera.interface import CameraFrame, DataFrame, Frame
 
-
-class DigitConfig(HWCameraSetConfig):
-    """
-    Configuration for the DIGIT device.
-    This class is used to define the settings for the DIGIT device.
-    """
-
-    stream_name: str = "QVGA"  # options: "QVGA" (60 and 30 fps), "VGA" (30 and 15 fps)
-
-
-class DigitCam(BaseHardwareCameraSet):
+class DigitCam(HardwareCamera):
     """
     This module provides an interface to interact with the DIGIT device.
     It allows for connecting to the device, changing settings, and retrieving information.
     """
 
-    def __init__(self, cfg: DigitConfig):
-        self._cfg = cfg
-        super().__init__()
+    def __init__(self, cameras: dict[str, BaseCameraConfig]):
+        self.cameras = cameras
         self._cameras: dict[str, Digit] = {}
-        self.initalize(self.config)
 
-    def initalize(self, cfg: HWCameraSetConfig):
+    def open(self):
         """
         Initialize the digit interface with the given configuration.
         :param cfg: Configuration for the DIGIT device.
         """
-        for name, serial in cfg.name_to_identifier.items():
-            digit = Digit(serial, name)
+        for name, camera in self.cameras.items():
+            digit = Digit(camera.identifier, name)
             digit.connect()
             self._cameras[name] = digit
 
-    def _poll_frame(self, camera_name: str) -> Frame:
+    def poll_frame(self, camera_name: str) -> Frame:
         """Polls the frame from the camera with the given name."""
         digit = self._cameras[camera_name]
         frame = digit.get_frame()
@@ -45,6 +34,13 @@ def _poll_frame(self, camera_name: str) -> Frame:
 
         return Frame(camera=cf)
 
-    @property
-    def config(self) -> DigitConfig:
-        return self._cfg
+    def close(self):
+        """
+        Closes the connection to the DIGIT device.
+        """
+        for digit in self._cameras.values():
+            digit.disconnect()
+        self._cameras = {}
+
+    def config(self, camera_name) -> BaseCameraConfig:
+        return self.cameras[camera_name]

python/rcs/camera/hw.py

@@ -1,47 +1,92 @@
 import logging
 import threading
 import typing
-from abc import ABC, abstractmethod
 from datetime import datetime
 from pathlib import Path
 from time import sleep
 
 import cv2
 import numpy as np
-from pydantic import Field
-from rcs.camera.interface import (
-    BaseCameraConfig,
-    BaseCameraSetConfig,
-    Frame,
-    FrameSet,
-    SimpleFrameRate,
-)
-
-
-class HWCameraSetConfig(BaseCameraSetConfig):
-    cameras: dict[str, BaseCameraConfig] = Field(default={})
-    warm_up_disposal_frames: int = 30  # frames
-    record_path: str = "camera_frames"
-    max_buffer_frames: int = 1000
-
-
-# TODO(juelg): refactor camera thread into their own class, to avoid a base hardware camera set class
-# TODO(juelg): add video recording
-class BaseHardwareCameraSet(ABC):
-    """This base class should have the ability to poll in a separate thread for all cameras and store them in a buffer.
-    Implements BaseCameraSet
+from rcs._core.common import BaseCameraConfig
+from rcs.camera.interface import BaseCameraSet, Frame, FrameSet, SimpleFrameRate
+
+
+class HardwareCamera(typing.Protocol):
+    """Implementation of a hardware camera potentially a set of cameras of the same kind."""
+
+    def open(self):
+        """Should open the camera and prepare it for polling."""
+
+    def close(self):
+        """Should close the camera and release all resources."""
+
+    def config(self, camera_name: str) -> BaseCameraConfig:
+        """Should return the configuration object of the cameras."""
+
+    def poll_frame(self, camera_name: str) -> Frame:
+        """Should return the latest frame from the camera with the given name.
+
+        This method should be thread safe.
+        """
+
+    @property
+    def camera_names(self) -> list[str]:
+        """Should return a list of the activated human readable names of the cameras."""
+
+
+class HardwareCameraSet(BaseCameraSet):
+    """This base class polls in a separate thread for all cameras and stores them in a buffer.
+
+    Cameras can consist of multiple cameras, e.g. RealSense cameras.
     """
 
-    def __init__(self):
-        self._buffer: list[FrameSet | None] = [None for _ in range(self.config.max_buffer_frames)]
+    def __init__(self, cameras: list[HardwareCamera], warm_up_disposal_frames: int = 30, max_buffer_frames: int = 1000):
+        self.cameras = cameras
+        self.camera_dict, self.camera_names = self._cameras_util()
+        self.name_to_identifier = self._name_to_identifier()
+        self.frame_rate = self._frames_rate()
+        self.rate_limiter = SimpleFrameRate(self.frame_rate)
+
+        self.warm_up_disposal_frames = warm_up_disposal_frames
+        self.max_buffer_frames = max_buffer_frames
+        self._buffer: list[FrameSet | None] = [None for _ in range(self.max_buffer_frames)]
         self._buffer_lock = threading.Lock()
         self.running = False
         self._thread: threading.Thread | None = None
         self._logger = logging.getLogger(__name__)
         self._next_ring_index = 0
         self._buffer_len = 0
         self.writer: dict[str, cv2.VideoWriter] = {}
-        self.rate = SimpleFrameRate()
+
+    def _name_to_identifier(self) -> dict[str, str]:
+        """Returns a dictionary mapping the camera names to their identifiers."""
+        name_to_id: dict[str, str] = {}
+        for camera in self.cameras:
+            for name in camera.camera_names:
+                name_to_id[name] = camera.config(name).identifier
+        return name_to_id
+
+    def _frames_rate(self) -> int:
+        """Checks if all cameras have the same frame rate."""
+        frame_rates = {camera.config(name).frame_rate for camera in self.cameras for name in camera.camera_names}
+        if len(frame_rates) > 1:
+            msg = "All cameras must have the same frame rate. Different frame rates are not supported."
+            raise ValueError(msg)
+        if len(frame_rates) == 0:
+            self._logger.warning("No camera found, empty polling with 1 fps.")
+            return 1
+        return frame_rates[0]
+
+    def _cameras_util(self) -> tuple[dict[str, HardwareCamera], list[str]]:
+        """Utility function to create a dictionary of cameras and a list of camera names."""
+        camera_dict: dict[str, HardwareCamera] = {}
+        camera_names: list[str] = []
+        for camera in self.cameras:
+            camera_names.extend(camera.camera_names)
+            for name in camera.camera_names:
+                assert name not in camera_dict, f"Camera name {name} not unique."
+                camera_dict[name] = camera
+        return camera_dict, camera_names
 
     def buffer_size(self) -> int:
         return len(self._buffer) - self._buffer.count(None)
@@ -64,7 +109,7 @@ def get_timestamp_frames(self, ts: datetime) -> FrameSet | None:
         # iterate through the buffer and find the closest timestamp
         with self._buffer_lock:
             for i in range(self._buffer_len):
-                idx = (self._next_ring_index - i - 1) % self.config.max_buffer_frames  # iterate backwards
+                idx = (self._next_ring_index - i - 1) % self.max_buffer_frames  # iterate backwards
                 assert self._buffer[idx] is not None
                 item: FrameSet = typing.cast(FrameSet, self._buffer[idx])
                 assert item.avg_timestamp is not None
@@ -82,6 +127,8 @@ def stop(self):
     def close(self):
         if self.running and self._thread is not None:
             self.stop()
+        for camera in self.cameras:
+            camera.close()
         self.stop_video()
 
     def start(self, warm_up: bool = True):
@@ -101,8 +148,8 @@ def record_video(self, path: Path, str_id: str):
                 str(path / f"episode_{str_id}_{camera}.mp4"),
                 # migh require to install ffmpeg
                 cv2.VideoWriter_fourcc(*"mp4v"),  # type: ignore
-                self.config.frame_rate,
-                (self.config.resolution_width, self.config.resolution_height),
+                self.frame_rate,
+                (self.config(camera).resolution_width, self.config(camera).resolution_height),
             )
 
     def recording_ongoing(self) -> bool:
@@ -117,32 +164,35 @@ def stop_video(self):
                 self.writer = {}
 
     def warm_up(self):
-        for _ in range(self.config.warm_up_disposal_frames):
+        for _ in range(self.warm_up_disposal_frames):
             for camera_name in self.camera_names:
-                self._poll_frame(camera_name)
-            self.rate(self.config.frame_rate)
+                self.poll_frame(camera_name)
+            self.rate_limiter()
 
     def polling_thread(self, warm_up: bool = True):
+        for camera in self.cameras:
+            camera.open()
         if warm_up:
             self.warm_up()
         while self.running:
             frame_set = self.poll_frame_set()
             # buffering
             with self._buffer_lock:
                 self._buffer[self._next_ring_index] = frame_set
-                self._next_ring_index = (self._next_ring_index + 1) % self.config.max_buffer_frames
-                self._buffer_len = max(self._buffer_len + 1, self.config.max_buffer_frames)
+                self._next_ring_index = (self._next_ring_index + 1) % self.max_buffer_frames
+                self._buffer_len = max(self._buffer_len + 1, self.max_buffer_frames)
             # video recording
             for camera_key, writer in self.writer.items():
                 if frame_set is not None:
                     writer.write(frame_set.frames[camera_key].camera.color.data[:, :, ::-1])
-            self.rate(self.config.frame_rate)
+            self.rate_limiter()
 
     def poll_frame_set(self) -> FrameSet:
         """Gather frames over all available cameras."""
         frames: dict[str, Frame] = {}
         for camera_name in self.camera_names:
-            frame = self._poll_frame(camera_name)
+            # callback
+            frame = self.poll_frame(camera_name)
             frames[camera_name] = frame
         # filter none
         timestamps: list[float] = [frame.avg_timestamp for frame in frames.values() if frame.avg_timestamp is not None]
@@ -151,29 +201,14 @@ def poll_frame_set(self) -> FrameSet:
     def clear_buffer(self):
         """Deletes all frames from the buffer."""
         with self._buffer_lock:
-            self._buffer = [None for _ in range(self.config.max_buffer_frames)]
+            self._buffer = [None for _ in range(self.max_buffer_frames)]
             self._next_ring_index = 0
             self._buffer_len = 0
         self.wait_for_frames()
 
-    @property
-    @abstractmethod
-    def config(self) -> HWCameraSetConfig:
-        """Should return the configuration object of the cameras."""
-
-    @abstractmethod
-    def _poll_frame(self, camera_name: str) -> Frame:
-        """Should return the latest frame from the camera with the given name.
-
-        This method should be thread safe.
-        """
-
-    @property
-    def camera_names(self) -> list[str]:
-        """Should return a list of the activated human readable names of the cameras."""
-        return list(self.config.cameras)
+    def config(self, camera_name: str) -> BaseCameraConfig:
+        """Returns the configuration object of the cameras."""
+        return self.camera_dict[camera_name].config(camera_name)
 
-    @property
-    def name_to_identifier(self) -> dict[str, str]:
-        # return {key: camera.identifier for key, camera in self._cfg.cameras.items()}
-        return self.config.name_to_identifier
+    def poll_frame(self, camera_name: str) -> Frame:
+        return self.camera_dict[camera_name].poll_frame(camera_name)

python/rcs/camera/interface.py

@@ -4,27 +4,29 @@
 from time import sleep, time
 from typing import Any, Protocol
 
-
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
 
 
 class SimpleFrameRate:
-    def __init__(self):
-        self.t = None
-        self._last_print = None
+    def __init__(self, frame_rate: int | float):
+        self.t: float  | None = None
+        self._last_print: float | None = None
+        self.frame_rate = frame_rate
 
     def reset(self):
         self.t = None
 
-    def __call__(self, frame_rate: int | float):
+    def __call__(self):
         if self.t is None:
             self.t = time()
             self._last_print = self.t
-            sleep(1 / frame_rate if isinstance(frame_rate, int) else frame_rate)
+            sleep(1 / self.frame_rate if isinstance(self.frame_rate, int) else self.frame_rate)
             return
         sleep_time = (
-            1 / frame_rate - (time() - self.t) if isinstance(frame_rate, int) else frame_rate - (time() - self.t)
+            1 / self.frame_rate - (time() - self.t)
+            if isinstance(self.frame_rate, int)
+            else self.frame_rate - (time() - self.t)
         )
         if sleep_time > 0:
             sleep(sleep_time)
@@ -35,18 +37,6 @@ def __call__(self, frame_rate: int | float):
         self.t = time()
 
 
-@dataclass(kw_only=True)
-class BaseCameraConfig:
-    identifier: str
-    resolution_width: int = 1280  # pixels
-    resolution_height: int = 720  # pixels
-    frame_rate: int = 15  # Hz
-
-    @property
-    def name_to_identifier(self):
-        return {key: camera.identifier for key, camera in self.cameras.items()}
-
-
 @dataclass(kw_only=True)
 class DataFrame:
     data: Any
@@ -100,10 +90,6 @@ def clear_buffer(self):
     def close(self):
         """Stops any running threads e.g. for exitting."""
 
-    @property
-    def config(self) -> BaseCameraSetConfig:
-        """Return the configuration object of the cameras."""
-
     @property
     def camera_names(self) -> list[str]:
         """Returns a list of the activated human readable names of the cameras."""