cesar
43cb1ff01c
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.
136 lines
3.7 KiB
Python
136 lines
3.7 KiB
Python
"""Tests that need neither a dataflow nor a display."""
|
|
|
|
import io
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import pytest
|
|
from PIL import Image
|
|
|
|
from opencv_plot import bbox, decode
|
|
from opencv_plot.plot import Bboxes, Plot
|
|
|
|
WIDTH, HEIGHT = 12, 8
|
|
|
|
|
|
class FakeStorage:
|
|
"""Stands in for the Arrow array dora hands to the node."""
|
|
|
|
def __init__(self, data: np.ndarray) -> None:
|
|
self._data = data
|
|
|
|
def to_numpy(self) -> np.ndarray:
|
|
return self._data
|
|
|
|
|
|
@pytest.fixture
|
|
def frame() -> np.ndarray:
|
|
rng = np.random.default_rng(seed=0)
|
|
return rng.integers(0, 255, (HEIGHT, WIDTH, 3), dtype=np.uint8)
|
|
|
|
|
|
def test_bgr8_round_trips_unchanged(frame):
|
|
decoded = decode.to_bgr(FakeStorage(frame.ravel()), "bgr8", WIDTH, HEIGHT)
|
|
|
|
np.testing.assert_array_equal(decoded, frame)
|
|
|
|
|
|
def test_rgb8_round_trips_unchanged(frame):
|
|
rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
|
|
|
decoded = decode.to_bgr(FakeStorage(rgb.ravel()), "rgb8", WIDTH, HEIGHT)
|
|
|
|
np.testing.assert_array_equal(decoded, frame)
|
|
|
|
|
|
@pytest.mark.parametrize(("suffix", "encoding"), [
|
|
(".png", "png"),
|
|
(".jpg", "jpeg"),
|
|
(".bmp", "bmp"),
|
|
(".webp", "webp"),
|
|
])
|
|
def test_self_describing_encodings_need_no_dimensions(frame, suffix, encoding):
|
|
ok, buffer = cv2.imencode(suffix, frame)
|
|
assert ok
|
|
|
|
decoded = decode.to_bgr(FakeStorage(buffer.ravel()), encoding, None, None)
|
|
|
|
assert decoded.shape == (HEIGHT, WIDTH, 3)
|
|
|
|
|
|
def test_yuv420_is_decoded(frame):
|
|
planes = cv2.cvtColor(frame, cv2.COLOR_BGR2YUV_I420)
|
|
|
|
decoded = decode.to_bgr(FakeStorage(planes.ravel()), "yuv420", WIDTH, HEIGHT)
|
|
|
|
assert decoded.shape == (HEIGHT, WIDTH, 3)
|
|
|
|
|
|
def test_avif_is_decoded(frame):
|
|
buffer = io.BytesIO()
|
|
Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)).save(buffer, format="AVIF")
|
|
payload = np.frombuffer(buffer.getvalue(), dtype=np.uint8)
|
|
|
|
decoded = decode.to_bgr(FakeStorage(payload), "avif", None, None)
|
|
|
|
assert decoded.shape == (HEIGHT, WIDTH, 3)
|
|
|
|
|
|
def test_decoded_frame_is_writable():
|
|
storage = FakeStorage(np.arange(2 * 3 * 3, dtype=np.uint8))
|
|
|
|
decoded = decode.to_bgr(storage, "bgr8", width=3, height=2)
|
|
|
|
# The frame is annotated in place, so it must not be a read-only view.
|
|
decoded[0, 0] = 0
|
|
|
|
|
|
def test_raw_encodings_need_dimensions():
|
|
storage = FakeStorage(np.zeros(18, dtype=np.uint8))
|
|
with pytest.raises(RuntimeError, match="width"):
|
|
decode.to_bgr(storage, "bgr8", width=None, height=None)
|
|
|
|
|
|
def test_short_payload_is_rejected():
|
|
storage = FakeStorage(np.zeros(4, dtype=np.uint8))
|
|
with pytest.raises(RuntimeError, match="expected at least"):
|
|
decode.to_bgr(storage, "bgr8", width=3, height=2)
|
|
|
|
|
|
def test_unsupported_encoding_is_rejected():
|
|
storage = FakeStorage(np.zeros(4, dtype=np.uint8))
|
|
with pytest.raises(RuntimeError, match="unsupported image encoding"):
|
|
decode.to_bgr(storage, "tiff", width=None, height=None)
|
|
|
|
|
|
def test_xywh_is_converted_to_corners():
|
|
import pyarrow as pa
|
|
|
|
value = pa.array(
|
|
[
|
|
{
|
|
"bbox": np.array([10.0, 20.0, 4.0, 6.0]),
|
|
"conf": np.array([0.5]),
|
|
"labels": np.array(["cat"]),
|
|
},
|
|
],
|
|
)
|
|
|
|
boxes = bbox.parse(value, "xywh")
|
|
|
|
np.testing.assert_allclose(boxes.corners, [[8.0, 17.0, 12.0, 23.0]])
|
|
assert list(boxes.labels) == ["cat"]
|
|
|
|
|
|
def test_headless_show_draws_without_a_window():
|
|
plot = Plot(headless=True, text="hello")
|
|
plot.bboxes = Bboxes(
|
|
corners=np.array([[1.0, 1.0, 5.0, 5.0]]),
|
|
conf=np.array([0.9]),
|
|
labels=np.array(["cat"]),
|
|
)
|
|
frame = np.zeros((32, 32, 3), dtype=np.uint8)
|
|
|
|
assert plot.show(frame) is False
|
|
assert frame.any(), "the overlays should have marked the frame"
|