CodeRabbitAI suggested fixes + miscellaneous fixes

depthai_nodes/node/parsers/fastsam.py

@@ -307,7 +307,7 @@ def run(self):
                 protos_output,
                 protos_len,
             ) = self.extract(output)
-            results_masks = self.compute(
+            results_masks, mask_count = self.compute(
                 outputs_values,
                 masks_outputs_values,
                 protos_output,
@@ -321,7 +321,7 @@ def run(self):
                 point_label=self.point_label,
                 bbox=self.bbox,
             )
-            self.emit(output, results_masks)
+            self.emit(output, results_masks, mask_count)
 
     def extract(
         self, output: dai.NNData
@@ -357,7 +357,7 @@ def compute(
         points: tuple[int, int] | None,
         point_label: int | None,
         bbox: tuple[int, int, int, int] | None,
-    ) -> np.ndarray:
+    ) -> tuple[np.ndarray, int]:
         return compute_fastsam_mask(
             outputs_values,
             masks_outputs_values,
@@ -373,16 +373,16 @@ def compute(
             bbox=bbox,
         )
 
-    def emit(self, output: dai.NNData, results_masks: np.ndarray) -> None:
+    def emit(
+        self, output: dai.NNData, results_masks: np.ndarray, mask_count: int
+    ) -> None:
         segmentation_message = create_segmentation_message(results_masks)
         transformation = output.getTransformation()
         if transformation is not None:
             segmentation_message.setTransformation(transformation)
         segmentation_message.setTimestamp(output.getTimestamp())
         segmentation_message.setSequenceNum(output.getSequenceNum())
         segmentation_message.setTimestampDevice(output.getTimestampDevice())
-        self._logger.debug(
-            f"Created segmentation message with {len(results_masks)} masks"
-        )
+        self._logger.debug(f"Created segmentation message with {mask_count} masks")
         self.out.send(segmentation_message)
         self._logger.debug("Segmentation message sent successfully")

depthai_nodes/node/parsers/scrfd.py

@@ -74,13 +74,16 @@ def __init__(
         self.num_anchors = num_anchors
         self.input_size = input_size
         self.label_names = ["Face"]
-        self._cached_anchors = compute_anchor_centers(
-            self.feat_stride_fpn, self.input_size, self.num_anchors
-        )
+        self._refresh_cached_anchors()
         self._logger.debug(
             f"SCRFDParser initialized with output_layer_names={output_layer_names}, conf_threshold={conf_threshold}, iou_threshold={iou_threshold}, max_det={max_det}, input_size={input_size}, feat_stride_fpn={feat_stride_fpn}, num_anchors={num_anchors}"
         )
 
+    def _refresh_cached_anchors(self) -> None:
+        self._cached_anchors = compute_anchor_centers(
+            self.feat_stride_fpn, self.input_size, self.num_anchors
+        )
+
     def setOutputLayerNames(self, output_layer_names: list[str]) -> None:
         """Sets the output layer name(s) for the parser.
 
@@ -105,6 +108,7 @@ def setInputSize(self, input_size: tuple[int, int]) -> None:
         if not all(isinstance(size, int) for size in input_size):
             raise ValueError("Input size must be a tuple of integers.")
         self.input_size = input_size
+        self._refresh_cached_anchors()
         self._logger.debug(f"Input size set to {self.input_size}")
 
     def setFeatStrideFPN(self, feat_stride_fpn: list[int]) -> None:
@@ -118,6 +122,7 @@ def setFeatStrideFPN(self, feat_stride_fpn: list[int]) -> None:
         if not all(isinstance(stride, int) for stride in feat_stride_fpn):
             raise ValueError("Feature stride must be a list of integers.")
         self.feat_stride_fpn = feat_stride_fpn
+        self._refresh_cached_anchors()
         self._logger.debug(f"Feature stride set to {self.feat_stride_fpn}")
 
     def setNumAnchors(self, num_anchors: int) -> None:
@@ -129,6 +134,7 @@ def setNumAnchors(self, num_anchors: int) -> None:
         if not isinstance(num_anchors, int):
             raise ValueError("Number of anchors must be an integer.")
         self.num_anchors = num_anchors
+        self._refresh_cached_anchors()
         self._logger.debug(f"Number of anchors set to {self.num_anchors}")
 
     def build(
@@ -160,9 +166,7 @@ def build(
         self.output_layer_names = output_layers
         self.feat_stride_fpn = head_config.get("feat_stride_fpn", self.feat_stride_fpn)
         self.num_anchors = head_config.get("num_anchors", self.num_anchors)
-        self._cached_anchors = compute_anchor_centers(
-            self.feat_stride_fpn, self.input_size, self.num_anchors
-        )
+        self._refresh_cached_anchors()
 
         self._logger.debug(
             f"SCRFDParser built with output_layer_names={self.output_layer_names}, feat_stride_fpn={self.feat_stride_fpn}, num_anchors={self.num_anchors}"

depthai_nodes/node/parsers/segmentation.py

@@ -54,6 +54,8 @@ class segmentation model. Default is False. If True, the parser will use
         self.output_layer_name = output_layer_name
         self.classes_in_one_layer = classes_in_one_layer
         self.background_class = background_class
+        self.class_names = None
+        self.n_classes = 0
         self._background_class_ignored_warning_sent = False
         self._logger.debug(
             "SegmentationParser initialized with "
@@ -106,6 +108,13 @@ def _warn_if_background_class_ignored(self) -> None:
             )
             self._background_class_ignored_warning_sent = True
 
+    def _get_logged_class_count(self, class_map: np.ndarray) -> int:
+        if self.class_names is not None:
+            return len(self.class_names)
+        if self.n_classes > 0:
+            return self.n_classes
+        return int(np.unique(class_map[class_map != 255]).size)
+
     def build(
         self,
         head_config: dict[str, Any],
@@ -130,6 +139,10 @@ def build(
         self.background_class = head_config.get(
             "background_class", self.background_class
         )
+        self.class_names = head_config.get("classes", self.class_names)
+        self.n_classes = head_config.get("n_classes", self.n_classes)
+        if self.n_classes == 0 and self.class_names is not None:
+            self.n_classes = len(self.class_names)
 
         self._logger.debug(
             "SegmentationParser built with "
@@ -196,7 +209,7 @@ def emit(self, output: dai.NNData, class_map: np.ndarray) -> None:
             mask_message.setTransformation(transformation)
 
         self._logger.debug(
-            f"Created segmentation message with {class_map.shape[0]} classes"
+            f"Created segmentation message with {self._get_logged_class_count(class_map)} classes"
         )
         self.out.send(mask_message)
         self._logger.debug("Segmentation message sent successfully")

depthai_nodes/node/parsers/utils/detection.py

@@ -13,9 +13,12 @@ def compute_detection_outputs(
     max_det: int,
 ) -> tuple[np.ndarray, np.ndarray]:
     """Filter detection outputs and convert boxes to center-width-height format."""
-    indices = nms_cv2(bboxes, scores, conf_threshold, iou_threshold, max_det)
+    nms_bboxes = np.column_stack((bboxes[:, :2], bboxes[:, 2:] - bboxes[:, :2]))
+    indices = np.asarray(
+        nms_cv2(nms_bboxes, scores, conf_threshold, iou_threshold, max_det)
+    ).reshape(-1)
 
-    if len(indices) == 0:
+    if indices.size == 0:
         return np.array([]), np.array([])
 
     filtered_bboxes = xyxy_to_xywh(bboxes[indices])

depthai_nodes/node/parsers/utils/fastsam.py

@@ -407,8 +407,8 @@ def compute_fastsam_mask(
     points: tuple[int, int] | None,
     point_label: int | None,
     bbox: tuple[int, int, int, int] | None,
-) -> np.ndarray:
-    """Decode FastSAM outputs into a merged segmentation mask."""
+) -> tuple[np.ndarray, int]:
+    """Decode FastSAM outputs into a merged segmentation mask and mask count."""
     width = outputs_values[0].shape[3] * 8
     height = outputs_values[0].shape[2] * 8
     input_shape = (width, height)
@@ -462,5 +462,8 @@ def compute_fastsam_mask(
 
     if len(results_masks) == 0:
         results_masks = np.full((1, height, width), -1, dtype=np.int16)
+        results_mask_count = 0
+    else:
+        results_mask_count = len(results_masks)
 
-    return merge_masks(results_masks)
+    return merge_masks(results_masks), results_mask_count

depthai_nodes/node/parsers/utils/image_output.py

@@ -7,7 +7,7 @@ def compute_image_output(output_image: np.ndarray) -> np.ndarray:
     """Convert a model image output tensor into an image array."""
     image = np.asarray(output_image)
 
-    if image.shape[0] == 1:
+    if image.ndim == 4 and image.shape[0] == 1:
         image = image[0]
 
     if image.ndim != 3:

depthai_nodes/node/parsers/utils/map_output.py

@@ -4,6 +4,16 @@
 def compute_map_output(map_tensor: np.ndarray) -> np.ndarray:
     """Return the model map output without the batch dimension."""
     map_output = np.asarray(map_tensor)
-    if map_output.shape[0] == 1:
+
+    while map_output.ndim > 2 and map_output.shape[0] == 1:
         map_output = map_output[0]
-    return map_output
+
+    if map_output.ndim == 2:
+        return map_output
+
+    if map_output.ndim == 3 and map_output.shape[-1] == 1:
+        return map_output[:, :, 0]
+
+    raise ValueError(
+        f"Expected HW, NHW, or HWN with singleton N; got {map_output.shape}."
+    )

depthai_nodes/node/parsers/utils/yolo.py

@@ -525,6 +525,12 @@ def decode_yolo_output(
     @return: NMS output.
     @rtype: np.ndarray
     """
+    if len(strides) != len(yolo_outputs):
+        raise ValueError(
+            "Number of `strides` must match number of YOLO outputs. "
+            f"Got {len(strides)} strides for {len(yolo_outputs)} outputs."
+        )
+
     num_outputs = num_classes + 5
 
     # 1. Parse and concatenate all head outputs efficiently
@@ -698,7 +704,8 @@ def compute_yolo_detections(
         max_instance_id = np.iinfo(np.int16).max
         if results.shape[0] > max_instance_id + 1:
             raise ValueError(
-                "YOLO segmentation can encode at most 32768 instances in SegmentationMask."
+                "YOLO segmentation can encode at most "
+                f"{max_instance_id + 1} instance masks in SegmentationMask."
             )
         final_mask = np.full(input_shape, -1, dtype=np.int16)