From 0e66a7127fde054c5e30ca295d2ff55d8dabc53a Mon Sep 17 00:00:00 2001 From: haixuanTao Date: Fri, 19 Jul 2024 15:42:47 +0200 Subject: [PATCH] Reformat folder so that they can be used in edit mode as well as use encoding to support multiple encoding --- examples/python-dataflow/dataflow.yml | 8 ++--- node-hub/opencv-plot/README.md | 35 +++++++++---------- .../opencv-plot/{ => opencv_plot}/main.py | 13 +++++-- node-hub/opencv-plot/pyproject.toml | 8 ++--- node-hub/opencv-video-capture/README.md | 34 +++++++++--------- .../{ => opencv_video_capture}/main.py | 13 ++++--- node-hub/opencv-video-capture/pyproject.toml | 8 ++--- node-hub/ultralytics-yolo/README.md | 8 ++--- node-hub/ultralytics-yolo/pyproject.toml | 8 ++--- .../{ => ultralytics_yolo}/main.py | 17 +++++++-- 10 files changed, 82 insertions(+), 70 deletions(-) rename node-hub/opencv-plot/{ => opencv_plot}/main.py (91%) rename node-hub/opencv-video-capture/{ => opencv_video_capture}/main.py (95%) rename node-hub/ultralytics-yolo/{ => ultralytics_yolo}/main.py (82%) diff --git a/examples/python-dataflow/dataflow.yml b/examples/python-dataflow/dataflow.yml index 035242c9..44cd8409 100644 --- a/examples/python-dataflow/dataflow.yml +++ b/examples/python-dataflow/dataflow.yml @@ -1,9 +1,9 @@ nodes: - id: camera - build: pip install ../../node-hub/opencv-video-capture + build: pip install -e ../../node-hub/opencv-video-capture path: opencv-video-capture inputs: - tick: dora/timer/millis/16 + tick: dora/timer/millis/20 outputs: - image env: @@ -12,7 +12,7 @@ nodes: IMAGE_HEIGHT: 480 - id: object-detection - build: pip install ../../node-hub/ultralytics-yolo + build: pip install -e ../../node-hub/ultralytics-yolo path: ultralytics-yolo inputs: image: @@ -24,7 +24,7 @@ nodes: MODEL: yolov8n.pt - id: plot - build: pip install ../../node-hub/opencv-plot + build: pip install -e ../../node-hub/opencv-plot path: opencv-plot inputs: image: diff --git a/node-hub/opencv-plot/README.md b/node-hub/opencv-plot/README.md index 266d5685..ee77c86c 100644 --- a/node-hub/opencv-plot/README.md +++ b/node-hub/opencv-plot/README.md @@ -5,17 +5,17 @@ This node is used to plot a text and a list of bbox on a base image (ideal for o # YAML ```yaml - - id: opencv-plot - build: pip install ../../node-hub/opencv-plot - path: opencv-plot - 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 +- id: opencv-plot + build: pip install ../../node-hub/opencv-plot + path: opencv-plot + 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 @@ -26,19 +26,18 @@ This node is used to plot a text and a list of bbox on a base image (ideal for o image: { "width": np.uint32, "height": np.uint32, - "channels": np.uint8, + "encoding": bytes, "data": np.array # flattened image data } encoded_image = pa.array([image]) decoded_image = { - "width": np.uint32(encoded_image[0]["width"].as_py()), - "height": np.uint32(encoded_image[0]["height"].as_py()), - "channels": np.uint8(encoded_image[0]["channels"].as_py()), + "width": np.uint32(encoded_image[0]["width"]), + "height": np.uint32(encoded_image[0]["height"]), + "encoding": encoded_image[0]["encoding"].as_py(), "data": encoded_image[0]["data"].values.to_numpy().astype(np.uint8) -} - +} ``` - `bbox`: an arrow array containing the bounding boxes, confidence scores, and class names of the detected objects @@ -68,7 +67,7 @@ text: str encoded_text = pa.array([text]) decoded_text = encoded_text[0].as_py() -``` +``` ## License diff --git a/node-hub/opencv-plot/main.py b/node-hub/opencv-plot/opencv_plot/main.py similarity index 91% rename from node-hub/opencv-plot/main.py rename to node-hub/opencv-plot/opencv_plot/main.py index 656b68ff..6b89026d 100644 --- a/node-hub/opencv-plot/main.py +++ b/node-hub/opencv-plot/opencv_plot/main.py @@ -130,11 +130,20 @@ def main(): if event_id == "image": arrow_image = event["value"][0] + + encoding = arrow_image["encoding"].as_py() + if encoding == "bgr8": + channels = 3 + storage_type = np.uint8 + else: + raise Exception(f"Unsupported image encoding: {encoding}") + image = { "width": np.uint32(arrow_image["width"].as_py()), "height": np.uint32(arrow_image["height"].as_py()), - "channels": np.uint8(arrow_image["channels"].as_py()), - "data": arrow_image["data"].values.to_numpy().astype(np.uint8), + "encoding": encoding, + "channels": channels, + "data": arrow_image["data"].values.to_numpy().astype(storage_type), } plot.frame = np.reshape( diff --git a/node-hub/opencv-plot/pyproject.toml b/node-hub/opencv-plot/pyproject.toml index a1f01c45..a6791fe3 100644 --- a/node-hub/opencv-plot/pyproject.toml +++ b/node-hub/opencv-plot/pyproject.toml @@ -3,14 +3,12 @@ name = "opencv-plot" version = "0.1" authors = [ "Haixuan Xavier Tao ", - "Enzo Le Van " + "Enzo Le Van ", ] description = "Dora Node for plotting text and bbox on image with OpenCV" readme = "README.md" -packages = [ - { include = "main.py", to = "opencv_plot" } -] +packages = [{ include = "opencv_plot" }] [tool.poetry.dependencies] dora-rs = "0.3.5" @@ -22,4 +20,4 @@ opencv-plot = "opencv_plot.main:main" [build-system] requires = ["poetry-core>=1.8.0"] -build-backend = "poetry.core.masonry.api" \ No newline at end of file +build-backend = "poetry.core.masonry.api" diff --git a/node-hub/opencv-video-capture/README.md b/node-hub/opencv-video-capture/README.md index 216a47ed..923b9400 100644 --- a/node-hub/opencv-video-capture/README.md +++ b/node-hub/opencv-video-capture/README.md @@ -5,19 +5,19 @@ This node is used to capture video from a camera using OpenCV. # YAML ```yaml - - id: opencv-video-capture - build: pip install ../../node-hub/opencv-video-capture - path: opencv-video-capture - inputs: - tick: dora/timer/millis/16 # try to capture at 60fps - outputs: - - image: # the captured image - - env: - PATH: 0 # optional, default is 0 - - IMAGE_WIDTH: 640 # optional, default is video capture width - IMAGE_HEIGHT: 480 # optional, default is video capture height +- id: opencv-video-capture + build: pip install ../../node-hub/opencv-video-capture + path: opencv-video-capture + inputs: + tick: dora/timer/millis/16 # try to capture at 60fps + outputs: + - image: # the captured image + + env: + PATH: 0 # optional, default is 0 + + IMAGE_WIDTH: 640 # optional, default is video capture width + IMAGE_HEIGHT: 480 # optional, default is video capture height ``` # Inputs @@ -33,16 +33,16 @@ This node is used to capture video from a camera using OpenCV. image: { "width": np.uint32, "height": np.uint32, - "channels": np.uint8, + "encoding": str, "data": np.array # flattened image data } encoded_image = pa.array([image]) decoded_image = { - "width": np.uint32(encoded_image[0]["width"].as_py()), - "height": np.uint32(encoded_image[0]["height"].as_py()), - "channels": np.uint8(encoded_image[0]["channels"].as_py()), + "width": np.uint32(encoded_image[0]["width"]), + "height": np.uint32(encoded_image[0]["height"]), + "encoding": encoded_image[0]["encoding"].as_py(), "data": encoded_image[0]["data"].values.to_numpy().astype(np.uint8) } ``` diff --git a/node-hub/opencv-video-capture/main.py b/node-hub/opencv-video-capture/opencv_video_capture/main.py similarity index 95% rename from node-hub/opencv-video-capture/main.py rename to node-hub/opencv-video-capture/opencv_video_capture/main.py index 892cc85b..54e3886a 100644 --- a/node-hub/opencv-video-capture/main.py +++ b/node-hub/opencv-video-capture/opencv_video_capture/main.py @@ -54,18 +54,21 @@ def main(): if isinstance(video_capture_path, str) and video_capture_path.isnumeric(): video_capture_path = int(video_capture_path) + video_capture = cv2.VideoCapture(video_capture_path) + image_width = os.getenv("IMAGE_WIDTH", args.image_width) - image_height = os.getenv("IMAGE_HEIGHT", args.image_height) if image_width is not None: if isinstance(image_width, str) and image_width.isnumeric(): image_width = int(image_width) + video_capture.set(cv2.CAP_PROP_FRAME_WIDTH, image_width) + image_height = os.getenv("IMAGE_HEIGHT", args.image_height) if image_height is not None: if isinstance(image_height, str) and image_height.isnumeric(): image_height = int(image_height) + video_capture.set(cv2.CAP_PROP_FRAME_HEIGHT, image_height) - video_capture = cv2.VideoCapture(video_capture_path) node = Node(args.name) start_time = time.time() @@ -105,16 +108,12 @@ def main(): image = { "width": np.uint32(frame.shape[1]), "height": np.uint32(frame.shape[0]), - "channels": np.uint8(frame.shape[2]), + "encoding": "bgr8", "data": frame.ravel(), } node.send_output("image", pa.array([image]), event["metadata"]) - if event_id == "stop": - video_capture.release() - break - elif event_type == "ERROR": raise Exception(event["error"]) diff --git a/node-hub/opencv-video-capture/pyproject.toml b/node-hub/opencv-video-capture/pyproject.toml index 00100acf..5d686ebd 100644 --- a/node-hub/opencv-video-capture/pyproject.toml +++ b/node-hub/opencv-video-capture/pyproject.toml @@ -3,14 +3,12 @@ name = "opencv-video-capture" version = "0.1" authors = [ "Haixuan Xavier Tao ", - "Enzo Le Van " + "Enzo Le Van ", ] description = "Dora Node for capturing video with OpenCV" readme = "README.md" -packages = [ - { include = "main.py", to = "opencv_video_capture" } -] +packages = [{ include = "opencv_video_capture" }] [tool.poetry.dependencies] dora-rs = "0.3.5" @@ -22,4 +20,4 @@ opencv-video-capture = "opencv_video_capture.main:main" [build-system] requires = ["poetry-core>=1.8.0"] -build-backend = "poetry.core.masonry.api" \ No newline at end of file +build-backend = "poetry.core.masonry.api" diff --git a/node-hub/ultralytics-yolo/README.md b/node-hub/ultralytics-yolo/README.md index 8a7bc8f2..fae47540 100644 --- a/node-hub/ultralytics-yolo/README.md +++ b/node-hub/ultralytics-yolo/README.md @@ -25,16 +25,16 @@ This node is used to detect objects in images using YOLOv8. image: { "width": np.uint32, "height": np.uint32, - "channels": np.uint8, + "encoding": str, "data": np.array # flattened image data } encoded_image = pa.array([image]) decoded_image = { - "width": np.uint32(encoded_image[0]["width"].as_py()), - "height": np.uint32(encoded_image[0]["height"].as_py()), - "channels": np.uint8(encoded_image[0]["channels"].as_py()), + "width": np.uint32(encoded_image[0]["width"]), + "height": np.uint32(encoded_image[0]["height"]), + "encoding": encoded_image[0]["encoding"].as_py(), "data": encoded_image[0]["data"].values.to_numpy().astype(np.uint8) } diff --git a/node-hub/ultralytics-yolo/pyproject.toml b/node-hub/ultralytics-yolo/pyproject.toml index 0d7feacf..2665168a 100644 --- a/node-hub/ultralytics-yolo/pyproject.toml +++ b/node-hub/ultralytics-yolo/pyproject.toml @@ -3,14 +3,12 @@ name = "ultralytics-yolo" version = "0.1" authors = [ "Haixuan Xavier Tao ", - "Enzo Le Van " + "Enzo Le Van ", ] description = "Dora Node for object detection with Ultralytics YOLOv8" readme = "README.md" -packages = [ - { include = "main.py", to = "ultralytics_yolo" } -] +packages = [{ include = "ultralytics_yolo" }] [tool.poetry.dependencies] dora-rs = "0.3.5" @@ -22,4 +20,4 @@ ultralytics-yolo = "ultralytics_yolo.main:main" [build-system] requires = ["poetry-core>=1.8.0"] -build-backend = "poetry.core.masonry.api" \ No newline at end of file +build-backend = "poetry.core.masonry.api" diff --git a/node-hub/ultralytics-yolo/main.py b/node-hub/ultralytics-yolo/ultralytics_yolo/main.py similarity index 82% rename from node-hub/ultralytics-yolo/main.py rename to node-hub/ultralytics-yolo/ultralytics_yolo/main.py index abb3e048..023742c8 100644 --- a/node-hub/ultralytics-yolo/main.py +++ b/node-hub/ultralytics-yolo/ultralytics_yolo/main.py @@ -46,18 +46,29 @@ def main(): if event_id == "image": arrow_image = event["value"][0] + encoding = arrow_image["encoding"].as_py() + + if encoding == "bgr8": + channels = 3 + storage_type = np.uint8 + else: + raise Exception(f"Unsupported image encoding: {encoding}") + image = { "width": np.uint32(arrow_image["width"].as_py()), "height": np.uint32(arrow_image["height"].as_py()), - "channels": np.uint8(arrow_image["channels"].as_py()), - "data": arrow_image["data"].values.to_numpy().astype(np.uint8), + "encoding": encoding, + "channels": channels, + "data": arrow_image["data"].values.to_numpy().astype(storage_type), } frame = image["data"].reshape( (image["height"], image["width"], image["channels"]) ) - frame = frame[:, :, ::-1] # OpenCV image (BGR to RGB) + if encoding == "bgr8": + frame = frame[:, :, ::-1] # OpenCV image (BGR to RGB) + results = model(frame, verbose=False) # includes NMS bboxes = np.array(results[0].boxes.xyxy.cpu())