First commit

This commit is contained in:
Cristhian Aguilera
2026-01-28 16:21:37 -03:00
parent 2da5baf508
commit 610c43e16d
17 changed files with 2113 additions and 2 deletions

View File

@@ -0,0 +1,95 @@
# Dora Node for plotting data with OpenCV
This node is used to plot a text and a list of bbox on a base image (ideal for object detection).
# YAML
```yaml
- id: dora-image-viewer
build: pip install ../../dora_image_viewer
path: dora_image_viewer/dora_image_viewer/main.py
inputs:
# image: Arrow array of size 1 containing the base image
# bbox: Arrow array of bbox
# text: Arrow array of size 1 containing the text to be plotted
env:
PLOT_WIDTH: 640 # optional, default is image input width
PLOT_HEIGHT: 480 # optional, default is image input height
```
# Inputs
- `image`: Arrow array containing the base image
```python
## Image data
image_data: UInt8Array # Example: pa.array(img.ravel())
metadata = {
"width": 640,
"height": 480,
"shape": [480, 640, 3], # optional alternative to width/height
"encoding": str, # bgr8, rgb8
}
## Example
node.send_output(
image_data, {"width": 640, "height": 480, "encoding": "bgr8"}
)
## Decoding
storage = event["value"]
metadata = event["metadata"]
encoding = metadata["encoding"]
width = metadata.get("width")
height = metadata.get("height")
shape = metadata.get("shape")
if encoding == "bgr8":
channels = 3
storage_type = np.uint8
frame = (
storage.to_numpy()
.astype(storage_type)
.reshape((height, width, channels))
)
```
- `bbox`: an arrow array containing the bounding boxes, confidence scores, and class names of the detected objects
```Python
bbox: {
"bbox": np.array, # flattened array of bounding boxes
"conf": np.array, # flat array of confidence scores
"labels": np.array, # flat array of class names
}
encoded_bbox = pa.array([bbox], {"format": "xyxy"})
decoded_bbox = {
"bbox": encoded_bbox[0]["bbox"].values.to_numpy().reshape(-1, 4),
"conf": encoded_bbox[0]["conf"].values.to_numpy(),
"labels": encoded_bbox[0]["labels"].values.to_numpy(zero_copy_only=False),
}
```
- `text`: Arrow array containing the text to be plotted
```python
text: str
encoded_text = pa.array([text])
decoded_text = encoded_text[0].as_py()
```
## Example
Check example at [examples/python-dataflow](examples/python-dataflow)
## License
This project is licensed under Apache-2.0. Check out [NOTICE.md](../../NOTICE.md) for more information.

View File

@@ -0,0 +1,13 @@
"""TODO: Add docstring."""
import os
# Define the path to the README file relative to the package directory
readme_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), "README.md")
# Read the content of the README file
try:
with open(readme_path, encoding="utf-8") as f:
__doc__ = f.read()
except FileNotFoundError:
__doc__ = "README file not found."

View File

@@ -0,0 +1,250 @@
"""TODO: Add docstring."""
import argparse
import io
import os
import cv2
import numpy as np
import pyarrow as pa
from dora import Node
from PIL import (
Image,
)
if True:
import pillow_avif # noqa # noqa
RUNNER_CI = True if os.getenv("CI") == "true" else False
class Plot:
"""TODO: Add docstring."""
frame: np.array = np.array([])
bboxes: dict = {
"bbox": np.array([]),
"conf": np.array([]),
"labels": np.array([]),
}
text: str = ""
width: np.uint32 = None
height: np.uint32 = None
def plot_frame(plot):
"""TODO: Add docstring."""
for bbox in zip(plot.bboxes["bbox"], plot.bboxes["conf"], plot.bboxes["labels"]):
[
[min_x, min_y, max_x, max_y],
confidence,
label,
] = bbox
cv2.rectangle(
plot.frame,
(int(min_x), int(min_y)),
(int(max_x), int(max_y)),
(0, 255, 0),
2,
)
cv2.putText(
plot.frame,
f"{label}, {confidence:0.2f}",
(int(max_x) - 120, int(max_y) - 10),
cv2.FONT_HERSHEY_SIMPLEX,
0.5,
(0, 255, 0),
1,
1,
)
cv2.putText(
plot.frame,
plot.text,
(20, 20),
cv2.FONT_HERSHEY_SIMPLEX,
0.5,
(255, 255, 255),
1,
1,
)
if plot.width is not None and plot.height is not None:
plot.frame = cv2.resize(plot.frame, (plot.width, plot.height))
if not RUNNER_CI:
if len(plot.frame.shape) >= 3:
cv2.imshow("Dora Node: dora-image-viewer", plot.frame)
def yuv420p_to_bgr_opencv(yuv_array, width, height):
"""TODO: Add docstring."""
yuv_array = yuv_array[: width * height * 3 // 2]
yuv = yuv_array.reshape((height * 3 // 2, width))
return cv2.cvtColor(yuv, cv2.COLOR_YUV420p2RGB)
def main():
# Handle dynamic nodes, ask for the name of the node in the dataflow, and the same values as the ENV variables.
"""TODO: Add docstring."""
parser = argparse.ArgumentParser(
description="OpenCV Plotter: This node is used to plot text and bounding boxes on an image.",
)
parser.add_argument(
"--name",
type=str,
required=False,
help="The name of the node in the dataflow.",
default="dora-image-viewer",
)
parser.add_argument(
"--plot-width",
type=int,
required=False,
help="The width of the plot.",
default=None,
)
parser.add_argument(
"--plot-height",
type=int,
required=False,
help="The height of the plot.",
default=None,
)
args = parser.parse_args()
plot_width = os.getenv("PLOT_WIDTH", args.plot_width)
plot_height = os.getenv("PLOT_HEIGHT", args.plot_height)
if plot_width is not None:
if isinstance(plot_width, str) and plot_width.isnumeric():
plot_width = int(plot_width)
if plot_height is not None:
if isinstance(plot_height, str) and plot_height.isnumeric():
plot_height = int(plot_height)
node = Node(
args.name,
) # provide the name to connect to the dataflow if dynamic node
plot = Plot()
plot.width = plot_width
plot.height = plot_height
pa.array([]) # initialize pyarrow array
for event in node:
event_type = event["type"]
if event_type == "INPUT":
event_id = event["id"]
if event_id == "image":
storage = event["value"]
metadata = event["metadata"]
encoding = metadata["encoding"].lower()
width = metadata.get("width")
height = metadata.get("height")
if (width is None or height is None) and "shape" in metadata:
shape = metadata["shape"]
if isinstance(shape, (list, tuple)) and len(shape) >= 2:
height = height if height is not None else int(shape[0])
width = width if width is not None else int(shape[1])
if width is None or height is None:
raise KeyError("width/height (or shape) missing from metadata")
if encoding == "bgr8":
channels = 3
storage_type = np.uint8
plot.frame = (
storage.to_numpy()
.astype(storage_type)
.reshape((height, width, channels))
.copy() # Copy So that we can add annotation on the image
)
elif encoding == "rgb8":
channels = 3
storage_type = np.uint8
frame = (
storage.to_numpy()
.astype(storage_type)
.reshape((height, width, channels))
)
plot.frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
elif encoding in ["jpeg", "jpg", "jpe", "bmp", "webp", "png"]:
channels = 3
storage_type = np.uint8
storage = storage.to_numpy()
plot.frame = cv2.imdecode(storage, cv2.IMREAD_COLOR)
elif encoding == "yuv420":
storage = storage.to_numpy()
# Convert back to BGR results in more saturated image.
channels = 3
storage_type = np.uint8
img_bgr_restored = yuv420p_to_bgr_opencv(storage, width, height)
plot.frame = img_bgr_restored
elif encoding == "avif":
# Convert AVIF to RGB
array = storage.to_numpy()
bytes = array.tobytes()
img = Image.open(io.BytesIO(bytes))
img = img.convert("RGB")
plot.frame = np.array(img)
plot.frame = cv2.cvtColor(plot.frame, cv2.COLOR_RGB2BGR)
else:
raise RuntimeError(f"Unsupported image encoding: {encoding}")
plot_frame(plot)
if not RUNNER_CI:
if cv2.waitKey(1) & 0xFF == ord("q"):
break
elif event_id == "bbox":
arrow_bbox = event["value"][0]
bbox_format = event["metadata"]["format"].lower()
if bbox_format == "xyxy":
bbox = arrow_bbox["bbox"].values.to_numpy().reshape(-1, 4)
elif bbox_format == "xywh":
original_bbox = arrow_bbox["bbox"].values.to_numpy().reshape(-1, 4)
bbox = np.array(
[
(
x - w / 2,
y - h / 2,
x + w / 2,
y + h / 2,
)
for [x, y, w, h] in original_bbox
],
)
else:
raise RuntimeError(f"Unsupported bbox format: {bbox_format}")
plot.bboxes = {
"bbox": bbox,
"conf": arrow_bbox["conf"].values.to_numpy(),
"labels": arrow_bbox["labels"].values.to_numpy(
zero_copy_only=False,
),
}
elif event_id == "text":
plot.text = event["value"][0].as_py()
elif event_type == "ERROR":
raise RuntimeError(event["error"])
if __name__ == "__main__":
main()

View File

@@ -0,0 +1,37 @@
[project]
name = "dora-image-viewer"
version = "0.4.1"
license = { file = "MIT" }
authors = [
{ name = "Haixuan Xavier Tao", email = "tao.xavier@outlook.com" },
{ name = "Enzo Le Van", email = "dev@enzo-le-van.fr" },
]
description = "Dora Node for plotting text and bbox on image with OpenCV"
requires-python = ">=3.8"
dependencies = [
"dora-rs >= 0.3.9",
"numpy < 2.0.0",
"opencv-python >= 4.1.1",
"pillow-avif-plugin>=1.5.1",
"pillow>=10.4.0",
]
[dependency-groups]
dev = ["pytest >=8.1.1", "ruff >=0.9.1"]
[project.scripts]
dora-image-viewer = "dora_image_viewer.main:main"
[tool.ruff.lint]
extend-select = [
"D", # pydocstyle
"UP", # Ruff's UP rule
"PERF", # Ruff's PERF rule
"RET", # Ruff's RET rule
"RSE", # Ruff's RSE rule
"NPY", # Ruff's NPY rule
"N", # Ruff's N rule
"I", # Ruff's I rule
]

View File

@@ -0,0 +1,12 @@
"""TODO: Add docstring."""
import pytest
def test_import_main():
"""TODO: Add docstring."""
from dora_image_viewer.main import main
# Check that everything is working, and catch dora Runtime Exception as we're not running in a dora dataflow.
with pytest.raises(RuntimeError):
main()