refactor: rename the node from opencv-plot to dora-plot
The repository is dora-plot, but the package, the entrypoint, the manifest name and the window title were all opencv-plot. Consumers wire the node by `path:`, so the split name was visible in every dataflow. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_018WLVZCFFcCH1EpTXGfS9aD
This commit is contained in:
Claude Opus 5
@@ -0,0 +1 @@
|
||||
"""Dora node that overlays bounding boxes and text on an image and displays it."""
|
||||
@@ -0,0 +1,37 @@
|
||||
"""Parsing of the `bbox` input: an Arrow struct carrying boxes and labels."""
|
||||
|
||||
import numpy as np
|
||||
|
||||
from dora_plot.plot import Bboxes
|
||||
|
||||
|
||||
def parse(value, bbox_format: str) -> Bboxes:
|
||||
"""Read a `bbox` payload. `bbox_format` is the `format` metadata parameter."""
|
||||
fields = value[0]
|
||||
corners = _flat(fields, "bbox").reshape(-1, 4)
|
||||
|
||||
if bbox_format == "xyxy":
|
||||
pass
|
||||
elif bbox_format == "xywh":
|
||||
# xywh is centre-based: x and y are the box centre, not a corner.
|
||||
x, y, w, h = corners.T
|
||||
corners = np.stack([x - w / 2, y - h / 2, x + w / 2, y + h / 2], axis=1)
|
||||
else:
|
||||
raise RuntimeError(f"unsupported bbox format: {bbox_format}")
|
||||
|
||||
return Bboxes(
|
||||
corners=corners,
|
||||
conf=_flat(fields, "conf"),
|
||||
labels=_flat(fields, "labels", numeric=False),
|
||||
)
|
||||
|
||||
|
||||
def _flat(fields, name: str, numeric: bool = True) -> np.ndarray:
|
||||
"""Read a struct field as a flat numpy array.
|
||||
|
||||
Senders wrap each field in a single-element list -- that is what
|
||||
`pa.array([{...}])` produces from numpy arrays -- so unwrap it when present.
|
||||
"""
|
||||
column = fields[name]
|
||||
column = getattr(column, "values", column)
|
||||
return column.to_numpy(zero_copy_only=False) if not numeric else column.to_numpy()
|
||||
@@ -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]
|
||||
@@ -0,0 +1,84 @@
|
||||
"""A dora display sink: overlays bounding boxes and text on an image input.
|
||||
|
||||
An `image` event decodes the frame and immediately redraws. `bbox` and `text`
|
||||
events only update stored state, which the next `image` event draws.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
|
||||
from dora import Node
|
||||
|
||||
from dora_plot import bbox, decode
|
||||
from dora_plot.plot import Plot
|
||||
|
||||
|
||||
def main() -> None:
|
||||
"""Connect to the dataflow and display every frame that arrives."""
|
||||
args = _parse_args()
|
||||
|
||||
plot = Plot(
|
||||
width=_size("PLOT_WIDTH", args.plot_width),
|
||||
height=_size("PLOT_HEIGHT", args.plot_height),
|
||||
headless=os.getenv("CI") == "true",
|
||||
)
|
||||
|
||||
# Passing the name lets the node also run standalone as a dynamic node; when
|
||||
# dora spawns it, the id from the dataflow wins.
|
||||
node = Node(args.name)
|
||||
|
||||
for event in node:
|
||||
if event["type"] == "ERROR":
|
||||
raise RuntimeError(event["error"])
|
||||
if event["type"] != "INPUT":
|
||||
continue
|
||||
|
||||
event_id = event["id"]
|
||||
metadata = event["metadata"]
|
||||
|
||||
if event_id == "image":
|
||||
frame = decode.to_bgr(
|
||||
event["value"],
|
||||
metadata.get("encoding", "bgr8").lower(),
|
||||
metadata.get("width"),
|
||||
metadata.get("height"),
|
||||
)
|
||||
if plot.show(frame):
|
||||
break
|
||||
elif event_id == "bbox":
|
||||
plot.bboxes = bbox.parse(event["value"], metadata.get("format", "xyxy").lower())
|
||||
elif event_id == "text":
|
||||
plot.text = event["value"][0].as_py()
|
||||
else:
|
||||
print(f"ignoring unexpected input `{event_id}`")
|
||||
|
||||
|
||||
def _parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="OpenCV plotter: draws text and bounding boxes on an image.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--name",
|
||||
default="dora-plot",
|
||||
help="The name of the node in the dataflow, used when running as a dynamic node.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--plot-width",
|
||||
type=int,
|
||||
help="Resize the displayed frame to this width. Overridden by PLOT_WIDTH.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--plot-height",
|
||||
type=int,
|
||||
help="Resize the displayed frame to this height. Overridden by PLOT_HEIGHT.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def _size(name: str, fallback: int | None) -> int | None:
|
||||
value = os.getenv(name)
|
||||
return int(value) if value and value.isnumeric() else fallback
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -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: dora-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