Merge pull request #139 from pollen-robotics/138-provide-raw-images-to-teleop

138 provide raw images to teleop

Author: FabienDanieau

.github/workflows/pytest.yml

@@ -42,7 +42,7 @@ jobs:
           echo "total=$TOTAL" >> $GITHUB_ENV
           echo "### Total coverage: ${TOTAL}%" >> $GITHUB_STEP_SUMMARY
       - name: Make badge
-        uses: schneegans/dynamic-badges-action@v1.6.0
+        uses: schneegans/dynamic-badges-action@v1.7.0
         with:
           # GIST_TOKEN is a GitHub personal access token with scope "gist".
           auth: ${{ secrets.GIST_TOKEN }}

examples/camera_wrappers_examples/teleopWrapper_example.py

@@ -27,7 +27,7 @@
 
 
 def spawn_procs(names: List[str]) -> Dict[str, sp.Popen]:  # type: ignore [type-arg]
-    width, height = 1280, 720
+    width, height = 960, 720
     command = [
         "ffplay",
         "-i",
@@ -60,9 +60,11 @@ def spawn_procs(names: List[str]) -> Dict[str, sp.Popen]:  # type: ignore [type-
 procs = spawn_procs(["left", "right"])
 
 while True:
-    data, lat, _ = w.get_data()
+    data, lat, _ = w.get_data_h264()
     logging.info(lat)
     for name, packets in data.items():
+        # if name == "left_raw" or name == "right_raw":
+        #    continue
         io = procs[name].stdin
         if io is not None:
             io.write(packets)

pollen_vision/pollen_vision/camera_wrappers/depthai/cam_config.py

@@ -1,8 +1,9 @@
 """Camera configuration class for depthai cameras."""
 
 import json
-from typing import Dict, Optional, Tuple
+from typing import Dict, List, Optional, Tuple
 
+import cv2
 import depthai as dai
 import numpy as np
 import numpy.typing as npt
@@ -66,6 +67,10 @@ def __init__(
         }
         self.calib: dai.CalibrationHandler = dai.CalibrationHandler()
 
+        # lazy init, camera needs to be connected to
+        self.P_left: Optional[cv2.UMat] = None
+        self.P_right: Optional[cv2.UMat] = None
+
     def get_device_info(self) -> dai.DeviceInfo:
         """Returns a dai.DeviceInfo object with the mx_id.
         This allows connecting to multiple devices plugged in the host machine at the same time,
@@ -146,3 +151,54 @@ def to_string(self) -> str:
         ret_string += "Undistort maps are: " + "set" if self.undistort_maps["left"] is not None else "not set"
 
         return ret_string
+
+    def compute_projection_matrices(self) -> Tuple[cv2.UMat, cv2.UMat]:
+        left_socket = get_socket_from_name("left", self.name_to_socket)
+        right_socket = get_socket_from_name("right", self.name_to_socket)
+
+        left_D = np.array(self.calib.getDistortionCoefficients(left_socket))
+        right_D = np.array(self.calib.getDistortionCoefficients(right_socket))
+
+        R = np.array(self.calib.getStereoRightRectificationRotation())
+
+        T = np.array(self.calib.getCameraTranslationVector(left_socket, right_socket))
+        T *= 0.01  # to meter for ROS
+
+        R1, R2, P1, P2, Q, _, _ = cv2.stereoRectify(
+            self.get_K_left(),
+            left_D,
+            self.get_K_right(),
+            right_D,
+            self.undistort_resolution,
+            R,
+            T,
+            flags=0,
+        )
+        return P1, P2
+
+    def to_ROS_msg(
+        self, side: str = "left"
+    ) -> Tuple[int, int, str, List[float], npt.NDArray[np.float32], npt.NDArray[np.float32], npt.NDArray[np.float32]]:
+        # as defined in https://docs.ros.org/en/melodic/api/sensor_msgs/html/msg/CameraInfo.html
+
+        height = self.resize_resolution[1]
+        width = self.resize_resolution[0]
+        distortion_model = "plumb_bob"
+        if self.calib.getDistortionModel(get_socket_from_name(side, self.name_to_socket)) == dai.CameraModel.Fisheye:
+            distortion_model = "equidistant"
+        D = self.calib.getDistortionCoefficients(get_socket_from_name(side, self.name_to_socket))
+
+        if self.P_left is None or self.P_right is None:
+            self.P_left, self.P_right = self.compute_projection_matrices()
+
+        if side == "left":
+            K = self.get_K_left().flatten()
+            R = np.array(self.calib.getStereoLeftRectificationRotation()).flatten()
+            P = np.array(self.P_left).flatten()
+
+        else:
+            K = self.get_K_right().flatten()
+            R = np.array(self.calib.getStereoRightRectificationRotation()).flatten()
+            P = np.array(self.P_right).flatten()
+
+        return height, width, distortion_model, D, K, R, P

pollen_vision/pollen_vision/camera_wrappers/depthai/teleop.py

@@ -30,18 +30,28 @@ def __init__(
         exposure_params: Optional[Tuple[int, int]] = None,
         mx_id: str = "",
     ) -> None:
+        self._data_h264: Dict[str, npt.NDArray[np.uint8]] = {}
+        self._latency_h264: Dict[str, float] = {}
+        self._ts_h264: Dict[str, timedelta] = {}
+
+        self._data_mjpeg: Dict[str, npt.NDArray[np.uint8]] = {}
+        self._latency_mjpeg: Dict[str, float] = {}
+        self._ts_mjpeg: Dict[str, timedelta] = {}
+
+        self._queues_mjpeg: Dict[str, dai.DataOutputQueue] = {}
+
         super().__init__(
             cam_config_json,
             fps,
             force_usb2=force_usb2,
-            resize=(1280, 720),
+            resize=(960, 720),
             rectify=rectify,
             exposure_params=exposure_params,
             mx_id=mx_id,
             isp_scale=(2, 3),
         )
 
-    def get_data(
+    def get_data_h264(
         self,
     ) -> Tuple[Dict[str, npt.NDArray[np.uint8]], Dict[str, float], Dict[str, timedelta]]:
         """Extends the get_data method of the Wrapper class to return the h264 encoded left and right images.
@@ -51,32 +61,57 @@ def get_data(
                 latencies and timestamps for each camera.
         """
 
-        data, latency, ts = super().get_data()
+        for name, queue in self.queues.items():
+            pkt = queue.get()
+            self._data_h264[name] = pkt.getData()
+            self._latency_h264[name] = dai.Clock.now() - pkt.getTimestamp()  # type: ignore[call-arg]
+            self._ts_h264[name] = pkt.getTimestamp()
+
+        return self._data_h264, self._latency_h264, self._ts_h264
+
+    def get_data_mjpeg(self) -> Tuple[Dict[str, npt.NDArray[np.uint8]], Dict[str, float], Dict[str, timedelta]]:
+        for name, queue in self._queues_mjpeg.items():
+            pkt = queue.get()
+            self._data_mjpeg[name] = pkt.getData()  # type: ignore[attr-defined]
+            self._latency_mjpeg[name] = dai.Clock.now() - pkt.getTimestamp()  # type: ignore[attr-defined, call-arg]
+            self._ts_mjpeg[name] = pkt.getTimestamp()  # type: ignore[attr-defined]
 
-        for name, pkt in data.items():
-            data[name] = pkt.getData()
+        return self._data_mjpeg, self._latency_mjpeg, self._ts_mjpeg
 
-        return data, latency, ts
+    def _create_output_streams(self, pipeline: dai.Pipeline) -> dai.Pipeline:
+        super()._create_output_streams(pipeline)
+
+        self.xout_left_mjpeg = pipeline.createXLinkOut()
+        self.xout_left_mjpeg.setStreamName("left_mjpeg")
+
+        self.xout_right_mjpeg = pipeline.createXLinkOut()
+        self.xout_right_mjpeg.setStreamName("right_mjpeg")
+
+        return pipeline
 
     def _link_pipeline(self, pipeline: dai.Pipeline) -> dai.Pipeline:
         """Overloads the base class abstract method to link the pipeline with the nodes together."""
 
         self.left.isp.link(self.left_manip.inputImage)
         self.left_manip.out.link(self.left_encoder.input)
-        self.right.isp.link(self.right_manip.inputImage)
-        self.right_manip.out.link(self.right_encoder.input)
-
+        self.left_manip.out.link(self.left_encoder_mjpeg.input)
         self.left_encoder.bitstream.link(self.xout_left.input)
         self.right_encoder.bitstream.link(self.xout_right.input)
 
+        self.right.isp.link(self.right_manip.inputImage)
+        self.right_manip.out.link(self.right_encoder.input)
+        self.right_manip.out.link(self.right_encoder_mjpeg.input)
+        self.left_encoder_mjpeg.bitstream.link(self.xout_left_mjpeg.input)
+        self.right_encoder_mjpeg.bitstream.link(self.xout_right_mjpeg.input)
+
         return pipeline
 
     def _create_encoders(self, pipeline: dai.Pipeline) -> dai.Pipeline:
         """Creates the h264 encoders for the left and right images."""
 
         profile = dai.VideoEncoderProperties.Profile.H264_BASELINE
         bitrate = 4000
-        numBFrames = 0  # gstreamer recommends 0 B frames
+        numBFrames = 0  # no B frames for streaming
         self.left_encoder = pipeline.create(dai.node.VideoEncoder)
         self.left_encoder.setDefaultProfilePreset(self.cam_config.fps, profile)
         self.left_encoder.setKeyframeFrequency(self.cam_config.fps)  # every 1s
@@ -91,6 +126,16 @@ def _create_encoders(self, pipeline: dai.Pipeline) -> dai.Pipeline:
         self.right_encoder.setBitrateKbps(bitrate)
         # self.right_encoder.setQuality(self.cam_config.encoder_quality)
 
+        profile = dai.VideoEncoderProperties.Profile.MJPEG
+
+        self.left_encoder_mjpeg = pipeline.create(dai.node.VideoEncoder)
+        self.left_encoder_mjpeg.setDefaultProfilePreset(self.cam_config.fps, profile)
+        # self.left_encoder_mjpeg.setLossless(True)
+
+        self.right_encoder_mjpeg = pipeline.create(dai.node.VideoEncoder)
+        self.right_encoder_mjpeg.setDefaultProfilePreset(self.cam_config.fps, profile)
+        # self.right_encoder_mjpeg.setLossless(True)
+
         return pipeline
 
     def _create_pipeline(self) -> dai.Pipeline:
@@ -111,7 +156,11 @@ def _create_queues(self) -> Dict[str, dai.DataOutputQueue]:
         """Extends the base class method _create_queues() to add the h264 encoded left and right images queues."""
 
         # config for video: https://docs.luxonis.com/projects/api/en/latest/components/device/#output-queue-maxsize-and-blocking
-        queues: Dict[str, dai.DataOutputQueue] = {}
+        queues_h264: Dict[str, dai.DataOutputQueue] = {}
         for name in ["left", "right"]:
-            queues[name] = self._device.getOutputQueue(name, maxSize=10, blocking=True)
-        return queues
+            queues_h264[name] = self._device.getOutputQueue(name, maxSize=30, blocking=True)
+
+        for name in ["left_mjpeg", "right_mjpeg"]:
+            self._queues_mjpeg[name] = self._device.getOutputQueue(name, maxSize=1, blocking=False)
+
+        return queues_h264

pollen_vision/pollen_vision/camera_wrappers/depthai/wrapper.py

@@ -99,6 +99,7 @@ def _prepare(self) -> None:
             self._set_undistort_maps()
 
         self.pipeline = self._create_pipeline()
+        self.pipeline.setXLinkChunkSize(0)  # better usb performance
 
         self._device.startPipeline(self.pipeline)
         self.queues = self._create_queues()

setup.cfg

@@ -36,7 +36,7 @@ checkpoints =
 all = mobile-sam @ git+https://github.com/pollen-robotics/MobileSAM
       ram @ git+https://github.com/pollen-robotics/recognize-anything
       datasets==2.18.0
-      depthai==2.25.0.0
+      depthai==2.27.0.0
       datasets==2.18.0
       supervision==0.20.0
       inference-gpu[yolo-world]==0.9.13
@@ -52,7 +52,7 @@ vision = mobile-sam @ git+https://github.com/pollen-robotics/MobileSAM
          scikit-learn==1.2.2
          transformers==4.40.2
          FramesViewer==1.0.2
-depthai_wrapper = depthai==2.25.0.0
+depthai_wrapper = depthai==2.27.0.0
 realsense_wrapper = pyrealsense2==2.55.1.6486
 gradio = gradio==4.21.0
 open3d = open3d==0.17.0