99 lines
2.6 KiB
Python
99 lines
2.6 KiB
Python
|
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
from lerobot.common.robot_devices.cameras.opencv import OpenCVCamera
|
||
|
|
from lerobot.common.robot_devices.utils import RobotDeviceNotConnectedError, RobotDeviceAlreadyConnectedError
|
||
|
|
|
||
|
|
|
||
|
|
def test_camera():
|
||
|
|
# Test instantiating with missing camera index raises an error
|
||
|
|
with pytest.raises(ValueError):
|
||
|
|
camera = OpenCVCamera()
|
||
|
|
|
||
|
|
# Test instantiating with a wrong camera index raises an error
|
||
|
|
with pytest.raises(ValueError):
|
||
|
|
camera = OpenCVCamera(-1)
|
||
|
|
|
||
|
|
# Test instantiating
|
||
|
|
camera = OpenCVCamera(0)
|
||
|
|
|
||
|
|
# Test reading, async reading, disconnecting before connecting raises an error
|
||
|
|
with pytest.raises(RobotDeviceNotConnectedError):
|
||
|
|
camera.read()
|
||
|
|
with pytest.raises(RobotDeviceNotConnectedError):
|
||
|
|
camera.async_read()
|
||
|
|
with pytest.raises(RobotDeviceNotConnectedError):
|
||
|
|
camera.disconnect()
|
||
|
|
|
||
|
|
# Test deleting the object without connecting first
|
||
|
|
del camera
|
||
|
|
|
||
|
|
# Test connecting
|
||
|
|
camera = OpenCVCamera(0)
|
||
|
|
camera.connect()
|
||
|
|
assert camera.is_connected
|
||
|
|
assert camera.fps is not None
|
||
|
|
assert camera.width is not None
|
||
|
|
assert camera.height is not None
|
||
|
|
|
||
|
|
# Test connecting twice raises an error
|
||
|
|
with pytest.raises(RobotDeviceAlreadyConnectedError):
|
||
|
|
camera.connect()
|
||
|
|
|
||
|
|
# Test reading from the camera
|
||
|
|
color_image = camera.read()
|
||
|
|
assert isinstance(color_image, np.ndarray)
|
||
|
|
assert color_image.ndim == 3
|
||
|
|
h, w, c = color_image.shape
|
||
|
|
assert c == 3
|
||
|
|
assert w > h
|
||
|
|
|
||
|
|
# Test reading asynchronously from the camera and image is similar
|
||
|
|
async_color_image = camera.async_read()
|
||
|
|
assert np.allclose(color_image, async_color_image)
|
||
|
|
|
||
|
|
# Test disconnecting
|
||
|
|
camera.disconnect()
|
||
|
|
assert camera.camera is None
|
||
|
|
assert camera.thread is None
|
||
|
|
|
||
|
|
# Test disconnecting with `__del__`
|
||
|
|
camera = OpenCVCamera(0)
|
||
|
|
camera.connect()
|
||
|
|
del camera
|
||
|
|
|
||
|
|
# Test acquiring a bgr image
|
||
|
|
camera = OpenCVCamera(0, color="bgr")
|
||
|
|
camera.connect()
|
||
|
|
assert camera.color == "bgr"
|
||
|
|
bgr_color_image = camera.read()
|
||
|
|
assert np.allclose(color_image, bgr_color_image[[2,1,0]])
|
||
|
|
del camera
|
||
|
|
|
||
|
|
# Test fps can be set
|
||
|
|
camera = OpenCVCamera(0, fps=60)
|
||
|
|
camera.connect()
|
||
|
|
assert camera.fps == 60
|
||
|
|
# TODO(rcadene): measure fps in nightly?
|
||
|
|
del camera
|
||
|
|
|
||
|
|
# Test width and height can be set
|
||
|
|
camera = OpenCVCamera(0, fps=30, width=1280, height=720)
|
||
|
|
camera.connect()
|
||
|
|
assert camera.fps == 30
|
||
|
|
assert camera.width == 1280
|
||
|
|
assert camera.height == 720
|
||
|
|
color_image = camera.read()
|
||
|
|
h, w, c = color_image.shape
|
||
|
|
assert h == 720
|
||
|
|
assert w == 1280
|
||
|
|
assert c == 3
|
||
|
|
del camera
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
|