add opencv-plot node for dora dataflows
A display sink that overlays bounding boxes and text on an image input and shows it in an OpenCV window. Based on opencv-plot from dora-hub, brought up to the dora 1.0 node contract and packaged with uv. Python rather than Rust: the opencv Rust crate generates its bindings from the OpenCV headers on the build machine, so it needs libopencv-dev and libclang-dev per machine and per architecture and does not cross-compile. Every dependency here ships prebuilt wheels for linux-x86_64, linux-aarch64 and macos-arm64. Covers every encoding the upstream node did — bgr8, rgb8, jpeg/jpg/jpe/bmp/ webp/png, yuv420 and avif — plus both bbox formats, the text overlay, PLOT_WIDTH/PLOT_HEIGHT and headless operation under CI=true.
This commit is contained in:
@@ -0,0 +1,75 @@
|
||||
"""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]
|
||||
Reference in New Issue
Block a user