Rescale detections to input image (luxonis#606)

botsandus · Oct 2, 2024 · 0e6d19c · 0e6d19c
1 parent 97d9854
commit 0e6d19c
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 3 deletions.
diff --git a/depthai_ros_driver/include/depthai_ros_driver/dai_nodes/nn/detection.hpp b/depthai_ros_driver/include/depthai_ros_driver/dai_nodes/nn/detection.hpp
@@ -65,6 +65,9 @@ class Detection : public BaseNode {
         if(ph->getParam<bool>("i_disable_resize")) {
             width = ph->getOtherNodeParam<int>(socketName, "i_preview_width");
             height = ph->getOtherNodeParam<int>(socketName, "i_preview_height");
+        } else if(ph->getParam<bool>("i_desqueeze_output")) {
+            width = ph->getOtherNodeParam<int>(socketName, "i_width");
+            height = ph->getOtherNodeParam<int>(socketName, "i_height");
         } else {
             width = imageManip->initialConfig.getResizeConfig().width;
             height = imageManip->initialConfig.getResizeConfig().height;
@@ -132,6 +135,7 @@ class Detection : public BaseNode {
             ptPub = setupOutput(pipeline, ptQName, [&](dai::Node::Input input) { detectionNode->passthrough.link(input); });
         }
     };
+
     /**
      * @brief      Closes the queues for the DetectionNetwork node and the passthrough.
      */

diff --git a/depthai_ros_driver/include/depthai_ros_driver/param_handlers/nn_param_handler.hpp b/depthai_ros_driver/include/depthai_ros_driver/param_handlers/nn_param_handler.hpp
@@ -40,6 +40,7 @@ class NNParamHandler : public BaseParamHandler {
     template <typename T>
     void declareParams(std::shared_ptr<T> nn, std::shared_ptr<dai::node::ImageManip> imageManip) {
         declareAndLogParam<bool>("i_disable_resize", false);
+        declareAndLogParam<bool>("i_desqueeze_output", false);
         declareAndLogParam<bool>("i_enable_passthrough", false);
         declareAndLogParam<bool>("i_enable_passthrough_depth", false);
         declareAndLogParam<bool>("i_get_base_device_timestamp", false);

diff --git a/depthai_ros_driver/src/dai_nodes/nn/nn_wrapper.cpp b/depthai_ros_driver/src/dai_nodes/nn/nn_wrapper.cpp
@@ -20,15 +20,15 @@ NNWrapper::NNWrapper(const std::string& daiNodeName,
     auto family = ph->getNNFamily();
     switch(family) {
         case param_handlers::nn::NNFamily::Yolo: {
-            nnNode = std::make_unique<dai_nodes::nn::Detection<dai::node::YoloDetectionNetwork>>(getName(), getROSNode(), pipeline);
+            nnNode = std::make_unique<dai_nodes::nn::Detection<dai::node::YoloDetectionNetwork>>(getName(), getROSNode(), pipeline, socket);
             break;
         }
         case param_handlers::nn::NNFamily::Mobilenet: {
-            nnNode = std::make_unique<dai_nodes::nn::Detection<dai::node::MobileNetDetectionNetwork>>(getName(), getROSNode(), pipeline);
+            nnNode = std::make_unique<dai_nodes::nn::Detection<dai::node::MobileNetDetectionNetwork>>(getName(), getROSNode(), pipeline, socket);
             break;
         }
         case param_handlers::nn::NNFamily::Segmentation: {
-            nnNode = std::make_unique<dai_nodes::nn::Segmentation>(getName(), getROSNode(), pipeline);
+            nnNode = std::make_unique<dai_nodes::nn::Segmentation>(getName(), getROSNode(), pipeline, socket);
             break;
         }
     }