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,72 @@
|
||||
"""The display state: the frame, the overlays drawn on it, and the window."""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
WINDOW = "Dora Node: opencv-plot"
|
||||
GREEN = (0, 255, 0)
|
||||
WHITE = (255, 255, 255)
|
||||
FONT = cv2.FONT_HERSHEY_SIMPLEX
|
||||
FONT_SCALE = 0.5
|
||||
|
||||
|
||||
@dataclass
|
||||
class Bboxes:
|
||||
"""One overlay's worth of boxes, always stored as xyxy corners."""
|
||||
|
||||
corners: np.ndarray = field(default_factory=lambda: np.empty((0, 4)))
|
||||
conf: np.ndarray = field(default_factory=lambda: np.empty(0))
|
||||
labels: np.ndarray = field(default_factory=lambda: np.empty(0, dtype=object))
|
||||
|
||||
|
||||
@dataclass
|
||||
class Plot:
|
||||
"""Overlay state, plus the window it is displayed in.
|
||||
|
||||
`headless` suppresses the window; the node still decodes every input.
|
||||
"""
|
||||
|
||||
width: int | None = None
|
||||
height: int | None = None
|
||||
headless: bool = False
|
||||
bboxes: Bboxes = field(default_factory=Bboxes)
|
||||
text: str = ""
|
||||
|
||||
def show(self, frame: np.ndarray) -> bool:
|
||||
"""Draw the overlays on `frame`, resize it, and display it.
|
||||
|
||||
Returns whether the viewer asked to quit by pressing `q`.
|
||||
"""
|
||||
self._draw(frame)
|
||||
|
||||
if self.width and self.height:
|
||||
frame = cv2.resize(frame, (self.width, self.height))
|
||||
|
||||
if self.headless:
|
||||
return False
|
||||
|
||||
cv2.imshow(WINDOW, frame)
|
||||
return cv2.waitKey(1) & 0xFF == ord("q")
|
||||
|
||||
def _draw(self, frame: np.ndarray) -> None:
|
||||
for corners, conf, label in zip(
|
||||
self.bboxes.corners,
|
||||
self.bboxes.conf,
|
||||
self.bboxes.labels,
|
||||
strict=False,
|
||||
):
|
||||
min_x, min_y, max_x, max_y = (int(value) for value in corners)
|
||||
cv2.rectangle(frame, (min_x, min_y), (max_x, max_y), GREEN, 2)
|
||||
cv2.putText(
|
||||
frame,
|
||||
f"{label}, {conf:0.2f}",
|
||||
(max_x - 120, max_y - 10),
|
||||
FONT,
|
||||
FONT_SCALE,
|
||||
GREEN,
|
||||
1,
|
||||
)
|
||||
|
||||
cv2.putText(frame, self.text, (20, 20), FONT, FONT_SCALE, WHITE, 1)
|
||||
Reference in New Issue
Block a user