rename color to color_mode

This commit is contained in:
Remi Cadene
2024-07-10 19:54:23 +02:00
parent 95de4b7454
commit 1bc13f39bf
2 changed files with 19 additions and 15 deletions

View File

@@ -71,7 +71,7 @@ def save_images_from_cameras(
camera = OpenCVCamera(cam_idx, fps=fps, width=width, height=height) camera = OpenCVCamera(cam_idx, fps=fps, width=width, height=height)
camera.connect() camera.connect()
print( print(
f"OpenCVCamera({camera.camera_index}, fps={camera.fps}, width={camera.width}, height={camera.height}, color={camera.color})" f"OpenCVCamera({camera.camera_index}, fps={camera.fps}, width={camera.width}, height={camera.height}, color_mode={camera.color_mode})"
) )
cameras.append(camera) cameras.append(camera)
@@ -134,11 +134,13 @@ class OpenCVCameraConfig:
fps: int | None = None fps: int | None = None
width: int | None = None width: int | None = None
height: int | None = None height: int | None = None
color: str = "rgb" color_mode: str = "rgb"
def __post_init__(self): def __post_init__(self):
if self.color not in ["rgb", "bgr"]: if self.color_mode not in ["rgb", "bgr"]:
raise ValueError(f"Expected color values are 'rgb' or 'bgr', but {self.color} is provided.") raise ValueError(
f"Expected color_mode values are 'rgb' or 'bgr', but {self.color_mode} is provided."
)
class OpenCVCamera: class OpenCVCamera:
@@ -155,7 +157,7 @@ class OpenCVCamera:
python lerobot/common/robot_devices/cameras/opencv.py --images-dir outputs/images_from_opencv_cameras python lerobot/common/robot_devices/cameras/opencv.py --images-dir outputs/images_from_opencv_cameras
``` ```
When an OpenCVCamera is instantiated, if no specific config is provided, the default fps, width, height and color When an OpenCVCamera is instantiated, if no specific config is provided, the default fps, width, height and color_mode
of the given camera will be used. of the given camera will be used.
Example of usage of the class: Example of usage of the class:
@@ -167,7 +169,7 @@ class OpenCVCamera:
camera.disconnect() camera.disconnect()
``` ```
Example of changing default fps, width, height and color: Example of changing default fps, width, height and color_mode:
```python ```python
camera = OpenCVCamera(0, fps=30, width=1280, height=720) camera = OpenCVCamera(0, fps=30, width=1280, height=720)
camera = connect() # applies the settings, might error out if these settings are not compatible with the camera camera = connect() # applies the settings, might error out if these settings are not compatible with the camera
@@ -175,7 +177,7 @@ class OpenCVCamera:
camera = OpenCVCamera(0, fps=90, width=640, height=480) camera = OpenCVCamera(0, fps=90, width=640, height=480)
camera = connect() camera = connect()
camera = OpenCVCamera(0, fps=90, width=640, height=480, color="bgr") camera = OpenCVCamera(0, fps=90, width=640, height=480, color_mode="bgr")
camera = connect() camera = connect()
``` ```
""" """
@@ -190,7 +192,7 @@ class OpenCVCamera:
self.fps = config.fps self.fps = config.fps
self.width = config.width self.width = config.width
self.height = config.height self.height = config.height
self.color = config.color self.color_mode = config.color_mode
if not isinstance(self.camera_index, int): if not isinstance(self.camera_index, int):
raise ValueError( raise ValueError(
@@ -262,7 +264,7 @@ class OpenCVCamera:
self.is_connected = True self.is_connected = True
def read(self, temporary_color: str | None = None) -> np.ndarray: def read(self, temporary_color_mode: str | None = None) -> np.ndarray:
"""Read a frame from the camera returned in the format (height, width, channels) """Read a frame from the camera returned in the format (height, width, channels)
(e.g. (640, 480, 3)), contrarily to the pytorch format which is channel first. (e.g. (640, 480, 3)), contrarily to the pytorch format which is channel first.
@@ -280,15 +282,17 @@ class OpenCVCamera:
if not ret: if not ret:
raise OSError(f"Can't capture color image from camera {self.camera_index}.") raise OSError(f"Can't capture color image from camera {self.camera_index}.")
requested_color = self.color if temporary_color is None else temporary_color requested_color_mode = self.color_mode if temporary_color_mode is None else temporary_color_mode
if requested_color not in ["rgb", "bgr"]: if requested_color_mode not in ["rgb", "bgr"]:
raise ValueError(f"Expected color values are 'rgb' or 'bgr', but {requested_color} is provided.") raise ValueError(
f"Expected color values are 'rgb' or 'bgr', but {requested_color_mode} is provided."
)
# OpenCV uses BGR format as default (blue, green red) for all operations, including displaying images. # OpenCV uses BGR format as default (blue, green red) for all operations, including displaying images.
# However, Deep Learning framework such as LeRobot uses RGB format as default to train neural networks, # However, Deep Learning framework such as LeRobot uses RGB format as default to train neural networks,
# so we convert the image color from BGR to RGB. # so we convert the image color from BGR to RGB.
if requested_color == "rgb": if requested_color_mode == "rgb":
color_image = cv2.cvtColor(color_image, cv2.COLOR_BGR2RGB) color_image = cv2.cvtColor(color_image, cv2.COLOR_BGR2RGB)
h, w, _ = color_image.shape h, w, _ = color_image.shape

View File

@@ -84,9 +84,9 @@ def test_camera(request):
del camera del camera
# Test acquiring a bgr image # Test acquiring a bgr image
camera = OpenCVCamera(CAMERA_INDEX, color="bgr") camera = OpenCVCamera(CAMERA_INDEX, color_mode="bgr")
camera.connect() camera.connect()
assert camera.color == "bgr" assert camera.color_mode == "bgr"
bgr_color_image = camera.read() bgr_color_image = camera.read()
assert np.allclose(color_image, bgr_color_image[:, :, [2, 1, 0]], rtol=1e-5, atol=MAX_PIXEL_DIFFERENCE) assert np.allclose(color_image, bgr_color_image[:, :, [2, 1, 0]], rtol=1e-5, atol=MAX_PIXEL_DIFFERENCE)
del camera del camera