Extract-compute-emit refactor for run() method of parsers

.github/workflows/e2e_tests.yaml

@@ -115,7 +115,6 @@ jobs:
             fi
           else
             exec hil_runner --platforms $PLATFORM --wait --reservation-name $RESERVATION_NAME --before-docker-pull "echo ${{ github.token }} | docker login ghcr.io -u ${{ github.actor }} --password-stdin" --docker-image ghcr.io/luxonis/depthai-nodes-testing --docker-run-args "--env LUXONIS_EXTRA_INDEX_URL=${{secrets.LUXONIS_EXTRA_INDEX_URL}} --env DEPTHAI_VERSION=\"3.8.0\" --env HUBAI_TEAM_SLUG=${{ secrets.HUBAI_TEAM_SLUG }} --env HUBAI_API_KEY=${{ secrets.HUBAI_API_KEY }} --env BRANCH=\"$SAFE_REF_NAME\" --env FLAGS=\"-all --platform $SAFE_MATRIX_PLATFORM\""
-
           fi
 
   HIL-test-windows:

depthai_nodes/node/parsers/classification.py

@@ -7,7 +7,7 @@
     create_classification_message,
 )
 from depthai_nodes.node.parsers.base_parser import BaseParser
-from depthai_nodes.node.parsers.utils import softmax
+from depthai_nodes.node.parsers.utils import compute_classification_scores
 
 
 class ClassificationParser(BaseParser):
@@ -131,36 +131,46 @@ def run(self):
             except dai.MessageQueue.QueueException:
                 break
 
-            layers = output.getAllLayerNames()
-            self._logger.debug(f"Processing input with layers: {layers}")
-            if len(layers) == 1 and self.output_layer_name == "":
-                self.output_layer_name = layers[0]
-            elif len(layers) != 1 and self.output_layer_name == "":
-                raise ValueError(
-                    f"Expected 1 output layer, got {len(layers)} layers. Please provide the output_layer_name."
-                )
+            scores = self.extract(output)
+            scores = self.compute(scores, is_softmax=self.is_softmax)
+            self.emit(output, scores)
 
-            scores = output.getTensor(self.output_layer_name, dequantize=True)
-            scores = np.array(scores).flatten()
-
-            if len(scores) != self.n_classes and self.n_classes != 0:
-                raise ValueError(
-                    f"Number of labels and scores mismatch. Provided {self.n_classes} class names and {len(scores)} scores."
-                )
-
-            if not self.is_softmax:
-                scores = softmax(scores)
+    def extract(self, output: dai.NNData) -> np.ndarray:
+        layers = output.getAllLayerNames()
+        self._logger.debug(f"Processing input with layers: {layers}")
+        if len(layers) == 1 and self.output_layer_name == "":
+            self.output_layer_name = layers[0]
+        elif len(layers) != 1 and self.output_layer_name == "":
+            raise ValueError(
+                f"Expected 1 output layer, got {len(layers)} layers. Please provide the output_layer_name."
+            )
 
-            msg = create_classification_message(self.classes, scores)
-            transformation = output.getTransformation()
-            if transformation is not None:
-                msg.setTransformation(transformation)
-            msg.setTimestamp(output.getTimestamp())
-            msg.setSequenceNum(output.getSequenceNum())
-            msg.setTimestampDevice(output.getTimestampDevice())
+        scores = output.getTensor(self.output_layer_name, dequantize=True).flatten()
 
-            self._logger.debug(f"Created message with {len(msg.classes)} classes")
+        if len(scores) != self.n_classes and self.n_classes != 0:
+            raise ValueError(
+                f"Number of labels and scores mismatch. Provided {self.n_classes} class names and {len(scores)} scores."
+            )
 
-            self.out.send(msg)
+        return scores
 
-            self._logger.debug("Classification message sent successfully")
+    @staticmethod
+    def compute(
+        scores: np.ndarray,
+        *,
+        is_softmax: bool = True,
+    ) -> np.ndarray:
+        return compute_classification_scores(scores, is_softmax=is_softmax)
+
+    def emit(self, output: dai.NNData, scores: np.ndarray) -> None:
+        msg = create_classification_message(self.classes, scores)
+        transformation = output.getTransformation()
+        if transformation is not None:
+            msg.setTransformation(transformation)
+        msg.setTimestamp(output.getTimestamp())
+        msg.setSequenceNum(output.getSequenceNum())
+        msg.setTimestampDevice(output.getTimestampDevice())
+
+        self._logger.debug(f"Created message with {len(msg.classes)} classes")
+        self.out.send(msg)
+        self._logger.debug("Classification message sent successfully")

depthai_nodes/node/parsers/classification_sequence.py

@@ -4,7 +4,9 @@
 import numpy as np
 
 from depthai_nodes.message.creators import create_classification_sequence_message
-from depthai_nodes.node.parsers.utils import softmax
+from depthai_nodes.node.parsers.utils.classification_sequence import (
+    compute_classification_sequence_scores,
+)
 
 from .classification import ClassificationParser
 
@@ -149,56 +151,46 @@ def run(self):
             except dai.MessageQueue.QueueException:
                 break
 
-            layers = output.getAllLayerNames()
-            self._logger.debug(f"Processing input with layers: {layers}")
-            if len(layers) == 1 and self.output_layer_name == "":
-                self.output_layer_name = layers[0]
-            elif len(layers) != 1 and self.output_layer_name == "":
-                raise ValueError(
-                    f"Expected 1 output layer, got {len(layers)} layers. Please provide the output_layer_name."
-                )
-
-            if self.n_classes == 0:
-                raise ValueError("Classes must be provided for classification.")
-
-            scores = output.getTensor(self.output_layer_name, dequantize=True).astype(
-                np.float32
-            )
-
-            if len(scores.shape) != 3 and len(scores.shape) != 2:
-                raise ValueError(
-                    f"Scores should be a 3D or 2D array, got shape {scores.shape}."
-                )
-
-            if len(scores.shape) == 3:
-                if scores.shape[0] == 1:
-                    scores = scores[0]
-                elif scores.shape[2] == 1:
-                    scores = scores[:, :, 0]
-                else:
-                    raise ValueError(
-                        "Scores should be a 3D array of shape (1, sequence_length, n_classes) or (sequence_length, n_classes, 1)."
-                    )
-
-            if not self.is_softmax:
-                scores = softmax(scores, axis=1, keep_dims=True)
-
-            msg = create_classification_sequence_message(
-                classes=self.classes,
-                scores=scores,
-                remove_duplicates=self.remove_duplicates,
-                ignored_indexes=self.ignored_indexes,
-                concatenate_classes=self.concatenate_classes,
+            scores = self.extract(output)
+            scores = self.compute(scores, is_softmax=self.is_softmax)
+            self.emit(output, scores)
+
+    def extract(self, output: dai.NNData) -> np.ndarray:
+        layers = output.getAllLayerNames()
+        self._logger.debug(f"Processing input with layers: {layers}")
+        if len(layers) == 1 and self.output_layer_name == "":
+            self.output_layer_name = layers[0]
+        elif len(layers) != 1 and self.output_layer_name == "":
+            raise ValueError(
+                f"Expected 1 output layer, got {len(layers)} layers. Please provide the output_layer_name."
             )
-            msg.setTimestamp(output.getTimestamp())
-            msg.setSequenceNum(output.getSequenceNum())
-            msg.setTimestampDevice(output.getTimestampDevice())
-            transformation = output.getTransformation()
-            if transformation is not None:
-                msg.setTransformation(transformation)
 
-            self._logger.debug(f"Created message with {len(msg.classes)} classes")
+        if self.n_classes == 0:
+            raise ValueError("Classes must be provided for classification.")
 
-            self.out.send(msg)
+        return output.getTensor(self.output_layer_name, dequantize=True).astype(
+            np.float32
+        )
 
-            self._logger.debug("Classification message sent successfully")
+    @staticmethod
+    def compute(scores: np.ndarray, *, is_softmax: bool = True) -> np.ndarray:
+        return compute_classification_sequence_scores(scores, is_softmax=is_softmax)
+
+    def emit(self, output: dai.NNData, scores: np.ndarray) -> None:
+        msg = create_classification_sequence_message(
+            classes=self.classes,
+            scores=scores,
+            remove_duplicates=self.remove_duplicates,
+            ignored_indexes=self.ignored_indexes,
+            concatenate_classes=self.concatenate_classes,
+        )
+        msg.setTimestamp(output.getTimestamp())
+        msg.setSequenceNum(output.getSequenceNum())
+        msg.setTimestampDevice(output.getTimestampDevice())
+        transformation = output.getTransformation()
+        if transformation is not None:
+            msg.setTransformation(transformation)
+
+        self._logger.debug(f"Created message with {len(msg.classes)} classes")
+        self.out.send(msg)
+        self._logger.debug("Classification message sent successfully")

depthai_nodes/node/parsers/detection.py

@@ -4,8 +4,7 @@
 from depthai_nodes.message.creators import (
     create_detection_message,
 )
-from depthai_nodes.node.parsers.utils.bbox_format_converters import xyxy_to_xywh
-from depthai_nodes.node.parsers.utils.nms import nms_cv2
+from depthai_nodes.node.parsers.utils.detection import compute_detection_outputs
 
 from .base_parser import BaseParser
 
@@ -137,54 +136,69 @@ def run(self):
             except dai.MessageQueue.QueueException:
                 break  # Pipeline was stopped
 
-            layers = output.getAllLayerNames()
-            if len(layers) != 2:
-                raise ValueError(
-                    f"Expected 2 output layers, got {len(layers)} layers. Please use different parser or create a new one."
-                )
-
-            bboxes = None
-            scores = None
-
-            for layer in layers:
-                tensor: np.ndarray = output.getTensor(layer, dequantize=True)
-                if tensor.shape[-1] == 4 and len(tensor.shape) != 1:
-                    bboxes = tensor
-                else:
-                    scores = tensor
-
-            bboxes = bboxes.reshape(-1, 4)
-            scores = scores.reshape(-1)
-
-            if bboxes is None or scores is None:
-                raise ValueError(
-                    "Bounding boxes or scores are missing in the output. Please check the NN model."
-                )
-
-            indices = nms_cv2(
-                bboxes, scores, self.conf_threshold, self.iou_threshold, self.max_det
+            bboxes, scores = self.extract(output)
+            bboxes, scores = self.compute(
+                bboxes,
+                scores,
+                conf_threshold=self.conf_threshold,
+                iou_threshold=self.iou_threshold,
+                max_det=self.max_det,
+            )
+            self.emit(output, bboxes, scores)
+
+    def extract(self, output: dai.NNData) -> tuple[np.ndarray, np.ndarray]:
+        layers = output.getAllLayerNames()
+        self._logger.debug(f"Processing input with layers: {layers}")
+        if len(layers) != 2:
+            raise ValueError(
+                f"Expected 2 output layers, got {len(layers)} layers. Please use different parser or create a new one."
             )
 
-            if len(indices) > 0:
-                bboxes = bboxes[indices]
-                scores = scores[indices]
-
-                bboxes = xyxy_to_xywh(bboxes)
+        bboxes = None
+        scores = None
 
+        for layer in layers:
+            tensor = np.asarray(output.getTensor(layer, dequantize=True))
+            if tensor.shape[-1] == 4 and tensor.ndim != 1:
+                bboxes = tensor
             else:
-                bboxes = np.array([])
-                scores = np.array([])
+                scores = tensor
 
-            message = create_detection_message(
-                bboxes=bboxes, scores=scores, label_names=self.label_names
+        if bboxes is None or scores is None:
+            raise ValueError(
+                "Bounding boxes or scores are missing in the output. Please check the NN model."
             )
-            transformation = output.getTransformation()
-            if transformation is not None:
-                message.setTransformation(transformation)
-            message.setTimestamp(output.getTimestamp())
-            message.setSequenceNum(output.getSequenceNum())
-            message.setTimestampDevice(output.getTimestampDevice())
-
-            self._logger.debug(f"Created detections message with {len(bboxes)} objects")
-            self.out.send(message)
-            self._logger.debug("Detections message sent successfully")
+
+        return bboxes.reshape(-1, 4), scores.reshape(-1)
+
+    @staticmethod
+    def compute(
+        bboxes: np.ndarray,
+        scores: np.ndarray,
+        *,
+        conf_threshold: float,
+        iou_threshold: float,
+        max_det: int,
+    ) -> tuple[np.ndarray, np.ndarray]:
+        return compute_detection_outputs(
+            bboxes,
+            scores,
+            conf_threshold=conf_threshold,
+            iou_threshold=iou_threshold,
+            max_det=max_det,
+        )
+
+    def emit(self, output: dai.NNData, bboxes: np.ndarray, scores: np.ndarray) -> None:
+        message = create_detection_message(
+            bboxes=bboxes, scores=scores, label_names=self.label_names
+        )
+        transformation = output.getTransformation()
+        if transformation is not None:
+            message.setTransformation(transformation)
+        message.setTimestamp(output.getTimestamp())
+        message.setSequenceNum(output.getSequenceNum())
+        message.setTimestampDevice(output.getTimestampDevice())
+
+        self._logger.debug(f"Created detections message with {len(bboxes)} objects")
+        self.out.send(message)
+        self._logger.debug("Detections message sent successfully")