diff --git a/depthai_nodes/node/parsers/fastsam.py b/depthai_nodes/node/parsers/fastsam.py index a9804db0..17cec1a6 100644 --- a/depthai_nodes/node/parsers/fastsam.py +++ b/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,7 +373,9 @@ 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: @@ -381,8 +383,6 @@ def emit(self, output: dai.NNData, results_masks: np.ndarray) -> None: 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") diff --git a/depthai_nodes/node/parsers/scrfd.py b/depthai_nodes/node/parsers/scrfd.py index b1d38308..a174e36e 100644 --- a/depthai_nodes/node/parsers/scrfd.py +++ b/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}" diff --git a/depthai_nodes/node/parsers/segmentation.py b/depthai_nodes/node/parsers/segmentation.py index 8624d8bd..91096784 100644 --- a/depthai_nodes/node/parsers/segmentation.py +++ b/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") diff --git a/depthai_nodes/node/parsers/utils/detection.py b/depthai_nodes/node/parsers/utils/detection.py index b019132d..21d91b6a 100644 --- a/depthai_nodes/node/parsers/utils/detection.py +++ b/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]) diff --git a/depthai_nodes/node/parsers/utils/fastsam.py b/depthai_nodes/node/parsers/utils/fastsam.py index 88d68d5e..805e166b 100644 --- a/depthai_nodes/node/parsers/utils/fastsam.py +++ b/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 diff --git a/depthai_nodes/node/parsers/utils/image_output.py b/depthai_nodes/node/parsers/utils/image_output.py index bd944ede..1276251a 100644 --- a/depthai_nodes/node/parsers/utils/image_output.py +++ b/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: diff --git a/depthai_nodes/node/parsers/utils/map_output.py b/depthai_nodes/node/parsers/utils/map_output.py index e04ee141..4e7d72a2 100644 --- a/depthai_nodes/node/parsers/utils/map_output.py +++ b/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}." + ) diff --git a/depthai_nodes/node/parsers/utils/yolo.py b/depthai_nodes/node/parsers/utils/yolo.py index 393a5db0..97bd94f0 100644 --- a/depthai_nodes/node/parsers/utils/yolo.py +++ b/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)