Merge pull request #1205 from luxonis/feature/node_io_set_dtype

Feature/node io set dtype
luxonis · Jan 6, 2025 · 1ef7c11 · 1ef7c11
2 parents d8a97c2 + c8d729c
commit 1ef7c11
Show file tree

Hide file tree

Showing 4 changed files with 68 additions and 3 deletions.
diff --git a/bindings/python/src/pipeline/node/NodeBindings.cpp b/bindings/python/src/pipeline/node/NodeBindings.cpp
@@ -313,6 +313,16 @@ void NodeBindings::bind(pybind11::module& m, void* pCallstack) {
              static_cast<Node& (Node::Input::*)()>(&Node::Input::getParent),
              py::return_value_policy::reference_internal,
              DOC(dai, Node, Input, getParent))
+        .def("getPossibleDatatypes", &Node::Input::getPossibleDatatypes, DOC(dai, Node, Input, getPossibleDatatypes))
+        .def("setPossibleDatatypes", &Node::Input::setPossibleDatatypes, py::arg("types"), DOC(dai, Node, Input, setPossibleDatatypes))
+        .def("setPossibleDatatypes", [](Node::Input& input, const std::vector<std::tuple<DatatypeEnum, bool>>& types) {
+            std::vector<Node::DatatypeHierarchy> converted;
+            converted.reserve(types.size());
+            for(const auto& t : types) {
+                converted.emplace_back(std::get<0>(t), std::get<1>(t));
+            }
+            input.setPossibleDatatypes(converted);
+        }, py::arg("types"), DOC(dai, Node, Input, setPossibleDatatypes))
         .def("setWaitForMessage", &Node::Input::setWaitForMessage, py::arg("waitForMessage"), DOC(dai, Node, Input, setWaitForMessage))
         .def("getWaitForMessage", &Node::Input::getWaitForMessage, DOC(dai, Node, Input, getWaitForMessage))
         .def("setReusePreviousMessage", &Node::Input::setReusePreviousMessage, py::arg("reusePreviousMessage"), DOC(dai, Node, Input, setReusePreviousMessage))
@@ -336,6 +346,15 @@ void NodeBindings::bind(pybind11::module& m, void* pCallstack) {
              py::arg("possibleDatatypes") = Node::OutputDescription{}.types,
              py::keep_alive<1, 0>())
         .def("getPossibleDatatypes", &Node::Output::getPossibleDatatypes, DOC(dai, Node, Output, getPossibleDatatypes))
+        .def("setPossibleDatatypes", &Node::Output::setPossibleDatatypes, py::arg("types"), DOC(dai, Node, Output, setPossibleDatatypes))
+        .def("setPossibleDatatypes", [](Node::Output& output, const std::vector<std::tuple<DatatypeEnum, bool>>& types) {
+            std::vector<Node::DatatypeHierarchy> converted;
+            converted.reserve(types.size());
+            for(const auto& t : types) {
+                converted.emplace_back(std::get<0>(t), std::get<1>(t));
+            }
+            output.setPossibleDatatypes(converted);
+        }, py::arg("types"), DOC(dai, Node, Output, setPossibleDatatypes))
         .def("getParent",
              static_cast<const Node& (Node::Output::*)() const>(&Node::Output::getParent),
              py::return_value_policy::reference_internal,

diff --git a/examples/python/HostNodes/threaded_host_nodes.py b/examples/python/HostNodes/threaded_host_nodes.py
@@ -8,6 +8,22 @@ def __init__(self, name: str):
         self.input = self.createInput()
         self.output = self.createOutput()
 
+        # Possible API 1:
+        self.input.setPossibleDatatypes([dai.Node.DatatypeHierarchy(dai.DatatypeEnum.ImgFrame, True)])
+        self.output.setPossibleDatatypes([dai.Node.DatatypeHierarchy(dai.DatatypeEnum.ImgFrame, True)])
+
+        # Possible API 2:
+        self.input.setPossibleDatatypes([
+            (dai.DatatypeEnum.ImgFrame, True),
+            (dai.DatatypeEnum.Buffer, True)
+        ])
+        self.output.setPossibleDatatypes([
+            (dai.DatatypeEnum.ImgFrame, True),
+            (dai.DatatypeEnum.Buffer, True)
+        ])
+
+
+
     def onStart(self):
         print("Hello, this is", self.name)
 
@@ -24,6 +40,7 @@ class TestSink(dai.node.ThreadedHostNode):
     def __init__(self, name: str):
         super().__init__()
         self.input = self.createInput()
+
         self.name = name
 
     def onStart(self):

diff --git a/include/depthai/pipeline/Node.hpp b/include/depthai/pipeline/Node.hpp
@@ -180,9 +180,12 @@ class Node : public std::enable_shared_from_this<Node> {
         /**
          * Get possible datatypes that can be sent
          */
-        std::vector<DatatypeHierarchy> getPossibleDatatypes() const {
-            return desc.types;
-        }
+        std::vector<DatatypeHierarchy> getPossibleDatatypes() const;
+
+        /**
+         * Set possible datatypes that can be sent
+         */
+        void setPossibleDatatypes(std::vector<DatatypeHierarchy> types);
 
         /**
          * Check if this output and given input are on the same pipeline.
@@ -378,6 +381,16 @@ class Node : public std::enable_shared_from_this<Node> {
          */
         bool getWaitForMessage() const;
 
+        /**
+         * Get possible datatypes that can be received
+         */
+        std::vector<DatatypeHierarchy> getPossibleDatatypes() const;
+
+        /**
+         * Set possible datatypes that can be received
+         */
+        void setPossibleDatatypes(std::vector<DatatypeHierarchy> types);
+
         /**
          * Equivalent to setWaitForMessage but with inverted logic.
          */

diff --git a/src/pipeline/Node.cpp b/src/pipeline/Node.cpp
@@ -731,4 +731,20 @@ std::vector<std::pair<Node::Input&, std::shared_ptr<Capability>>> Node::getRequi
     DAI_CHECK_V(false, "Node '{}' doesn't support node to node linking. Please link outputs <--> inputs manually.", getName());
 }
 
+void Node::Output::setPossibleDatatypes(std::vector<Node::DatatypeHierarchy> types) {
+    desc.types = std::move(types);
+}
+
+std::vector<Node::DatatypeHierarchy> Node::Output::getPossibleDatatypes() const {
+    return desc.types;
+}
+
+void Node::Input::setPossibleDatatypes(std::vector<Node::DatatypeHierarchy> types) {
+    possibleDatatypes = std::move(types);
+}
+
+std::vector<Node::DatatypeHierarchy> Node::Input::getPossibleDatatypes() const {
+    return possibleDatatypes;
+}
+
 }  // namespace dai