First commit
This commit is contained in:
95
dora_image_viewer/README.md
Normal file
95
dora_image_viewer/README.md
Normal 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.
|
||||
13
dora_image_viewer/dora_image_viewer/__init__.py
Normal file
13
dora_image_viewer/dora_image_viewer/__init__.py
Normal 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."
|
||||
250
dora_image_viewer/dora_image_viewer/main.py
Normal file
250
dora_image_viewer/dora_image_viewer/main.py
Normal 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()
|
||||
37
dora_image_viewer/pyproject.toml
Normal file
37
dora_image_viewer/pyproject.toml
Normal 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
|
||||
]
|
||||
12
dora_image_viewer/tests/test_dora_image_viewer.py
Normal file
12
dora_image_viewer/tests/test_dora_image_viewer.py
Normal 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()
|
||||
Reference in New Issue
Block a user