diff --git a/src/lerobot/cameras/realsense/camera_realsense.py b/src/lerobot/cameras/realsense/camera_realsense.py index a7d0b786dc1..ccdcc412ee0 100644 --- a/src/lerobot/cameras/realsense/camera_realsense.py +++ b/src/lerobot/cameras/realsense/camera_realsense.py @@ -749,11 +749,11 @@ def _build_frame_snapshot(self, frames: Any) -> _FrameSnapshot: depth = None color_frame = frames.get_color_frame() - if color_frame is not None: + if color_frame: color = np.ascontiguousarray(_decode_color_frame_to_rgb(color_frame)) depth_frame = frames.get_depth_frame() - if depth_frame is not None: + if depth_frame: depth = np.ascontiguousarray(np.asanyarray(depth_frame.get_data()).copy()) return _FrameSnapshot(color=color, depth=depth) diff --git a/tests/cameras/test_realsense_frame_snapshot.py b/tests/cameras/test_realsense_frame_snapshot.py new file mode 100644 index 00000000000..87ff95719c3 --- /dev/null +++ b/tests/cameras/test_realsense_frame_snapshot.py @@ -0,0 +1,82 @@ +#!/usr/bin/env python + +# Copyright 2024 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest +from unittest.mock import MagicMock, patch + +import numpy as np + +from lerobot.cameras.realsense import D405_MODEL_TRAITS, SharedRealSenseManager + + +class TestBuildFrameSnapshot(unittest.TestCase): + """Regression tests for falsey pyrealsense2 frames returned by disabled streams.""" + + def setUp(self): + self.manager = SharedRealSenseManager(serial_number="042", model_traits=D405_MODEL_TRAITS) + + def test_ignores_invalid_color_frame(self): + invalid_color_frame = MagicMock() + invalid_color_frame.__bool__.return_value = False + frames = MagicMock() + frames.get_color_frame.return_value = invalid_color_frame + frames.get_depth_frame.return_value = None + + with patch( + "lerobot.cameras.realsense.camera_realsense._decode_color_frame_to_rgb" + ) as decode_color_frame: + snapshot = self.manager._build_frame_snapshot(frames) + + decode_color_frame.assert_not_called() + invalid_color_frame.get_data.assert_not_called() + self.assertIsNone(snapshot.color) + self.assertIsNone(snapshot.depth) + + def test_ignores_invalid_depth_frame(self): + invalid_depth_frame = MagicMock() + invalid_depth_frame.__bool__.return_value = False + frames = MagicMock() + frames.get_color_frame.return_value = None + frames.get_depth_frame.return_value = invalid_depth_frame + + snapshot = self.manager._build_frame_snapshot(frames) + + invalid_depth_frame.get_data.assert_not_called() + self.assertIsNone(snapshot.color) + self.assertIsNone(snapshot.depth) + + def test_decodes_valid_frames(self): + color_frame = MagicMock() + color_frame.__bool__.return_value = True + depth_frame = MagicMock() + depth_frame.__bool__.return_value = True + depth_data = np.arange(4, dtype=np.uint16).reshape(2, 2) + depth_frame.get_data.return_value = depth_data + frames = MagicMock() + frames.get_color_frame.return_value = color_frame + frames.get_depth_frame.return_value = depth_frame + expected_color = np.arange(12, dtype=np.uint8).reshape(2, 2, 3) + + with patch( + "lerobot.cameras.realsense.camera_realsense._decode_color_frame_to_rgb", + return_value=expected_color, + ) as decode_color_frame: + snapshot = self.manager._build_frame_snapshot(frames) + + decode_color_frame.assert_called_once_with(color_frame) + depth_frame.get_data.assert_called_once_with() + np.testing.assert_array_equal(snapshot.color, expected_color) + np.testing.assert_array_equal(snapshot.depth, depth_data)