cristhian
8a7bda5867
Both inputs now declare dora's Arrow struct {width, height, encoding, data}
rather than a flat buffer plus loose metadata. The payload describes itself, so
nothing about its shape has to be agreed out of band, and the frame travels as
one Arrow value instead of a buffer whose meaning lives in a sidecar dict.
Depth uses the same type. dora's registry has no depth type, and a depth map is
an image whose pixels are distances -- the reading ROS takes as well -- so it
arrives with a single-channel encoding: 32FC1/float32 or 16UC1/mono16/uint16.
`units` stays in message metadata because the struct has no field for it.
decode.IMAGE_TYPE is the single definition of the struct; the example source
and the tests build payloads from it rather than respelling the schema.
The older shapes are still decoded -- a flat UInt8 buffer with
width/height/encoding for image, a flat Float32 buffer with shape/units for
depth -- so a producer's payload can move before its manifest does. That is
only half a migration path: type compatibility is directional, so a producer
still declaring std/core/v1/UInt8 fails validate against an
std/media/v1/Image input regardless of what it sends, and has to update its
manifest too.
Verified: 41 tests, validate clean with 4 contracts injected, the dataflow runs
to a clean exit, and a composite rendered through the struct path is identical
to the one rendered before the migration.
127 lines
4.7 KiB
Python
127 lines
4.7 KiB
Python
"""Decoding of the `image` input into a BGR frame that can be drawn on.
|
|
|
|
Two payload shapes are accepted, because dora has both in the field. The typed
|
|
one is `std/media/v1/Image`: an Arrow struct carrying width, height, encoding
|
|
and data together, so the frame is self-describing. The older one is a flat
|
|
UInt8 buffer whose dimensions and encoding travel beside it in the message
|
|
metadata, which is what a node declaring `std/core/v1/UInt8` sends -- dora-zed
|
|
among them. Everything after unpacking is identical.
|
|
"""
|
|
|
|
import io
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import pillow_avif # noqa: F401 # registers the AVIF plugin with Pillow
|
|
import pyarrow as pa
|
|
from PIL import Image
|
|
|
|
# Encodings that carry their own dimensions, so `width`/`height` metadata is not
|
|
# consulted for them.
|
|
SELF_DESCRIBING = frozenset({"jpeg", "jpg", "jpe", "bmp", "webp", "png"})
|
|
|
|
SUPPORTED = SELF_DESCRIBING | {"bgr8", "rgb8", "yuv420", "avif"}
|
|
|
|
|
|
# std/media/v1/Image. A payload carrying these fields describes itself, so its
|
|
# message metadata is not consulted. Exported as an Arrow type so a producer --
|
|
# the example source, the tests -- builds it from one definition.
|
|
IMAGE_TYPE = pa.struct([
|
|
("width", pa.uint32()),
|
|
("height", pa.uint32()),
|
|
("encoding", pa.utf8()),
|
|
("data", pa.large_binary()),
|
|
])
|
|
IMAGE_FIELDS = tuple(field.name for field in IMAGE_TYPE)
|
|
|
|
|
|
def to_bgr(storage, metadata: dict) -> np.ndarray:
|
|
"""Decode one `image` payload into a writable BGR array.
|
|
|
|
Accepts either an `std/media/v1/Image` struct or a flat buffer described by
|
|
the message metadata.
|
|
"""
|
|
buffer, encoding, width, height = unpack(storage, metadata)
|
|
|
|
if encoding == "bgr8":
|
|
# Copy so the frame can be annotated: to_numpy() may be a view on the
|
|
# shared-memory payload.
|
|
return _reshape(buffer, width, height, encoding).copy()
|
|
|
|
if encoding == "rgb8":
|
|
return cv2.cvtColor(_reshape(buffer, width, height, encoding), cv2.COLOR_RGB2BGR)
|
|
|
|
if encoding == "yuv420":
|
|
# A planar YUV 4:2:0 frame is 1.5 bytes per pixel, laid out as a
|
|
# single-channel image of `height * 3 // 2` rows.
|
|
width, height = _dimensions(width, height, encoding)
|
|
plane = _take(buffer, width * height * 3 // 2, encoding).reshape(height * 3 // 2, width)
|
|
# COLOR_YUV420p2RGB, not ...2BGR: the upstream node converts to RGB and
|
|
# then displays it as if it were BGR. The channel swap is part of the
|
|
# behaviour callers already see.
|
|
return cv2.cvtColor(plane, cv2.COLOR_YUV420p2RGB)
|
|
|
|
if encoding in SELF_DESCRIBING:
|
|
frame = cv2.imdecode(buffer, cv2.IMREAD_COLOR)
|
|
if frame is None:
|
|
raise RuntimeError(f"could not decode {encoding} image")
|
|
return frame
|
|
|
|
if encoding == "avif":
|
|
image = Image.open(io.BytesIO(buffer.tobytes())).convert("RGB")
|
|
return cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
|
|
|
|
raise RuntimeError(f"unsupported image encoding: {encoding}")
|
|
|
|
|
|
def unpack(storage, metadata: dict) -> tuple[np.ndarray, str, int | None, int | None]:
|
|
"""Split a payload of either shape into buffer, encoding, width and height."""
|
|
if is_image_struct(storage):
|
|
# One image per message: the struct is a single-row array.
|
|
row = storage[0].as_py()
|
|
return (
|
|
np.frombuffer(row["data"], dtype=np.uint8),
|
|
str(row["encoding"]).lower(),
|
|
row["width"],
|
|
row["height"],
|
|
)
|
|
|
|
return (
|
|
storage.to_numpy().astype(np.uint8),
|
|
str(metadata.get("encoding", "bgr8")).lower(),
|
|
metadata.get("width"),
|
|
metadata.get("height"),
|
|
)
|
|
|
|
|
|
def is_image_struct(storage) -> bool:
|
|
"""Whether the payload is an Arrow struct with the std/media/v1/Image fields."""
|
|
fields = getattr(getattr(storage, "type", None), "fields", None)
|
|
|
|
if fields is None:
|
|
return False
|
|
|
|
return all(name in {field.name for field in fields} for name in IMAGE_FIELDS)
|
|
|
|
|
|
def _reshape(buffer: np.ndarray, width: int | None, height: int | None, encoding: str):
|
|
width, height = _dimensions(width, height, encoding)
|
|
return _take(buffer, width * height * 3, encoding).reshape(height, width, 3)
|
|
|
|
|
|
def _dimensions(width: int | None, height: int | None, encoding: str) -> tuple[int, int]:
|
|
if not width or not height:
|
|
raise RuntimeError(
|
|
f"the {encoding} encoding needs positive `width` and `height` metadata",
|
|
)
|
|
return width, height
|
|
|
|
|
|
def _take(buffer: np.ndarray, expected: int, encoding: str) -> np.ndarray:
|
|
if buffer.size < expected:
|
|
raise RuntimeError(
|
|
f"{encoding} payload is {buffer.size} bytes, "
|
|
f"expected at least {expected} for the declared size",
|
|
)
|
|
return buffer[:expected]
|