cristhian
a6fa3ce32b
dora's typed image is an Arrow struct of width, height, encoding and data -- it carries inside the payload what this node has been reading out of message metadata. The decoder now takes either shape: a struct describes itself and its metadata is ignored, a flat UInt8 buffer is described by width/height/ encoding beside it, and everything after unpacking is common. The declared contract stays std/core/v1/UInt8, deliberately. Type compatibility is directional -- Image satisfies a Bytes input, but a std/core/v1/UInt8 output wired into an std/media/v1/Image input is a validate error -- and dora-zed declares image_bgr as std/core/v1/UInt8. Flipping this manifest now would break the camera. Accepting both shapes first is what lets the producers move before the contract does. output_framing is not part of this. It selects the wire format, raw or arrow-ipc, and is not required for images: a struct payload was verified to arrive identically under both framings, and validate does not ask for arrow-ipc when an output is annotated std/media/v1/Image. It earns its keep when the wire must be self-describing, such as crossing languages or recording.
119 lines
4.4 KiB
Python
119 lines
4.4 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
|
|
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"}
|
|
|
|
|
|
# The fields of std/media/v1/Image. A payload carrying them describes itself and
|
|
# its metadata is not consulted.
|
|
IMAGE_FIELDS = ("width", "height", "encoding", "data")
|
|
|
|
|
|
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]
|