"""ROS image message conversion to NumPy/OpenCV arrays."""
import base64
import time
from dataclasses import dataclass
from typing import Any, Dict, Optional
import cv2
import numpy as np
[docs]
@dataclass(frozen=True, slots=True)
class Image:
"""Decoded image data; color images use OpenCV BGR layout."""
array: np.ndarray
encoding: str = ""
format: str = ""
header: Optional[Dict[str, Any]] = None
is_bigendian: bool = False
def _message_bytes(message: Dict[str, Any]) -> bytes:
data = message.get("data")
if data is None:
raise ValueError("Image message has no data field")
if isinstance(data, str):
return base64.b64decode(data)
if isinstance(data, (bytes, bytearray)):
return bytes(data)
raise ValueError(f"Unsupported image data type: {type(data)!r}")
def compressed_message_to_frame(message: Dict[str, Any]) -> Image:
"""Decode a sensor_msgs/CompressedImage dict to an Image."""
encoded = _message_bytes(message)
buffer = np.frombuffer(encoded, dtype=np.uint8)
array = cv2.imdecode(buffer, cv2.IMREAD_COLOR)
if array is None:
image_format = message.get("format", "unknown")
raise ValueError(f"Failed to decode compressed image ({image_format})")
header = message.get("header")
return Image(
array=np.ascontiguousarray(array),
format=str(message.get("format", "")),
header=dict(header) if isinstance(header, dict) else None,
)
def image_message_to_frame(message: Dict[str, Any]) -> Image:
"""Decode a sensor_msgs/Image dict to an Image."""
raw = _message_bytes(message)
height = int(message["height"])
width = int(message["width"])
step = int(message.get("step") or 0)
encoding = str(message.get("encoding", "")).lower()
is_bigendian = bool(message.get("is_bigendian", 0))
depth = _depth_array(raw, height, width, step, encoding, is_bigendian)
if depth is not None:
header = message.get("header")
return Image(
array=np.ascontiguousarray(depth),
encoding=str(message.get("encoding", "")),
header=dict(header) if isinstance(header, dict) else None,
is_bigendian=is_bigendian,
)
channels = _channel_count(encoding, raw, height, width, step)
row_bytes = step if step > 0 else width * channels
expected = row_bytes * height
if len(raw) < expected:
raise ValueError(
f"Image data too short for {width}x{height} ({encoding}): "
f"expected at least {expected} bytes, got {len(raw)}"
)
buffer = np.frombuffer(raw, dtype=np.uint8)
if row_bytes == width * channels:
array = buffer[: expected].reshape(height, width, channels)
else:
rows = [buffer[y * row_bytes : y * row_bytes + width * channels] for y in range(height)]
array = np.vstack(rows).reshape(height, width, channels)
array = _normalize_encoding(array, encoding)
header = message.get("header")
return Image(
array=np.ascontiguousarray(array),
encoding=str(message.get("encoding", "")),
header=dict(header) if isinstance(header, dict) else None,
is_bigendian=is_bigendian,
)
def _depth_array(
raw: bytes,
height: int,
width: int,
step: int,
encoding: str,
is_bigendian: bool,
) -> np.ndarray | None:
"""Decode supported depth encodings while preserving their native dtype."""
if encoding in {"16uc1", "mono16"}:
dtype = np.dtype(">u2" if is_bigendian else "<u2")
itemsize = 2
elif encoding == "32fc1":
dtype = np.dtype(">f4" if is_bigendian else "<f4")
itemsize = 4
else:
return None
row_bytes = step if step > 0 else width * itemsize
expected = row_bytes * height
if len(raw) < expected:
raise ValueError(
f"Image data too short for {width}x{height} ({encoding}): "
f"expected at least {expected} bytes, got {len(raw)}"
)
rows = [
np.frombuffer(raw[y * row_bytes : y * row_bytes + width * itemsize], dtype=dtype, count=width)
for y in range(height)
]
native = np.vstack(rows)
return native.astype(dtype.newbyteorder("="), copy=False)
def _channel_count(encoding: str, raw: bytes, height: int, width: int, step: int) -> int:
if encoding in {"mono8", "8uc1"}:
return 1
if encoding in {"bgr8", "rgb8", "8uc3"}:
return 3
if encoding in {"bgra8", "rgba8"}:
return 4
if step > 0 and height > 0:
return max(1, step // width)
if width > 0 and height > 0 and len(raw) >= height * width:
return max(1, len(raw) // (height * width))
return 3
def _normalize_encoding(array: np.ndarray, encoding: str) -> np.ndarray:
if encoding == "rgb8":
return cv2.cvtColor(array, cv2.COLOR_RGB2BGR)
if encoding in {"mono8", "8uc1"}:
if array.ndim == 3 and array.shape[2] == 1:
array = array[:, :, 0]
return cv2.cvtColor(array, cv2.COLOR_GRAY2BGR)
if encoding == "bgra8":
return cv2.cvtColor(array, cv2.COLOR_BGRA2BGR)
if encoding == "rgba8":
return cv2.cvtColor(array, cv2.COLOR_RGBA2BGR)
if array.ndim == 3 and array.shape[2] == 1:
return cv2.cvtColor(array[:, :, 0], cv2.COLOR_GRAY2BGR)
if array.ndim == 2:
return cv2.cvtColor(array, cv2.COLOR_GRAY2BGR)
return array
# Maps a raw ROS encoding to the OpenCV BGR->target conversion (None means no
# conversion) and the resulting channel count. Image.array is BGR by convention,
# so encoding an Image reverses the decode-side normalization.
_ENCODE_CONVERSIONS: Dict[str, tuple] = {
"bgr8": (None, 3),
"rgb8": (cv2.COLOR_BGR2RGB, 3),
"mono8": (cv2.COLOR_BGR2GRAY, 1),
"bgra8": (cv2.COLOR_BGR2BGRA, 4),
"rgba8": (cv2.COLOR_BGR2RGBA, 4),
}
def _message_header(image: "Image") -> Dict[str, Any]:
"""Return the image header, auto-stamping with the current time when absent."""
if image.header is not None:
return dict(image.header)
now = time.time()
sec = int(now)
nanosec = int(round((now - sec) * 1e9))
return {"stamp": {"sec": sec, "nanosec": nanosec}, "frame_id": ""}
def image_to_message(image: "Image") -> Dict[str, Any]:
"""Encode an Image (BGR array) to a sensor_msgs/Image dict.
The target ROS encoding comes from ``image.encoding`` (default ``bgr8``).
Pixel data remains bytes for the native Core binary payload.
"""
encoding = image.encoding or "bgr8"
depth_message = _depth_image_to_message(image, encoding)
if depth_message is not None:
return depth_message
conversion = _ENCODE_CONVERSIONS.get(encoding.lower())
if conversion is None:
raise ValueError(f"Unsupported image encoding for publish: {encoding!r}")
convert_code, channels = conversion
array = np.ascontiguousarray(image.array, dtype=np.uint8)
if convert_code is not None:
array = cv2.cvtColor(array, convert_code)
array = np.ascontiguousarray(array, dtype=np.uint8)
height, width = int(array.shape[0]), int(array.shape[1])
return {
"header": _message_header(image),
"height": height,
"width": width,
"encoding": encoding,
"is_bigendian": 1 if image.is_bigendian else 0,
"step": width * channels,
"data": array.tobytes(),
}
def _depth_image_to_message(image: "Image", encoding: str) -> Dict[str, Any] | None:
"""Encode native depth arrays for supported ROS encodings."""
lower = encoding.lower()
if lower not in {"16uc1", "mono16", "32fc1"}:
return None
if image.array.ndim != 2:
raise ValueError(f"Depth image encoding {encoding!r} requires a 2D array")
if lower in {"16uc1", "mono16"}:
dtype = np.dtype(">u2" if image.is_bigendian else "<u2")
else:
dtype = np.dtype(">f4" if image.is_bigendian else "<f4")
array = np.ascontiguousarray(image.array.astype(dtype, copy=False))
height, width = int(array.shape[0]), int(array.shape[1])
return {
"header": _message_header(image),
"height": height,
"width": width,
"encoding": encoding,
"is_bigendian": 1 if image.is_bigendian else 0,
"step": width * dtype.itemsize,
"data": array.tobytes(),
}
def compressed_image_to_message(image: "Image") -> Dict[str, Any]:
"""Encode an Image (BGR array) to a sensor_msgs/CompressedImage dict.
The compression format comes from ``image.format``: PNG when it contains
"png", otherwise JPEG (the default when ``format`` is empty). Encoded data
remains bytes for the native Core binary payload.
"""
fmt = image.format or "jpeg"
if "png" in fmt.lower():
ext = ".png"
else:
ext = ".jpg"
array = np.ascontiguousarray(image.array, dtype=np.uint8)
ok, encoded = cv2.imencode(ext, array)
if not ok:
raise ValueError(f"Failed to encode compressed image ({fmt})")
return {
"header": _message_header(image),
"format": fmt,
"data": encoded.tobytes(),
}
__all__ = [
"Image",
"compressed_image_to_message",
"compressed_message_to_frame",
"image_message_to_frame",
"image_to_message",
]