From 4c09bc58768bfd1b86a0e67d9577bd2f10e9c805 Mon Sep 17 00:00:00 2001 From: vangeliq Date: Fri, 18 Oct 2024 20:56:04 -0700 Subject: [PATCH] Added inference node. --- .../ros2_ws/src/python_workspace/README.md | 37 +++++++-- .../python_workspace/inference_node.py | 65 ++++++++++++++++ .../python_workspace/scripts/__init__.py | 0 .../python_workspace/scripts/utils.py | 76 +++++++++++++++++++ .../ros2_ws/src/python_workspace/setup.py | 7 +- 5 files changed, 176 insertions(+), 9 deletions(-) create mode 100644 workspace_python/ros2_ws/src/python_workspace/python_workspace/inference_node.py create mode 100644 workspace_python/ros2_ws/src/python_workspace/python_workspace/scripts/__init__.py create mode 100644 workspace_python/ros2_ws/src/python_workspace/python_workspace/scripts/utils.py diff --git a/workspace_python/ros2_ws/src/python_workspace/README.md b/workspace_python/ros2_ws/src/python_workspace/README.md index 2cdcb37..c4e8825 100644 --- a/workspace_python/ros2_ws/src/python_workspace/README.md +++ b/workspace_python/ros2_ws/src/python_workspace/README.md @@ -1,11 +1,34 @@ -Running the camera node: -ros2 run python_workspace camera_node --ros-args -p static_image_path:='/home/user/Desktop/ROS/Models/Maize Model/sample_maize_images' -p loop:=-1 -p frame_rate:=10 -p model_dimensions:=[640,640] -Running the picture node: -`ros2 run python_workspace picture_node --ros-args -p static_image_path:='./../assets/maize' -p loop:=-1 -p fram -e_rate:=1` +### List of Nodes +| Node Type | Subscribes To | Publishes To | Example command | +|-------------------|----------------|-----------------------------------|------| +| Picture Node | `/picture/command` | `/input_image` | `ros2 run python_workspace picture_node --ros-args -p static_image_path:='./../assets/maize' -p loop:=-1 -p frame_rate:=1`| +| Inference Node | `/input_image` | -`/bounding_box_coordinates`
- `/output_img` | `ros2 run python_workspace inference_node --ros-args -p weights_path:='./ros2_ws/src/python_workspace/python_workspace/scripts/yolo11n.pt'`| -ros2 run python_workspace jetson_node -ros2 run python_workspace extermination_node \ No newline at end of file +### List of Topics and Data Types + +| Topic Name | Description | Data Type | +|-----------------------------|--------------------------------------|--------------------| +| `/input_image` | Input image for processing | `sensor_msgs/Image`| +| `/bounding_box_coordinates` | Coordinates of detected objects (xyxy) | `std_msgs/Float32MultiArray` | +| `/output_img` | Processed output image | `sensor_msgs/Image`| + + + +### Other commands +#### Running the camera node: +`ros2 run python_workspace camera_node --ros-args -p static_image_path:='/home/user/Desktop/ROS/Models/Maize Model/sample_maize_images' -p loop:=-1 -p frame_rate:=10 -p model_dimensions:=[640,640]` + +#### Running the picture node: + +`ros2 run python_workspace picture_node --ros-args -p static_image_path:='./../assets/maize' -p loop:=-1 -p frame_rate:=1` + +#### Running the inference node: + +The path for weights (.pt) is relative to the ROS/workspace_python/ros2_ws directory. + +```bash +ros2 run python_workspace inference_node --ros-args -p weights_path:='./ros2_ws/src/python_workspace/python_workspace/scripts/yolo11n.pt' +``` diff --git a/workspace_python/ros2_ws/src/python_workspace/python_workspace/inference_node.py b/workspace_python/ros2_ws/src/python_workspace/python_workspace/inference_node.py new file mode 100644 index 0000000..eadc94d --- /dev/null +++ b/workspace_python/ros2_ws/src/python_workspace/python_workspace/inference_node.py @@ -0,0 +1,65 @@ +import os +from cv_bridge import CvBridge + +import rclpy +from rclpy.node import Node +from sensor_msgs.msg import Image +from std_msgs.msg import Int32MultiArray, Float32MultiArray + +from .scripts.utils import ModelInference + +class InferenceNode(Node): + def __init__(self): + super().__init__('inference_node') + + self.declare_parameter('weights_path', './src/python_workspace/python_workspace/scripts/best.onnx') + self.declare_parameter('precision', 'fp32') # fp32, fp16 # todo: do something with strip_weights and precision + + self.weights_path = self.get_parameter('weights_path').get_parameter_value().string_value + if not os.path.isabs(self.weights_path): + self.weights_path = os.path.join(os.getcwd(), self.weights_path) + + + self.precision = self.get_parameter('precision').get_parameter_value().string_value + + # instantiate the model here + self.model = ModelInference(self.weights_path, self.precision) + + self.bridge = CvBridge() + + self.subscription = self.create_subscription(Image, 'input_image', self.image_callback, 10) + + # create a publisher for the output image/boxes/extermination data + self.box_publisher = self.create_publisher(Float32MultiArray,'bounding_box_coordinates', 10) + + self.output_image_publisher = self.create_publisher(Image, 'output_img', 10) + + + def image_callback(self, msg): + # print("============================================") + opencv_img = self.bridge.imgmsg_to_cv2(msg, desired_encoding='passthrough') + output_img, named_classes, confidences, boxes = self.model.inference(opencv_img) + + output_msg = Float32MultiArray() + output_msg.data = [] + if len(boxes) != 0: + output_msg.data = boxes + self.box_publisher.publish(output_msg) + + + # convert the output image to a ROS2 image message + output_msg = self.bridge.cv2_to_imgmsg(output_img, encoding='rgb8') + + self.output_image_publisher.publish(output_msg) + + + +def main(args=None): + rclpy.init(args=args) + inference_node = InferenceNode() + rclpy.spin(inference_node) + inference_node.destroy_node() + rclpy.shutdown() + +if __name__ == '__main__': + main() \ No newline at end of file diff --git a/workspace_python/ros2_ws/src/python_workspace/python_workspace/scripts/__init__.py b/workspace_python/ros2_ws/src/python_workspace/python_workspace/scripts/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/workspace_python/ros2_ws/src/python_workspace/python_workspace/scripts/utils.py b/workspace_python/ros2_ws/src/python_workspace/python_workspace/scripts/utils.py new file mode 100644 index 0000000..02a8350 --- /dev/null +++ b/workspace_python/ros2_ws/src/python_workspace/python_workspace/scripts/utils.py @@ -0,0 +1,76 @@ +import cv2 +import time, os +import tensorrt as trt +import pycuda.driver as cuda +import cupy as cp +import numpy as np +from sensor_msgs.msg import Image +from ultralytics import YOLO +from ultralytics.engine.results import Results + + +#utility class for the python workspace + + +class ModelInference: + """ + class that performs inference using a model + """ + def __init__(self, weights_path, precision): + + # Initialize instance variables + self.precision = precision + + if not os.path.exists(weights_path): + print(f"weights file not found at {weights_path}") + raise FileNotFoundError(f"weights file not found at {weights_path}") + + self.yolo = YOLO(weights_path) + + def preprocess(self, image: np.ndarray): + """ + Takes in a numpy array that has been preprocessed + No preprocessing is nededed for YOLO + """ + return image + + + def inference(self, image_array: np.ndarray) -> Results: + # https://docs.ultralytics.com/modes/predict/#inference-arguments + """ + perform inference on the image array. + Returns list of predicted classes, their confidences and bounding boxes + """ + height, width, _ = image_array.shape + size = max(height, width) + results =self.yolo.predict(image_array,imgsz=size) #we are only predicting one image + + + if results[0] is not None: + out_img, named_classes, confidences, boxes = self.postprocess(results[0]) + + return out_img,named_classes, confidences, boxes + + return None, [], [], [] + + def postprocess(self, result:Results) -> np.ndarray: + """ + postprocess the Result that we got from prediction + returns the image with bounding boxes + class names, confidences and bounding boxes + + Since YOLO does this for us we just need to convert the BGR image to RGB. + """ + out_img = cv2.cvtColor(result.plot(), cv2.COLOR_BGR2RGB) # Convert BGR to RGB + + class_mapping = self.yolo.model.names + named_classes = [class_mapping[int(i)] for i in result.boxes.cls] + confidences = result.boxes.conf + print(f"Bounding boxes (xyxy format): {result.boxes.xyxy}") + boxes = result.boxes.xyxy.flatten().tolist() + + + return out_img,named_classes, confidences, boxes + + def print_info(self): + print(self.yolo.info()) \ No newline at end of file diff --git a/workspace_python/ros2_ws/src/python_workspace/setup.py b/workspace_python/ros2_ws/src/python_workspace/setup.py index 8ba6f10..48fd24c 100644 --- a/workspace_python/ros2_ws/src/python_workspace/setup.py +++ b/workspace_python/ros2_ws/src/python_workspace/setup.py @@ -1,13 +1,14 @@ import os from glob import glob -from setuptools import setup +from setuptools import setup, find_packages package_name = 'python_workspace' +subfolder = f'{package_name}/scripts' setup( name=package_name, version='0.0.0', - packages=[package_name], + packages=[package_name,subfolder], data_files=[ ('share/ament_index/resource_index/packages', ['resource/' + package_name]), @@ -24,8 +25,10 @@ 'video_node = python_workspace.video_node:main', 'camera_node = python_workspace.zed_camera_node:main', 'picture_node = python_workspace.picture_node:main', + 'inference_node = python_workspace.inference_node:main', 'jetson_node = python_workspace.jetson_node:main', 'extermination_node = python_workspace.extermination_node:main' ], }, + ) \ No newline at end of file