[#54504] src: cvnode_base: Adapt C++ nodes to updated messages API

Adapts the CVNodeBase class implemented in C++ to address modifications in 'kenning_computer_vision_msgs' ROS2 package. Inference now is performed only in a 'online' mode, where the system should only react to a single frame at a time. Signed-off-by: Illia Vysochyn <[email protected]>
antmicro · Feb 7, 2024 · ee306a8 · ee306a8
1 parent 711e905
commit ee306a8
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 18 deletions.
diff --git a/include/cvnode_base/cvnode_base.hpp b/include/cvnode_base/cvnode_base.hpp
@@ -10,6 +10,7 @@
 #include <string>
 
 #include <kenning_computer_vision_msgs/msg/segmentation_msg.hpp>
+#include <kenning_computer_vision_msgs/msg/video_frame_msg.hpp>
 #include <kenning_computer_vision_msgs/srv/manage_cv_node.hpp>
 #include <kenning_computer_vision_msgs/srv/segment_cv_node_srv.hpp>
 
@@ -107,8 +108,8 @@ class BaseCVNode : public rclcpp::Node
      *
      * @return Inference output.
      */
-    virtual std::vector<kenning_computer_vision_msgs::msg::SegmentationMsg>
-    run_inference(std::vector<sensor_msgs::msg::Image> &X) = 0;
+    virtual kenning_computer_vision_msgs::msg::SegmentationMsg
+    run_inference(kenning_computer_vision_msgs::msg::VideoFrameMsg &X) = 0;
 
     /**
      * Cleanup allocated model resources.

diff --git a/include/cvnode_base/nodes/mask_rcnn_torchscript.hpp b/include/cvnode_base/nodes/mask_rcnn_torchscript.hpp
@@ -70,12 +70,12 @@ class MaskRCNNTorchScript : public BaseCVNode
     /**
      * Run inference on input images.
      *
-     * @param X Vector of input images.
+     * @param X Input image.
      *
-     * @return Vector of instance segmentation results.
+     * @return Instance segmentation predictions.
      */
-    std::vector<kenning_computer_vision_msgs::msg::SegmentationMsg>
-    run_inference(std::vector<sensor_msgs::msg::Image> &X) override;
+    kenning_computer_vision_msgs::msg::SegmentationMsg
+    run_inference(kenning_computer_vision_msgs::msg::VideoFrameMsg &X) override;
 
     /**
      * Preprocess input image for inference.

diff --git a/src/cvnode_base.cpp b/src/cvnode_base.cpp
@@ -38,7 +38,7 @@ void BaseCVNode::process_callback(
     SegmentCVNodeSrv::Response::SharedPtr response)
 {
     RCLCPP_DEBUG(get_logger(), "Received request to process input data");
-    response->output = run_inference(request->input);
+    response->segmentation = run_inference(request->input);
     response->success = true;
     RCLCPP_DEBUG(get_logger(), "Processed the input data");
 }

diff --git a/src/nodes/mask_rcnn_torchscript.cpp b/src/nodes/mask_rcnn_torchscript.cpp
@@ -19,6 +19,7 @@ namespace cvnode_base
 {
 
 using SegmentationMsg = kenning_computer_vision_msgs::msg::SegmentationMsg;
+using VideoFrameMsg = kenning_computer_vision_msgs::msg::VideoFrameMsg;
 
 MaskRCNNTorchScript::MaskRCNNTorchScript(const rclcpp::NodeOptions &options)
     : BaseCVNode("mask_rcnn_torchscript_node", options)
@@ -84,18 +85,14 @@ bool MaskRCNNTorchScript::prepare()
     return true;
 }
 
-std::vector<SegmentationMsg> MaskRCNNTorchScript::run_inference(std::vector<sensor_msgs::msg::Image> &X)
+SegmentationMsg MaskRCNNTorchScript::run_inference(VideoFrameMsg &X)
 {
-    std::vector<SegmentationMsg> result;
-    for (auto &frame : X)
+    c10::IValue input = preprocess(X.frame);
+    MaskRCNNOutputs prediction = predict(input);
+    SegmentationMsg result = postprocess(prediction, X.frame);
+    if (device.is_cuda())
     {
-        c10::IValue input = preprocess(frame);
-        MaskRCNNOutputs prediction = predict(input);
-        result.push_back(postprocess(prediction, frame));
-        if (device.is_cuda())
-        {
-            c10::cuda::CUDACachingAllocator::emptyCache();
-        }
+        c10::cuda::CUDACachingAllocator::emptyCache();
     }
     return result;
 }
@@ -129,7 +126,6 @@ SegmentationMsg MaskRCNNTorchScript::postprocess(MaskRCNNOutputs &prediction, se
     using MaskMsg = kenning_computer_vision_msgs::msg::MaskMsg;
     using BoxMsg = kenning_computer_vision_msgs::msg::BoxMsg;
     SegmentationMsg msg;
-    msg.frame = frame;
     std::transform(
         prediction.classes.data_ptr<int64_t>(),
         prediction.classes.data_ptr<int64_t>() + prediction.classes.numel(),