From b630afa6ba008752612df183532f8047e7b9b3ac Mon Sep 17 00:00:00 2001
From: Myron Rodrigues <rodriguesmyron47@gmail.com>
Date: Fri, 27 Sep 2024 19:03:34 +0000
Subject: [PATCH] ugliest way to set these parameters, cannot log as well coz
 no ros node

---
 .../dai_nodes/nn/nn_wrapper.hpp                  |  1 +
 .../src/dai_nodes/nn/nn_wrapper.cpp              | 16 ++++++++++++++++
 .../src/dai_nodes/sensors/sensor_wrapper.cpp     |  6 ++++++
 3 files changed, 23 insertions(+)

diff --git a/depthai_ros_driver/include/depthai_ros_driver/dai_nodes/nn/nn_wrapper.hpp b/depthai_ros_driver/include/depthai_ros_driver/dai_nodes/nn/nn_wrapper.hpp
index cfab6975..dbdcfa5e 100644
--- a/depthai_ros_driver/include/depthai_ros_driver/dai_nodes/nn/nn_wrapper.hpp
+++ b/depthai_ros_driver/include/depthai_ros_driver/dai_nodes/nn/nn_wrapper.hpp
@@ -33,6 +33,7 @@ class NNWrapper : public BaseNode {
     ~NNWrapper();
     void updateParams(const std::vector<rclcpp::Parameter>& params) override;
     void setupQueues(std::shared_ptr<dai::Device> device) override;
+    void setInputDimensions(int width, int height);
     void link(dai::Node::Input in, int linkType = 0) override;
     dai::Node::Input getInput(int linkType = 0) override;
     virtual void setNames() override;
diff --git a/depthai_ros_driver/src/dai_nodes/nn/nn_wrapper.cpp b/depthai_ros_driver/src/dai_nodes/nn/nn_wrapper.cpp
index b7d56f80..7fdc89ad 100644
--- a/depthai_ros_driver/src/dai_nodes/nn/nn_wrapper.cpp
+++ b/depthai_ros_driver/src/dai_nodes/nn/nn_wrapper.cpp
@@ -48,6 +48,22 @@ void NNWrapper::closeQueues() {
     nnNode->closeQueues();
 }
 
+void NNWrapper::setInputDimensions(int width, int height)
+{
+    auto family = ph->getNNFamily();
+    if (family == param_handlers::nn::NNFamily::Yolo) {
+        auto detectionNode = dynamic_cast<dai_nodes::nn::Detection<dai::node::YoloDetectionNetwork>*>(nnNode.get());
+        if(detectionNode) {
+            detectionNode->setInputDimensions(width, height);
+        }
+    } else if (family == param_handlers::nn::NNFamily::Mobilenet) {
+        auto detectionNode = dynamic_cast<dai_nodes::nn::Detection<dai::node::MobileNetDetectionNetwork>*>(nnNode.get());
+        if(detectionNode) {
+            detectionNode->setInputDimensions(width, height);
+        }
+    }
+}
+
 void NNWrapper::link(dai::Node::Input in, int linkType) {
     nnNode->link(in, linkType);
 }
diff --git a/depthai_ros_driver/src/dai_nodes/sensors/sensor_wrapper.cpp b/depthai_ros_driver/src/dai_nodes/sensors/sensor_wrapper.cpp
index 6693d579..f1276832 100644
--- a/depthai_ros_driver/src/dai_nodes/sensors/sensor_wrapper.cpp
+++ b/depthai_ros_driver/src/dai_nodes/sensors/sensor_wrapper.cpp
@@ -69,6 +69,12 @@ SensorWrapper::SensorWrapper(const std::string& daiNodeName,
     }
     if(ph->getParam<bool>("i_enable_nn")) {
         nnNode = std::make_unique<NNWrapper>(daiNodeName + std::string("_nn"), node, pipeline, static_cast<dai::CameraBoardSocket>(socketID));
+        auto sensorWidth = ph->getParam<int>("i_width");
+        auto sensorHeight = ph->getParam<int>("i_height");
+        auto nn_wrapper = dynamic_cast<NNWrapper*>(nnNode.get());
+        if(nn_wrapper) {
+            nn_wrapper->setInputDimensions(sensorWidth, sensorHeight);
+        } 
         sensorNode->link(nnNode->getInput(), static_cast<int>(link_types::RGBLinkType::preview));
     }
     RCLCPP_DEBUG(node->get_logger(), "Base node %s created", daiNodeName.c_str());