From e901c30b8d812b9a619f2777bae7ddbea14b3f7a Mon Sep 17 00:00:00 2001 From: Petros Toupas Date: Thu, 5 Dec 2024 16:14:12 +0200 Subject: [PATCH 1/3] add getParser method to ParsingNeuralNetwork --- depthai_nodes/parsing_neural_network.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/depthai_nodes/parsing_neural_network.py b/depthai_nodes/parsing_neural_network.py index db757b74..3e97958a 100644 --- a/depthai_nodes/parsing_neural_network.py +++ b/depthai_nodes/parsing_neural_network.py @@ -127,6 +127,13 @@ def setBackend(self, setBackend: str) -> None: """Sets the backend of the NeuralNetwork node.""" self._nn.setBackend(setBackend) + def getParser(self, parserID: int = 0) -> BaseParser: + """Returns the parser node for the specified model head.""" + assert ( + parserID in self._parsers + ), f"Parser with ID {parserID} not found. Available parser IDs: {list(self._parsers.keys())}" + return self._parsers[parserID] + def setBackendProperties(self, setBackendProperties: Dict[str, str]) -> None: """Sets the backend properties of the NeuralNetwork node.""" self._nn.setBackendProperties(setBackendProperties) From 458fa1382e0ad093a11ddb007586e0d1592fcef4 Mon Sep 17 00:00:00 2001 From: Petros Toupas Date: Fri, 6 Dec 2024 09:09:02 +0200 Subject: [PATCH 2/3] change assertion to if with KeyError --- depthai_nodes/parsing_neural_network.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/depthai_nodes/parsing_neural_network.py b/depthai_nodes/parsing_neural_network.py index 3e97958a..9794c620 100644 --- a/depthai_nodes/parsing_neural_network.py +++ b/depthai_nodes/parsing_neural_network.py @@ -129,9 +129,10 @@ def setBackend(self, setBackend: str) -> None: def getParser(self, parserID: int = 0) -> BaseParser: """Returns the parser node for the specified model head.""" - assert ( - parserID in self._parsers - ), f"Parser with ID {parserID} not found. Available parser IDs: {list(self._parsers.keys())}" + if parserID not in self._parsers: + raise KeyError( + f"Parser with ID {parserID} not found. Available parser IDs: {list(self._parsers.keys())}" + ) return self._parsers[parserID] def setBackendProperties(self, setBackendProperties: Dict[str, str]) -> None: From 71f054d89d83e96412cf3d41ddb7c87903029e20 Mon Sep 17 00:00:00 2001 From: Petros Toupas Date: Fri, 6 Dec 2024 13:41:29 +0200 Subject: [PATCH 3/3] update docstring and change the input variable name --- depthai_nodes/parsing_neural_network.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/depthai_nodes/parsing_neural_network.py b/depthai_nodes/parsing_neural_network.py index 9794c620..7126ed30 100644 --- a/depthai_nodes/parsing_neural_network.py +++ b/depthai_nodes/parsing_neural_network.py @@ -127,13 +127,16 @@ def setBackend(self, setBackend: str) -> None: """Sets the backend of the NeuralNetwork node.""" self._nn.setBackend(setBackend) - def getParser(self, parserID: int = 0) -> BaseParser: - """Returns the parser node for the specified model head.""" - if parserID not in self._parsers: + def getParser(self, index: int = 0) -> BaseParser: + """Returns the parser node for the given model head index. + + If index is not provided, the first parser node is returned by default. + """ + if index not in self._parsers: raise KeyError( - f"Parser with ID {parserID} not found. Available parser IDs: {list(self._parsers.keys())}" + f"Parser with ID {index} not found. Available parser IDs: {list(self._parsers.keys())}" ) - return self._parsers[parserID] + return self._parsers[index] def setBackendProperties(self, setBackendProperties: Dict[str, str]) -> None: """Sets the backend properties of the NeuralNetwork node."""