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

.github/workflows/e2e_tests.yaml

@@ -23,9 +23,9 @@ on:
         required: true
         default: 'main'
       depthai-version:
-        description: 'Version of depthai to install. Default: 3.3.0'
+        description: 'Version of depthai to install. Default: 3.7.1'
         required: true
-        default: '3.3.0'
+        default: '3.7.1'
       platform:
         description: 'Platform to use for testing. By default both RVC2 and RVC4 are used but you can use only one by passing only `rvc2` or `rvc4`(use lowercase).'
         default: "all"
@@ -111,7 +111,8 @@ jobs:
               exec hil_runner --platforms $PLATFORM --wait $RESERVATION_OPTION --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=${{ github.event.inputs.depthai-version }} --env HUBAI_TEAM_SLUG=${{ secrets.HUBAI_TEAM_SLUG }} --env HUBAI_API_KEY=${{ secrets.HUBAI_API_KEY }} --env BRANCH=${{ github.event.inputs.depthai-nodes-version }} --env FLAGS=\"${{ github.event.inputs.additional-parameter }} -v ${{ github.event.inputs.depthai-nodes-version }} --platform ${{ matrix.platform }}\""
             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.3.0\" --env HUBAI_TEAM_SLUG=${{ secrets.HUBAI_TEAM_SLUG }} --env HUBAI_API_KEY=${{ secrets.HUBAI_API_KEY }} --env BRANCH=${{ github.ref_name }} --env FLAGS=\"-all --platform ${{ matrix.platform }}\""
+            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.7.1\" --env HUBAI_TEAM_SLUG=${{ secrets.HUBAI_TEAM_SLUG }} --env HUBAI_API_KEY=${{ secrets.HUBAI_API_KEY }} --env BRANCH=${{ github.ref_name }} --env FLAGS=\"-all --platform ${{ matrix.platform }}\""
+
           fi
 
   HIL-test-windows:
@@ -198,5 +199,5 @@ jobs:
               exec hil_runner --models 'oak4_s or oak4_pro or oak4_d' --wait $RESERVATION_OPTION -os mac --basic-sanity --sync-workspace --commands "cd /tmp/depthai-nodes && ./tests/end_to_end/run_end_to_end.sh ${{ secrets.HUBAI_API_KEY }} ${{ secrets.HUBAI_TEAM_SLUG }} ${{ github.event.inputs.depthai-version }} ${{secrets.LUXONIS_EXTRA_INDEX_URL}} rvc4"
             fi
           else
-            exec hil_runner --models 'oak4_s or oak4_pro or oak4_d' --wait --basic-sanity --reservation-name $RESERVATION_NAME -os mac --sync-workspace --commands "cd /tmp/depthai-nodes && ./tests/end_to_end/run_end_to_end.sh ${{ secrets.HUBAI_API_KEY }} ${{ secrets.HUBAI_TEAM_SLUG }} 3.3.0 ${{secrets.LUXONIS_EXTRA_INDEX_URL}} rvc4"
+            exec hil_runner --models 'oak4_s or oak4_pro or oak4_d' --wait --basic-sanity --reservation-name $RESERVATION_NAME -os mac --sync-workspace --commands "cd /tmp/depthai-nodes && ./tests/end_to_end/run_end_to_end.sh ${{ secrets.HUBAI_API_KEY }} ${{ secrets.HUBAI_TEAM_SLUG }} 3.7.1 ${{secrets.LUXONIS_EXTRA_INDEX_URL}} rvc4"
           fi

.github/workflows/stability_tests.yaml

@@ -8,9 +8,9 @@ on:
         required: true
         default: 'main'
       depthai-version:
-        description: 'Version of depthai to install. Default: 3.3.0'
+        description: 'Version of depthai to install. Default: 3.7.1'
         required: true
-        default: '3.3.0'
+        default: '3.7.1'
       duration:
         description: 'Duration of each test in seconds. Default: 10'
         required: true

depthai_nodes/message/README.md

@@ -27,8 +27,6 @@ Here are the custom message types that we introduce in this package. They are us
     - [Attributes](#attributes-9)
   - [Predictions](#predictions)
     - [Attributes](#attributes-10)
-  - [SegmentationMask](#segmentationmask)
-    - [Attributes](#attributes-11)
   - [SnapData](#snapdata)
     - [Attributes](#attributes-12)
 
@@ -136,14 +134,6 @@ Predictions class for storing predictions.
 
 - **predictions** (List\[[Prediction](#prediction)\]): List of predictions.
 
-## SegmentationMask
-
-SegmentationMask class for a single- or multi-object segmentation mask. Background is represented with "-1" and foreground classes with non-negative integers.
-
-### Attributes
-
-- **mask** (NDArray\[np.int16\]): Segmentation mask.
-
 ## SnapData
 
 SnapData class for representing a single snap event to be uploaded to DepthAI Hub.

depthai_nodes/message/__init__.py

@@ -6,14 +6,12 @@
 from .lines import Line, Lines
 from .map import Map2D
 from .prediction import Prediction, Predictions
-from .segmentation import SegmentationMask
 from .snap_data import SnapData
 
 __all__ = [
     "Line",
     "Lines",
     "Classifications",
-    "SegmentationMask",
     "Map2D",
     "Clusters",
     "Cluster",

depthai_nodes/message/creators/segmentation.py

@@ -1,19 +1,18 @@
 import numpy as np
-
-from depthai_nodes import SegmentationMask
+from depthai import SegmentationMask
 
 
 def create_segmentation_message(mask: np.ndarray) -> SegmentationMask:
     """Create a DepthAI message for segmentation mask.
 
     @param mask: Segmentation map array of shape (H, W) where each value represents a
-        segmented object class.
+        segmented object class. Index 255 represents background.
     @type mask: np.array
     @return: Segmentation mask message.
     @rtype: SegmentationMask
     @raise ValueError: If mask is not a numpy array.
     @raise ValueError: If mask is not 2D.
-    @raise ValueError: If mask is not of type int16.
+    @raise ValueError: If mask is not of type uint8.
     """
 
     if not isinstance(mask, np.ndarray):
@@ -22,10 +21,10 @@ def create_segmentation_message(mask: np.ndarray) -> SegmentationMask:
     if len(mask.shape) != 2:
         raise ValueError(f"Expected 2D input, got {len(mask.shape)}D input.")
 
-    if mask.dtype != np.int16:
-        raise ValueError(f"Expected int16 input, got {mask.dtype}.")
+    if mask.dtype != np.uint8:
+        raise ValueError(f"Expected uint8 input, got {mask.dtype}.")
 
     mask_msg = SegmentationMask()
-    mask_msg.mask = mask
+    mask_msg.setCvMask(mask)
 
     return mask_msg

depthai_nodes/message/utils/copy_message.py

@@ -40,7 +40,9 @@ def _copy_metadata(msg: dai.Buffer) -> dai.Buffer:
         if hasattr(msg, "getTimestampDevice"):
             msg_copy.setTimestampDevice(msg.getTimestampDevice())
         if hasattr(msg, "getTransformation"):
-            msg_copy.setTransformation(msg.getTransformation())
+            transformation = msg.getTransformation()
+            if transformation is not None:
+                msg_copy.setTransformation(transformation)
         return msg_copy
 
     def _copy_img_frame(img_frame: dai.ImgFrame) -> dai.ImgFrame:
@@ -49,6 +51,15 @@ def _copy_img_frame(img_frame: dai.ImgFrame) -> dai.ImgFrame:
         img_frame_copy.setCategory(img_frame.getCategory())
         return img_frame_copy
 
+    def _copy_segmentation_mask(
+        segmentation_mask: dai.SegmentationMask,
+    ) -> dai.SegmentationMask:
+        segmentation_mask_copy = _copy_metadata(segmentation_mask)
+        assert isinstance(segmentation_mask_copy, dai.SegmentationMask)
+        segmentation_mask_copy.setCvMask(segmentation_mask.getCvMask().copy())
+        segmentation_mask_copy.setLabels(segmentation_mask.getLabels())
+        return segmentation_mask_copy
+
     def _copy_img_detection(
         img_det: dai.ImgDetection | dai.SpatialImgDetection,
     ) -> dai.ImgDetection | dai.SpatialImgDetection:
@@ -125,7 +136,9 @@ def _copy_rotated_rect(rotated_rect: dai.RotatedRect) -> dai.RotatedRect:
         rotated_rect_copy.angle = rotated_rect.angle
         return rotated_rect_copy
 
-    if isinstance(msg, dai.ImgFrame):
+    if isinstance(msg, dai.SegmentationMask):
+        return _copy_segmentation_mask(msg)
+    elif isinstance(msg, dai.ImgFrame):
         return _copy_img_frame(msg)
     elif isinstance(msg, (dai.ImgDetection, dai.SpatialImgDetection)):
         return _copy_img_detection(msg)

depthai_nodes/node/README.md

@@ -37,8 +37,8 @@ Parser nodes are used to parse the output of a neural network. The main purpose
 
 ### Segmentation
 
-- `SegmentationParser`: Parser for parsing the output of the segmentation models. It will output the [`depthai_nodes.message.SegmentationMask`](../message/README.md#segmentationmask) message.
-- `FastSAMParser`: Parser for parsing the output of the FastSAM model. It will output the [`depthai_nodes.message.SegmentationMask`](../message/README.md#segmentationmask) message.
+- `SegmentationParser`: Parser for parsing the output of the segmentation models. It will output the [`dai.SegmentationMask`](https://docs.luxonis.com/software-v3/depthai/depthai-components/messages/segmentation_mask) message.
+- `FastSAMParser`: Parser for parsing the output of the FastSAM model. It will output the [`dai.SegmentationMask`](https://docs.luxonis.com/software-v3/depthai/depthai-components/messages/segmentation_mask) message.
 
 ### Feature Matching
 

depthai_nodes/node/apply_colormap.py

@@ -2,7 +2,7 @@
 import depthai as dai
 import numpy as np
 
-from depthai_nodes.message import Map2D, SegmentationMask
+from depthai_nodes.message import Map2D
 from depthai_nodes.message.utils import copy_message
 from depthai_nodes.node.base_host_node import BaseHostNode
 
@@ -25,7 +25,7 @@ class ApplyColormap(BaseHostNode):
 
     Inputs
     ------
-    frame : dai.ImgFrame | Map2D | dai.ImgDetections | SegmentationMask
+    frame : dai.ImgFrame | Map2D | dai.ImgDetections | dai.SegmentationMask
         Input message containing a 2D array to be colorized.
 
     Outputs
@@ -129,8 +129,9 @@ def _get_input_map(msg: dai.Buffer) -> np.ndarray:
 
         msg_copy = copy_message(msg)
 
-        if isinstance(msg_copy, SegmentationMask):
-            return msg_copy.mask
+        if isinstance(msg_copy, dai.SegmentationMask):
+            mask = msg_copy.getCvMask()
+            return np.where(mask == 255, 0, mask + 1)
 
         if isinstance(msg_copy, Map2D):
             return msg_copy.map
@@ -143,7 +144,7 @@ def _get_input_map(msg: dai.Buffer) -> np.ndarray:
         raise ValueError(
             f"Unsupported input type {type(msg_copy)}. "
             "ApplyColormap only accepts image-like inputs: "
-            "dai.ImgFrame, SegmentationMask, Map2D and dai.ImgDetections."
+            "dai.ImgFrame, dai.SegmentationMask, Map2D and dai.ImgDetections."
         )
 
     def _colorize(self, input_map: np.ndarray) -> np.ndarray:

depthai_nodes/node/parser_generator.py

@@ -79,7 +79,7 @@ def build(
                     parser_name = self._getHostParserName(parser_name)
                 else:
                     parser = pipeline.create(dai.node.DetectionParser)
-                    parser.setNNArchive(nnArchive)
+                    parser.setNNArchiveHead(head)
                     parsers[index] = parser
                     continue
             elif parser_name == "YOLOExtendedParser":
@@ -91,7 +91,7 @@ def build(
                         and not self._has_non_default_yolo_strides(head, yolo_subtype)
                     ):
                         parser = pipeline.create(dai.node.DetectionParser)
-                        parser.setNNArchive(nnArchive)
+                        parser.setNNArchiveHead(head)
                         if head.metadata.maskOutputs is not None and is_rvc2_device:
                             self._logger.warning(
                                 "Segmentation based model detected with RVC2 device. Parsing will be done on the host machine."
@@ -100,6 +100,18 @@ def build(
                         parsers[index] = parser
                         continue
 
+            elif parser_name == "SegmentationParser" and not hostOnly:
+                parser = pipeline.create(dai.node.SegmentationParser)
+                parser.setNNArchiveHead(head)
+                if is_rvc2_device:
+                    self._logger.warning(
+                        "Segmentation model detected with RVC2 device. Parsing will "
+                        "be done on the host machine."
+                    )
+                    parser.setRunOnHost(True)
+                parsers[index] = parser
+                continue
+
             parser = globals().get(parser_name)
 
             if parser is None: