cristhian
00dc4c7b71
The repository, the Python package, the entrypoint and the manifest name all now agree. `plot` was carried over from the upstream opencv-plot node and no longer described anything here, so the module, the class and the environment variables follow the node: dora_plot/ -> dora_image_view/ dora_plot/plot.py -> dora_image_view/view.py Plot -> View PLOT_WIDTH/HEIGHT -> VIEW_WIDTH/HEIGHT --plot-width/-height -> --view-width/-height Breaking for anyone wiring `path: dora-plot` or setting PLOT_WIDTH; nothing is tagged yet.
76 lines
2.9 KiB
Python
76 lines
2.9 KiB
Python
"""Decoding of the `image` input into a BGR frame that can be drawn on."""
|
|
|
|
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"}
|
|
|
|
|
|
def to_bgr(storage, encoding: str, width: int | None, height: int | None) -> np.ndarray:
|
|
"""Decode one `image` payload into a writable BGR array.
|
|
|
|
`width` and `height` come from the message metadata and are only required by
|
|
the raw encodings, which carry no dimensions of their own.
|
|
"""
|
|
buffer = storage.to_numpy().astype(np.uint8)
|
|
|
|
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 _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]
|