From 607d093c200efa7bb9fccecdbea652cd7e5e8df6 Mon Sep 17 00:00:00 2001 From: Dani Rogmans Date: Mon, 29 Jun 2026 23:18:44 +0200 Subject: [PATCH 01/11] extract-compute-emit pattern for classificationparser --- depthai_nodes/node/parsers/classification.py | 70 +++++++++++-------- depthai_nodes/node/parsers/utils/__init__.py | 2 + .../node/parsers/utils/classification.py | 15 ++++ depthai_nodes/node/parsers/utils/yolo.py | 8 +-- depthai_nodes/node/parsers/yolo.py | 14 ++-- 5 files changed, 70 insertions(+), 39 deletions(-) create mode 100644 depthai_nodes/node/parsers/utils/classification.py diff --git a/depthai_nodes/node/parsers/classification.py b/depthai_nodes/node/parsers/classification.py index 61b69155..6b08ac45 100644 --- a/depthai_nodes/node/parsers/classification.py +++ b/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,48 @@ 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 = np.asarray( + 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") diff --git a/depthai_nodes/node/parsers/utils/__init__.py b/depthai_nodes/node/parsers/utils/__init__.py index 86d07455..3195d22e 100644 --- a/depthai_nodes/node/parsers/utils/__init__.py +++ b/depthai_nodes/node/parsers/utils/__init__.py @@ -7,6 +7,7 @@ xywh_to_xyxy, xyxy_to_xywh, ) +from .classification import compute_classification_scores from .decode_head import decode_head from .denormalize import unnormalize_image @@ -19,6 +20,7 @@ "xyxy_to_xywh", "normalize_bboxes", "top_left_wh_to_xywh", + "compute_classification_scores", "softmax", "sigmoid", ] diff --git a/depthai_nodes/node/parsers/utils/classification.py b/depthai_nodes/node/parsers/utils/classification.py new file mode 100644 index 00000000..bc802d5a --- /dev/null +++ b/depthai_nodes/node/parsers/utils/classification.py @@ -0,0 +1,15 @@ +import numpy as np + +from .activations import softmax + + +def compute_classification_scores( + scores: np.ndarray, + *, + is_softmax: bool = True, +) -> np.ndarray: + """Return classification scores, applying softmax when needed.""" + computed_scores = np.asarray(scores).flatten() + if not is_softmax: + computed_scores = softmax(computed_scores) + return computed_scores diff --git a/depthai_nodes/node/parsers/utils/yolo.py b/depthai_nodes/node/parsers/utils/yolo.py index 409a2082..d8552559 100644 --- a/depthai_nodes/node/parsers/utils/yolo.py +++ b/depthai_nodes/node/parsers/utils/yolo.py @@ -400,10 +400,10 @@ def decode_yolo26( ) -> tuple[np.ndarray, np.ndarray | None]: """Decode YOLO26 output for detection, segmentation, or pose. - YOLO26 end2end output is already decoded (xyxy in pixels) with a pre-computed - confidence score (ReduceMax over class scores). Needs topk and conf thresholding. - Optionally filters an auxiliary tensor (mask coefficients or keypoints) with the - detections. + YOLO26 end2end output is already decoded to xyxy pixel boxes and includes a + pre-computed confidence score (ReduceMax over class scores) in column 4. This + path only applies confidence thresholding and top-k filtering. It can also filter + an auxiliary tensor such as mask coefficients or keypoints using the kept rows. @param raw: Raw detection tensor (N, A, 5+nc) where columns are [x1, y1, x2, y2, conf, cls_0, ..., cls_nc-1]. diff --git a/depthai_nodes/node/parsers/yolo.py b/depthai_nodes/node/parsers/yolo.py index 21ac7da5..01ec1d81 100644 --- a/depthai_nodes/node/parsers/yolo.py +++ b/depthai_nodes/node/parsers/yolo.py @@ -320,12 +320,14 @@ def build( ) from err if self.subtype == YOLOSubtype.V26: - # YOLO26 end2end: task type inferred from output layer names at runtime. - # Boxes are decoded to xyxy pixels, class scores are sigmoided. No NMS needed. - # - Detection: 'output_yolo26' (N, A, 5+nc) - # - Segmentation: 'output_yolo26' (N, A, 5+nc), 'output_masks' (N, A, nm), - # 'protos_output' (N, nm, proto_h, proto_w) - # - Pose: 'output_yolo26' (N, A, 5+nc), 'kpt_output' (N, A, nk) + # YOLO26 end2end: task type is inferred from output layer names at runtime. + # The detection tensor is already decoded and has shape (N, A, 5+nc): + # [x1, y1, x2, y2, conf, cls_0, ..., cls_nc-1]. + # Auxiliary outputs depend on the task: + # - Detection: 'output_yolo26' + # - Segmentation: 'output_yolo26' plus 'output_masks' (N, A, nm) and + # 'protos_output' (N, nm, proto_h, proto_w) + # - Pose: 'output_yolo26' plus 'kpt_output' (N, A, nk) kps_layer_names = [name for name in output_layers if "kpt_output" in name] masks_layer_names = [ name for name in output_layers if "output_masks" in name From 0f4bd7e8f3b3464aee6812e0934b40460671f890 Mon Sep 17 00:00:00 2001 From: Dani Rogmans Date: Tue, 30 Jun 2026 11:24:58 +0200 Subject: [PATCH 02/11] extract-compute-emit for all parsers in dai nodes --- .../node/parsers/classification_sequence.py | 83 ++-- depthai_nodes/node/parsers/detection.py | 108 +++--- depthai_nodes/node/parsers/embeddings.py | 41 +- depthai_nodes/node/parsers/fastsam.py | 161 ++++---- depthai_nodes/node/parsers/hrnet.py | 98 ++--- depthai_nodes/node/parsers/image_output.py | 80 ++-- depthai_nodes/node/parsers/keypoints.py | 84 ++-- depthai_nodes/node/parsers/lane_detection.py | 83 ++-- depthai_nodes/node/parsers/map_output.py | 54 +-- .../node/parsers/mediapipe_palm_detection.py | 164 ++++---- depthai_nodes/node/parsers/mlsd.py | 84 ++-- depthai_nodes/node/parsers/ppdet.py | 96 +++-- depthai_nodes/node/parsers/regression.py | 58 +-- depthai_nodes/node/parsers/rf_detr.py | 209 ++++------ depthai_nodes/node/parsers/scrfd.py | 166 ++++---- depthai_nodes/node/parsers/segmentation.py | 107 ++--- .../node/parsers/superanimal_landmarker.py | 87 +++-- .../parsers/utils/classification_sequence.py | 34 ++ depthai_nodes/node/parsers/utils/detection.py | 23 ++ .../node/parsers/utils/embeddings.py | 3 + depthai_nodes/node/parsers/utils/fastsam.py | 73 ++++ depthai_nodes/node/parsers/utils/hrnet.py | 27 ++ .../node/parsers/utils/image_output.py | 16 + depthai_nodes/node/parsers/utils/keypoints.py | 20 + .../node/parsers/utils/lane_detection.py | 22 ++ .../node/parsers/utils/map_output.py | 9 + depthai_nodes/node/parsers/utils/medipipe.py | 83 ++++ depthai_nodes/node/parsers/utils/mlsd.py | 17 + depthai_nodes/node/parsers/utils/ppdet.py | 19 + .../node/parsers/utils/regression.py | 7 + depthai_nodes/node/parsers/utils/rf_detr.py | 86 +++++ depthai_nodes/node/parsers/utils/scrfd.py | 33 ++ .../node/parsers/utils/segmentation.py | 46 +++ .../node/parsers/utils/superanimal.py | 17 + depthai_nodes/node/parsers/utils/xfeat.py | 29 ++ depthai_nodes/node/parsers/utils/yolo.py | 199 +++++++++- depthai_nodes/node/parsers/utils/yunet.py | 62 +++ depthai_nodes/node/parsers/xfeat.py | 237 +++++++----- depthai_nodes/node/parsers/yolo.py | 365 ++++++------------ depthai_nodes/node/parsers/yunet.py | 252 +++++------- 40 files changed, 2075 insertions(+), 1367 deletions(-) create mode 100644 depthai_nodes/node/parsers/utils/classification_sequence.py create mode 100644 depthai_nodes/node/parsers/utils/detection.py create mode 100644 depthai_nodes/node/parsers/utils/embeddings.py create mode 100644 depthai_nodes/node/parsers/utils/hrnet.py create mode 100644 depthai_nodes/node/parsers/utils/image_output.py create mode 100644 depthai_nodes/node/parsers/utils/lane_detection.py create mode 100644 depthai_nodes/node/parsers/utils/map_output.py create mode 100644 depthai_nodes/node/parsers/utils/regression.py create mode 100644 depthai_nodes/node/parsers/utils/rf_detr.py create mode 100644 depthai_nodes/node/parsers/utils/segmentation.py diff --git a/depthai_nodes/node/parsers/classification_sequence.py b/depthai_nodes/node/parsers/classification_sequence.py index 4cdef04d..e5697f1f 100644 --- a/depthai_nodes/node/parsers/classification_sequence.py +++ b/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 @@ -158,47 +160,48 @@ def run(self): 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 + 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." ) - 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, - ) - msg.setTimestamp(output.getTimestamp()) - msg.setSequenceNum(output.getSequenceNum()) - msg.setTimestampDevice(output.getTimestampDevice()) - transformation = output.getTransformation() - if transformation is not None: - msg.setTransformation(transformation) + if self.n_classes == 0: + raise ValueError("Classes must be provided for classification.") - self._logger.debug(f"Created message with {len(msg.classes)} classes") + return output.getTensor(self.output_layer_name, dequantize=True).astype( + np.float32 + ) - self.out.send(msg) + @staticmethod + def compute(scores: np.ndarray, *, is_softmax: bool = True) -> np.ndarray: + return compute_classification_sequence_scores( + scores, is_softmax=is_softmax + ) - self._logger.debug("Classification message sent successfully") + 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") diff --git a/depthai_nodes/node/parsers/detection.py b/depthai_nodes/node/parsers/detection.py index fcb30a7d..7f9a8d96 100644 --- a/depthai_nodes/node/parsers/detection.py +++ b/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") diff --git a/depthai_nodes/node/parsers/embeddings.py b/depthai_nodes/node/parsers/embeddings.py index 4c780345..3dae4e4a 100644 --- a/depthai_nodes/node/parsers/embeddings.py +++ b/depthai_nodes/node/parsers/embeddings.py @@ -3,6 +3,7 @@ import depthai as dai from depthai_nodes.node.parsers.base_parser import BaseParser +from depthai_nodes.node.parsers.utils.embeddings import compute_embeddings_output class EmbeddingsParser(BaseParser): @@ -68,20 +69,30 @@ def run(self): except dai.MessageQueue.QueueException: break # Pipeline was stopped, no more data - # Get all the layer names - output_names = self.output_layer_name or output.getAllLayerNames() - self._logger.debug(f"Processing input with layers: {output_names}") + extracted = self.extract(output) + computed = self.compute(extracted) + self.emit(computed) - assert ( - len(output_names) == 1 - ), "Embeddings head should have only one output layer" - output.setSequenceNum(output.getSequenceNum()) - output.setTimestamp(output.getTimestamp()) - output.setTimestampDevice(output.getTimestampDevice()) - transformation = output.getTransformation() - if transformation is not None: - output.setTransformation(transformation) + def extract(self, output: dai.NNData) -> dai.NNData: + output_names = self.output_layer_name or output.getAllLayerNames() + self._logger.debug(f"Processing input with layers: {output_names}") - self.out.send(output) - - self._logger.debug("Message sent successfully") + assert ( + len(output_names) == 1 + ), "Embeddings head should have only one output layer" + return output + + @staticmethod + def compute(output: dai.NNData) -> dai.NNData: + return compute_embeddings_output(output) + + def emit(self, output: dai.NNData) -> None: + output.setSequenceNum(output.getSequenceNum()) + output.setTimestamp(output.getTimestamp()) + output.setTimestampDevice(output.getTimestampDevice()) + transformation = output.getTransformation() + if transformation is not None: + output.setTransformation(transformation) + + self.out.send(output) + self._logger.debug("Message sent successfully") diff --git a/depthai_nodes/node/parsers/fastsam.py b/depthai_nodes/node/parsers/fastsam.py index 5a519302..c453ec89 100644 --- a/depthai_nodes/node/parsers/fastsam.py +++ b/depthai_nodes/node/parsers/fastsam.py @@ -8,6 +8,7 @@ from depthai_nodes.node.parsers.utils.fastsam import ( box_prompt, build_mask_coeffs, + compute_fastsam_mask, decode_fastsam_output, merge_masks, point_prompt, @@ -306,93 +307,89 @@ def run(self): except dai.MessageQueue.QueueException: break # Pipeline was stopped, no more data - outputs_names = sorted([name for name in self.yolo_outputs]) - self._logger.debug(f"Processing input with layers: {outputs_names}") - outputs_values = [ - output.getTensor( - o, dequantize=True, storageOrder=dai.TensorInfo.StorageOrder.NCHW - ).astype(np.float32) - for o in outputs_names - ] - - # Get the segmentation outputs ( + outputs_values, masks_outputs_values, protos_output, protos_len, - ) = get_segmentation_outputs(output, self.mask_outputs, self.protos_output) - - # determine the input shape of the model from the first output - width = outputs_values[0].shape[3] * 8 - height = outputs_values[0].shape[2] * 8 - input_shape = (width, height) - - # Decode the outputs - results = decode_fastsam_output( + ) = self.extract(output) + results_masks = self.compute( outputs_values, - [8, 16, 32], - [None, None, None], - img_shape=input_shape[::-1], - conf_thres=self.conf_threshold, - iou_thres=self.iou_threshold, - num_classes=self.n_classes, - ) - - if results.shape[0] == 0: - results_bboxes = np.array([]) - results_masks = np.array([]) - else: - results_bboxes = np.concatenate( - [ - results[:, :4].astype(int), - results[:, 4:5], - results[:, 5:6].astype(int), - ], - axis=1, - ) - mask_coeffs = build_mask_coeffs( - parsed_results=results, - masks_outputs_values=masks_outputs_values, - protos_len=protos_len, - ) - results_masks = process_masks( - parsed_results=results, - mask_coeffs=mask_coeffs, - protos=protos_output[0], - orig_shape=input_shape, - mask_conf=self.mask_conf, - ) - - if self.prompt == "bbox": - results_masks = box_prompt( - results_masks, bbox=self.bbox, orig_shape=input_shape[::-1] - ) - - elif self.prompt == "point": - results_masks = point_prompt( - results_bboxes, - results_masks, - points=self.points, - pointlabel=self.point_label, - orig_shape=input_shape[::-1], - ) - - if len(results_masks) == 0: - results_masks = np.full((1, height, width), -1, dtype=np.int16) - results_masks = merge_masks(results_masks) - - 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" + masks_outputs_values, + protos_output, + protos_len, + conf_threshold=self.conf_threshold, + n_classes=self.n_classes, + iou_threshold=self.iou_threshold, + mask_conf=self.mask_conf, + prompt=self.prompt, + points=self.points, + point_label=self.point_label, + bbox=self.bbox, ) + self.emit(output, results_masks) + + def extract( + self, output: dai.NNData + ) -> tuple[list[np.ndarray], list[np.ndarray], np.ndarray, int]: + outputs_names = sorted([name for name in self.yolo_outputs]) + self._logger.debug(f"Processing input with layers: {outputs_names}") + outputs_values = [ + output.getTensor( + o, dequantize=True, storageOrder=dai.TensorInfo.StorageOrder.NCHW + ).astype(np.float32) + for o in outputs_names + ] + + ( + masks_outputs_values, + protos_output, + protos_len, + ) = get_segmentation_outputs(output, self.mask_outputs, self.protos_output) + return outputs_values, masks_outputs_values, protos_output, protos_len + + @staticmethod + def compute( + outputs_values: list[np.ndarray], + masks_outputs_values: list[np.ndarray], + protos_output: np.ndarray, + protos_len: int, + *, + conf_threshold: float, + n_classes: int, + iou_threshold: float, + mask_conf: float, + prompt: str, + points: tuple[int, int] | None, + point_label: int | None, + bbox: tuple[int, int, int, int] | None, + ) -> np.ndarray: + return compute_fastsam_mask( + outputs_values, + masks_outputs_values, + protos_output, + protos_len, + conf_threshold=conf_threshold, + n_classes=n_classes, + iou_threshold=iou_threshold, + mask_conf=mask_conf, + prompt=prompt, + points=points, + point_label=point_label, + bbox=bbox, + ) - self.out.send(segmentation_message) + def emit(self, output: dai.NNData, results_masks: np.ndarray) -> 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("Segmentation message sent successfully") + self._logger.debug( + f"Created segmentation message with {len(results_masks)} masks" + ) + self.out.send(segmentation_message) + self._logger.debug("Segmentation message sent successfully") diff --git a/depthai_nodes/node/parsers/hrnet.py b/depthai_nodes/node/parsers/hrnet.py index fdf108e9..9f77fe45 100644 --- a/depthai_nodes/node/parsers/hrnet.py +++ b/depthai_nodes/node/parsers/hrnet.py @@ -5,6 +5,7 @@ from depthai_nodes.message.creators import create_keypoints_message from depthai_nodes.node.parsers.keypoints import KeypointParser +from depthai_nodes.node.parsers.utils.hrnet import compute_hrnet_keypoints class HRNetParser(KeypointParser): @@ -95,63 +96,48 @@ def run(self): except dai.MessageQueue.QueueException: break # Pipeline was stopped - 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." - ) - - heatmaps = output.getTensor( - self.output_layer_name, - dai.TensorInfo.StorageOrder.NCHW, # this signals DAI to return output as NCHW - dequantize=True, + heatmaps = self.extract(output) + keypoints, scores = self.compute(heatmaps) + self.emit(output, keypoints, 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." ) - if heatmaps.shape[0] == 1: - heatmaps = heatmaps[0] # remove batch dimension - - if len(heatmaps.shape) != 3: - raise ValueError( - f"Expected 3D output tensor, got {len(heatmaps.shape)}D." - ) - - self.n_keypoints, map_h, map_w = heatmaps.shape - - scores = np.array([np.max(heatmap) for heatmap in heatmaps]) - scores = np.clip(scores, 0, 1) # TODO: check why scores are sometimes >1 - - keypoints = np.array( - [ - np.unravel_index(heatmap.argmax(), heatmap.shape) - for heatmap in heatmaps - ] - ) - keypoints = keypoints.astype(np.float32) - keypoints = keypoints[:, ::-1] / np.array( - [map_w, map_h] - ) # normalize keypoints to [0, 1] - - keypoints_message = create_keypoints_message( - keypoints=keypoints, - scores=scores, - confidence_threshold=self.score_threshold, - edges=self.edges, - label_names=self.label_names, - ) - keypoints_message.setTimestamp(output.getTimestamp()) - keypoints_message.setSequenceNum(output.getSequenceNum()) - keypoints_message.setTimestampDevice(output.getTimestampDevice()) - transformation = output.getTransformation() - if transformation is not None: - keypoints_message.setTransformation(transformation) - - self._logger.debug( - f"Created keypoints message with {len(keypoints)} points" - ) + return output.getTensor( + self.output_layer_name, + dai.TensorInfo.StorageOrder.NCHW, + dequantize=True, + ) - self.out.send(keypoints_message) + def compute(self, heatmaps: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + keypoints, scores = compute_hrnet_keypoints(heatmaps) + self.n_keypoints = len(keypoints) + return keypoints, scores - self._logger.debug("Keypoint output sent successfully") + def emit( + self, output: dai.NNData, keypoints: np.ndarray, scores: np.ndarray + ) -> None: + keypoints_message = create_keypoints_message( + keypoints=keypoints, + scores=scores, + confidence_threshold=self.score_threshold, + edges=self.edges, + label_names=self.label_names, + ) + keypoints_message.setTimestamp(output.getTimestamp()) + keypoints_message.setSequenceNum(output.getSequenceNum()) + keypoints_message.setTimestampDevice(output.getTimestampDevice()) + transformation = output.getTransformation() + if transformation is not None: + keypoints_message.setTransformation(transformation) + + self._logger.debug(f"Created keypoints message with {len(keypoints)} points") + self.out.send(keypoints_message) + self._logger.debug("Keypoint output sent successfully") diff --git a/depthai_nodes/node/parsers/image_output.py b/depthai_nodes/node/parsers/image_output.py index 962c77dc..4cdd9eaf 100644 --- a/depthai_nodes/node/parsers/image_output.py +++ b/depthai_nodes/node/parsers/image_output.py @@ -4,7 +4,7 @@ from depthai_nodes.message.creators import create_image_message from depthai_nodes.node.parsers.base_parser import BaseParser -from depthai_nodes.node.parsers.utils import unnormalize_image +from depthai_nodes.node.parsers.utils.image_output import compute_image_output class ImageOutputParser(BaseParser): @@ -102,45 +102,43 @@ def run(self): except dai.MessageQueue.QueueException: break # Pipeline was stopped - 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." - ) - - output_image = output.getTensor(self.output_layer_name, dequantize=True) - - if output_image.shape[0] == 1: - output_image = output_image[0] # remove batch dimension - - if len(output_image.shape) != 3: - raise ValueError( - f"Expected 3D output tensor, got {len(output_image.shape)}D." - ) - - image = unnormalize_image(output_image) - - image_message = create_image_message( - image=image, - is_bgr=self.output_is_bgr, - img_frame_type=( - dai.ImgFrame.Type.BGR888p - if self._platform == "RVC2" - else dai.ImgFrame.Type.BGR888i - ), - ) - image_message.setTimestamp(output.getTimestamp()) - image_message.setSequenceNum(output.getSequenceNum()) - image_message.setTimestampDevice(output.getTimestampDevice()) - transformation = output.getTransformation() - if transformation is not None: - image_message.setTransformation(transformation) - - self._logger.debug(f"Created image message with shape {image.shape}") + output_image = self.extract(output) + image = self.compute(output_image) + self.emit(output, image) - self.out.send(image_message) + def extract(self, output: dai.NNData): + 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." + ) - self._logger.debug("Image message sent successfully") + return output.getTensor(self.output_layer_name, dequantize=True) + + @staticmethod + def compute(output_image): + return compute_image_output(output_image) + + def emit(self, output: dai.NNData, image) -> None: + image_message = create_image_message( + image=image, + is_bgr=self.output_is_bgr, + img_frame_type=( + dai.ImgFrame.Type.BGR888p + if self._platform == "RVC2" + else dai.ImgFrame.Type.BGR888i + ), + ) + image_message.setTimestamp(output.getTimestamp()) + image_message.setSequenceNum(output.getSequenceNum()) + image_message.setTimestampDevice(output.getTimestampDevice()) + transformation = output.getTransformation() + if transformation is not None: + image_message.setTransformation(transformation) + + self._logger.debug(f"Created image message with shape {image.shape}") + self.out.send(image_message) + self._logger.debug("Image message sent successfully") diff --git a/depthai_nodes/node/parsers/keypoints.py b/depthai_nodes/node/parsers/keypoints.py index d2de801c..4a3a97b9 100644 --- a/depthai_nodes/node/parsers/keypoints.py +++ b/depthai_nodes/node/parsers/keypoints.py @@ -5,6 +5,7 @@ from depthai_nodes.message.creators import create_keypoints_message from depthai_nodes.node.parsers.base_parser import BaseParser +from depthai_nodes.node.parsers.utils.keypoints import compute_keypoints class KeypointParser(BaseParser): @@ -209,45 +210,52 @@ def run(self): except dai.MessageQueue.QueueException: break # Pipeline was stopped - 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." - ) - - keypoints = output.getTensor( - self.output_layer_name, dequantize=True - ).astype(np.float32) - num_coords = int(np.prod(keypoints.shape) / self.n_keypoints) - - if num_coords not in [2, 3]: - raise ValueError( - f"Expected 2 or 3 coordinates per keypoint, got {num_coords}." - ) - - keypoints = keypoints.reshape(self.n_keypoints, num_coords) - - keypoints /= self.scale_factor - - keypoints = np.clip(keypoints, 0, 1) - - msg = create_keypoints_message( - keypoints, edges=self.edges, label_names=self.label_names + keypoints = self.extract(output) + keypoints = self.compute( + keypoints, + n_keypoints=self.n_keypoints, + scale_factor=self.scale_factor, ) - 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 keypoints message with {len(keypoints)} points" + self.emit(output, keypoints) + + 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." ) - self.out.send(msg) + return output.getTensor(self.output_layer_name, dequantize=True).astype( + np.float32 + ) + + @staticmethod + def compute( + keypoints: np.ndarray, + *, + n_keypoints: int, + scale_factor: float = 1.0, + ) -> np.ndarray: + return compute_keypoints( + keypoints, + n_keypoints=n_keypoints, + scale_factor=scale_factor, + ) - self._logger.debug("Keypoint output sent successfully") + def emit(self, output: dai.NNData, keypoints: np.ndarray) -> None: + msg = create_keypoints_message( + keypoints, edges=self.edges, label_names=self.label_names + ) + 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 keypoints message with {len(keypoints)} points") + self.out.send(msg) + self._logger.debug("Keypoint output sent successfully") diff --git a/depthai_nodes/node/parsers/lane_detection.py b/depthai_nodes/node/parsers/lane_detection.py index 97803266..9084f849 100644 --- a/depthai_nodes/node/parsers/lane_detection.py +++ b/depthai_nodes/node/parsers/lane_detection.py @@ -5,7 +5,9 @@ from depthai_nodes.message.creators import create_cluster_message from depthai_nodes.node.parsers.base_parser import BaseParser -from depthai_nodes.node.parsers.utils.ufld import decode_ufld +from depthai_nodes.node.parsers.utils.lane_detection import ( + compute_lane_detection_points, +) class LaneDetectionParser(BaseParser): @@ -192,41 +194,56 @@ def run(self): except dai.MessageQueue.QueueException: break # Pipeline was stopped - 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." - ) - - tensor = output.getTensor(self.output_layer_name, dequantize=True).astype( - np.float32 - ) - y = tensor[0] - - points = decode_ufld( - anchors=self.row_anchors, + tensor = self.extract(output) + points = self.compute( + tensor, + row_anchors=self.row_anchors, griding_num=self.griding_num, cls_num_per_lane=self.cls_num_per_lane, - input_width=self.input_size[0], - input_height=self.input_size[1], - y=y, + input_size=self.input_size, ) - - msg = create_cluster_message(points) - 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 lane detection message with {len(points)} points" + self.emit(output, points) + + 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." ) - self.out.send(msg) + return output.getTensor(self.output_layer_name, dequantize=True).astype( + np.float32 + ) + + @staticmethod + def compute( + tensor: np.ndarray, + *, + row_anchors: list[int], + griding_num: int, + cls_num_per_lane: int, + input_size: tuple[int, int], + ) -> list[list[tuple[int, int]]]: + return compute_lane_detection_points( + tensor, + row_anchors=row_anchors, + griding_num=griding_num, + cls_num_per_lane=cls_num_per_lane, + input_size=input_size, + ) - self._logger.debug("Lane detection message sent successfully") + def emit(self, output: dai.NNData, points) -> None: + msg = create_cluster_message(points) + 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 lane detection message with {len(points)} points") + self.out.send(msg) + self._logger.debug("Lane detection message sent successfully") diff --git a/depthai_nodes/node/parsers/map_output.py b/depthai_nodes/node/parsers/map_output.py index 20501f2c..d86f0ed8 100644 --- a/depthai_nodes/node/parsers/map_output.py +++ b/depthai_nodes/node/parsers/map_output.py @@ -4,6 +4,7 @@ from depthai_nodes.message.creators import create_map_message from depthai_nodes.node.parsers.base_parser import BaseParser +from depthai_nodes.node.parsers.utils.map_output import compute_map_output class MapOutputParser(BaseParser): @@ -99,32 +100,37 @@ def run(self): except dai.MessageQueue.QueueException: break # Pipeline was stopped - 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." - ) + map_tensor = self.extract(output) + map_output = self.compute(map_tensor) + self.emit(output, map_output) - map = output.getTensor(self.output_layer_name, dequantize=True) - - if map.shape[0] == 1: - map = map[0] # remove batch dimension - - map_message = create_map_message( - map=map, min_max_scaling=self.min_max_scaling + def extract(self, output: dai.NNData): + 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." ) - map_message.setTimestamp(output.getTimestamp()) - map_message.setTimestampDevice(output.getTimestampDevice()) - map_message.setSequenceNum(output.getSequenceNum()) - transformation = output.getTransformation() - if transformation is not None: - map_message.setTransformation(transformation) - self._logger.debug("Created Map message.") + return output.getTensor(self.output_layer_name, dequantize=True) - self.out.send(map_message) + @staticmethod + def compute(map_tensor): + return compute_map_output(map_tensor) - self._logger.debug("Map message sent successfully") + def emit(self, output: dai.NNData, map_output) -> None: + map_message = create_map_message( + map=map_output, min_max_scaling=self.min_max_scaling + ) + map_message.setTimestamp(output.getTimestamp()) + map_message.setTimestampDevice(output.getTimestampDevice()) + map_message.setSequenceNum(output.getSequenceNum()) + transformation = output.getTransformation() + if transformation is not None: + map_message.setTransformation(transformation) + + self._logger.debug("Created Map message.") + self.out.send(map_message) + self._logger.debug("Map message sent successfully") diff --git a/depthai_nodes/node/parsers/mediapipe_palm_detection.py b/depthai_nodes/node/parsers/mediapipe_palm_detection.py index 6851d8a0..4b8b570d 100644 --- a/depthai_nodes/node/parsers/mediapipe_palm_detection.py +++ b/depthai_nodes/node/parsers/mediapipe_palm_detection.py @@ -7,6 +7,7 @@ from depthai_nodes.message.creators import create_detection_message from depthai_nodes.node.parsers.detection import DetectionParser from depthai_nodes.node.parsers.utils.medipipe import ( + compute_mediapipe_palm_detections, decode, generate_handtracker_anchors, ) @@ -142,97 +143,88 @@ def run(self): except dai.MessageQueue.QueueException: break # Pipeline was stopped - all_tensors = output.getAllLayerNames() - - self._logger.debug(f"Processing input with layers: {all_tensors}") - - bboxes = None - scores = None - - for tensor_name in all_tensors: - tensor = np.array( - output.getTensor(tensor_name, dequantize=True), dtype=np.float32 - ) - - if bboxes is None: - bboxes = tensor - scores = tensor - else: - bboxes = bboxes if tensor.shape[-1] < bboxes.shape[-1] else tensor - scores = tensor if tensor.shape[-1] < scores.shape[-1] else scores - - bboxes = bboxes.reshape(-1, 18) - scores = scores.reshape(-1) - - if bboxes is None or scores is None: - raise ValueError("No valid output tensors found.") - - decoded_bboxes = decode( - bboxes=bboxes, - scores=scores, - anchors=self._anchors, - threshold=self.conf_threshold, - scale=self.scale, - ) - - bboxes = [] - scores = [] - angles = [] - for hand in decoded_bboxes: - extended_points = np.array(hand.rect_points) - - x_dist = extended_points[3][0] - extended_points[0][0] - y_dist = extended_points[3][1] - extended_points[0][1] - - angle = np.degrees(np.arctan2(y_dist, x_dist)) - x_center, y_center = np.mean(extended_points, axis=0) - width = np.linalg.norm(extended_points[0] - extended_points[3]) - height = np.linalg.norm(extended_points[0] - extended_points[1]) - - bboxes.append([x_center, y_center, width, height]) - angles.append(angle) - scores.append(hand.pd_score) - - indices = cv2.dnn.NMSBoxes( + bboxes, scores = self.extract(output) + bboxes, scores, angles, labels, label_names = self.compute( bboxes, scores, - self.conf_threshold, - self.iou_threshold, - top_k=self.max_det, + anchors=self._anchors, + conf_threshold=self.conf_threshold, + iou_threshold=self.iou_threshold, + max_det=self.max_det, + scale=self.scale, + label_names=self.label_names, ) - bboxes = np.array(bboxes)[indices] - scores = np.array(scores)[indices] - angles = np.array(angles)[indices] - bboxes = bboxes.astype(float) / self.scale + self.emit(output, bboxes, scores, angles, labels, label_names) - bboxes = np.clip(bboxes, 0, 1) - angles = np.round(angles, 0) + def extract(self, output: dai.NNData) -> tuple[np.ndarray, np.ndarray]: + all_tensors = output.getAllLayerNames() + self._logger.debug(f"Processing input with layers: {all_tensors}") - labels = np.array([0] * len(bboxes)) + bboxes = None + scores = None - label_names = ( - [self.label_names[label] for label in labels] - if self.label_names - else None - ) - detections_msg = create_detection_message( - bboxes=bboxes, - scores=scores, - angles=angles, - labels=labels, - label_names=label_names, + for tensor_name in all_tensors: + tensor = np.array( + output.getTensor(tensor_name, dequantize=True), dtype=np.float32 ) - detections_msg.setTimestamp(output.getTimestamp()) - detections_msg.setSequenceNum(output.getSequenceNum()) - detections_msg.setTimestampDevice(output.getTimestampDevice()) - transformation = output.getTransformation() - if transformation is not None: - detections_msg.setTransformation(transformation) - - self._logger.debug( - f"Created detection message with {len(bboxes)} detections" - ) - - self.out.send(detections_msg) + if bboxes is None: + bboxes = tensor + scores = tensor + else: + bboxes = bboxes if tensor.shape[-1] < bboxes.shape[-1] else tensor + scores = tensor if tensor.shape[-1] < scores.shape[-1] else scores + + if bboxes is None or scores is None: + raise ValueError("No valid output tensors found.") + + return bboxes.reshape(-1, 18), scores.reshape(-1) + + @staticmethod + def compute( + bboxes: np.ndarray, + scores: np.ndarray, + *, + anchors: np.ndarray, + conf_threshold: float, + iou_threshold: float, + max_det: int, + scale: int, + label_names: list[str] | None = None, + ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, list[str] | None]: + return compute_mediapipe_palm_detections( + bboxes, + scores, + anchors=anchors, + conf_threshold=conf_threshold, + iou_threshold=iou_threshold, + max_det=max_det, + scale=scale, + label_names=label_names, + ) - self._logger.debug("Detection message sent successfully") + def emit( + self, + output: dai.NNData, + bboxes: np.ndarray, + scores: np.ndarray, + angles: np.ndarray, + labels: np.ndarray, + label_names: list[str] | None, + ) -> None: + detections_msg = create_detection_message( + bboxes=bboxes, + scores=scores, + angles=angles, + labels=labels, + label_names=label_names, + ) + detections_msg.setTimestamp(output.getTimestamp()) + detections_msg.setSequenceNum(output.getSequenceNum()) + detections_msg.setTimestampDevice(output.getTimestampDevice()) + transformation = output.getTransformation() + if transformation is not None: + detections_msg.setTransformation(transformation) + + self._logger.debug(f"Created detection message with {len(bboxes)} detections") + self.out.send(detections_msg) + self._logger.debug("Detection message sent successfully") diff --git a/depthai_nodes/node/parsers/mlsd.py b/depthai_nodes/node/parsers/mlsd.py index dee93c07..5851a501 100644 --- a/depthai_nodes/node/parsers/mlsd.py +++ b/depthai_nodes/node/parsers/mlsd.py @@ -5,7 +5,7 @@ from depthai_nodes.message.creators import create_line_detection_message from depthai_nodes.node.parsers.base_parser import BaseParser -from depthai_nodes.node.parsers.utils.mlsd import decode_scores_and_points, get_lines +from depthai_nodes.node.parsers.utils.mlsd import compute_mlsd_lines class MLSDParser(BaseParser): @@ -165,38 +165,54 @@ def run(self): except dai.MessageQueue.QueueException: break # Pipeline was stopped - self._logger.debug( - f"Processing input with layers: {output.getAllLayerNames()}" + tpMap, heat_np = self.extract(output) + lines, scores = self.compute( + tpMap, + heat_np, + topk_n=self.topk_n, + score_thr=self.score_thr, + dist_thr=self.dist_thr, ) - tpMap = output.getTensor( - self.output_layer_tpmap, - dequantize=True, - storageOrder=dai.TensorInfo.StorageOrder.NCHW, - ).astype(np.float32) - heat_np = output.getTensor(self.output_layer_heat, dequantize=True).astype( - np.float32 - ) - - if len(tpMap.shape) != 4: - raise ValueError("Invalid shape of the tpMap tensor. Should be 4D.") - - pts, pts_score, vmap = decode_scores_and_points(tpMap, heat_np, self.topk_n) - lines, scores = get_lines( - pts, pts_score, vmap, self.score_thr, self.dist_thr - ) - - message = create_line_detection_message(lines, np.array(scores)) - message.setTimestamp(output.getTimestamp()) - message.setSequenceNum(output.getSequenceNum()) - message.setTimestampDevice(output.getTimestampDevice()) - transformation = output.getTransformation() - if transformation is not None: - message.setTransformation(transformation) - - self._logger.debug( - f"Created line detection message with {len(lines)} lines" - ) - - self.out.send(message) + self.emit(output, lines, scores) + + def extract(self, output: dai.NNData) -> tuple[np.ndarray, np.ndarray]: + self._logger.debug(f"Processing input with layers: {output.getAllLayerNames()}") + tpMap = output.getTensor( + self.output_layer_tpmap, + dequantize=True, + storageOrder=dai.TensorInfo.StorageOrder.NCHW, + ).astype(np.float32) + heat_np = output.getTensor(self.output_layer_heat, dequantize=True).astype( + np.float32 + ) + return tpMap, heat_np + + @staticmethod + def compute( + tpMap: np.ndarray, + heat_np: np.ndarray, + *, + topk_n: int, + score_thr: float, + dist_thr: float, + ) -> tuple[np.ndarray, np.ndarray]: + return compute_mlsd_lines( + tpMap, + heat_np, + topk_n=topk_n, + score_thr=score_thr, + dist_thr=dist_thr, + ) - self._logger.debug("Line detection message sent successfully") + def emit(self, output: dai.NNData, lines: np.ndarray, scores: np.ndarray) -> None: + message = create_line_detection_message(lines, scores) + message.setTimestamp(output.getTimestamp()) + message.setSequenceNum(output.getSequenceNum()) + message.setTimestampDevice(output.getTimestampDevice()) + transformation = output.getTransformation() + if transformation is not None: + message.setTransformation(transformation) + + self._logger.debug(f"Created line detection message with {len(lines)} lines") + self.out.send(message) + self._logger.debug("Line detection message sent successfully") diff --git a/depthai_nodes/node/parsers/ppdet.py b/depthai_nodes/node/parsers/ppdet.py index 604be8a2..2da5c656 100644 --- a/depthai_nodes/node/parsers/ppdet.py +++ b/depthai_nodes/node/parsers/ppdet.py @@ -5,7 +5,7 @@ from depthai_nodes.message.creators import create_detection_message from depthai_nodes.node.parsers.detection import DetectionParser -from depthai_nodes.node.parsers.utils.ppdet import parse_paddle_detection_outputs +from depthai_nodes.node.parsers.utils.ppdet import compute_pp_text_detections class PPTextDetectionParser(DetectionParser): @@ -106,47 +106,65 @@ def run(self): except dai.MessageQueue.QueueException: break # Pipeline was stopped - 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." - ) - - predictions = np.array( - output.getTensor( - self.output_layer_name, - dequantize=True, - storageOrder=dai.TensorInfo.StorageOrder.NCHW, - ) - ) - - _, _, height, width = predictions.shape - - bboxes, angles, scores = parse_paddle_detection_outputs( + predictions = self.extract(output) + bboxes, angles, scores = self.compute( predictions, - self.mask_threshold, - self.conf_threshold, - self.max_det, - width=width, - height=height, + mask_threshold=self.mask_threshold, + conf_threshold=self.conf_threshold, + max_det=self.max_det, ) - message = create_detection_message( - bboxes=bboxes, scores=scores, angles=angles + self.emit(output, bboxes, angles, 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." ) - message.setTimestamp(output.getTimestamp()) - message.setSequenceNum(output.getSequenceNum()) - message.setTimestampDevice(output.getTimestampDevice()) - transformation = output.getTransformation() - if transformation is not None: - message.setTransformation(transformation) - - self._logger.debug( - f"Created text detection message with {len(bboxes)} detections" + + return np.array( + output.getTensor( + self.output_layer_name, + dequantize=True, + storageOrder=dai.TensorInfo.StorageOrder.NCHW, ) + ) - self.out.send(message) + @staticmethod + def compute( + predictions: np.ndarray, + *, + mask_threshold: float, + conf_threshold: float, + max_det: int, + ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + return compute_pp_text_detections( + predictions, + mask_threshold=mask_threshold, + conf_threshold=conf_threshold, + max_det=max_det, + ) + + def emit( + self, + output: dai.NNData, + bboxes: np.ndarray, + angles: np.ndarray, + scores: np.ndarray, + ) -> None: + message = create_detection_message(bboxes=bboxes, scores=scores, angles=angles) + message.setTimestamp(output.getTimestamp()) + message.setSequenceNum(output.getSequenceNum()) + message.setTimestampDevice(output.getTimestampDevice()) + transformation = output.getTransformation() + if transformation is not None: + message.setTransformation(transformation) - self._logger.debug("Text detection message sent successfully") + self._logger.debug( + f"Created text detection message with {len(bboxes)} detections" + ) + self.out.send(message) + self._logger.debug("Text detection message sent successfully") diff --git a/depthai_nodes/node/parsers/regression.py b/depthai_nodes/node/parsers/regression.py index 96852142..9403fe99 100644 --- a/depthai_nodes/node/parsers/regression.py +++ b/depthai_nodes/node/parsers/regression.py @@ -5,6 +5,9 @@ from depthai_nodes.message.creators import create_regression_message from depthai_nodes.node.parsers.base_parser import BaseParser +from depthai_nodes.node.parsers.utils.regression import ( + compute_regression_predictions, +) class RegressionParser(BaseParser): @@ -82,32 +85,35 @@ def run(self): except dai.MessageQueue.QueueException: break # Pipeline was stopped - 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." - ) - - predictions = output.getTensor( - self.output_layer_name, dequantize=True - ).squeeze() - predictions = np.atleast_1d(predictions).tolist() - - regression_message = create_regression_message(predictions=predictions) - regression_message.setTimestamp(output.getTimestamp()) - regression_message.setTimestampDevice(output.getTimestampDevice()) - regression_message.setSequenceNum(output.getSequenceNum()) - transformation = output.getTransformation() - if transformation is not None: - regression_message.setTransformation(transformation) - - self._logger.debug( - f"Created regression message with {len(predictions)} values" + predictions = self.extract(output) + predictions = self.compute(predictions) + self.emit(output, predictions) + + 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." ) - self.out.send(regression_message) + return output.getTensor(self.output_layer_name, dequantize=True) + + @staticmethod + def compute(predictions: np.ndarray) -> list[float]: + return compute_regression_predictions(predictions) + + def emit(self, output: dai.NNData, predictions: list[float]) -> None: + regression_message = create_regression_message(predictions=predictions) + regression_message.setTimestamp(output.getTimestamp()) + regression_message.setTimestampDevice(output.getTimestampDevice()) + regression_message.setSequenceNum(output.getSequenceNum()) + transformation = output.getTransformation() + if transformation is not None: + regression_message.setTransformation(transformation) - self._logger.debug("Regression message sent successfully") + self._logger.debug(f"Created regression message with {len(predictions)} values") + self.out.send(regression_message) + self._logger.debug("Regression message sent successfully") diff --git a/depthai_nodes/node/parsers/rf_detr.py b/depthai_nodes/node/parsers/rf_detr.py index a0b6fa79..7dbe533b 100644 --- a/depthai_nodes/node/parsers/rf_detr.py +++ b/depthai_nodes/node/parsers/rf_detr.py @@ -5,10 +5,8 @@ from depthai_nodes.message.creators import create_detection_message from depthai_nodes.node.parsers.base_parser import BaseParser -from depthai_nodes.node.parsers.utils import sigmoid, xyxy_to_xywh -from depthai_nodes.node.parsers.utils.bbox_format_converters import xywh_to_xyxy -from depthai_nodes.node.parsers.utils.masks_utils import ( - process_single_mask_rfdetr, +from depthai_nodes.node.parsers.utils.rf_detr import ( + compute_rfdetr_detections, ) @@ -201,132 +199,91 @@ def run(self): except dai.MessageQueue.QueueException: break # Pipeline was stopped - # Get output layers - layer_names = self.output_layer_names or output.getAllLayerNames() - self._logger.debug(f"Processing input with layers: {layer_names}") + boxes_tensor, logits_tensor, masks_tensor = self.extract(output) + boxes, scores, labels, label_names_list, final_mask = self.compute( + boxes_tensor, + logits_tensor, + conf_threshold=self.conf_threshold, + max_det=self.max_det, + label_names=self.label_names, + mask_conf=self.mask_conf, + input_shape=self.input_shape, + masks_tensor=masks_tensor, + ) + self.emit(output, boxes, scores, labels, label_names_list, final_mask) - if len(layer_names) < 2 or len(layer_names) > 3: - raise ValueError( - "Expected 2 or 3 output layers " - f"(boxes, logits, optional masks), got {len(layer_names)} layers." - ) + def extract( + self, output: dai.NNData + ) -> tuple[np.ndarray, np.ndarray, np.ndarray | None]: + layer_names = self.output_layer_names or output.getAllLayerNames() + self._logger.debug(f"Processing input with layers: {layer_names}") - # outputs[0]: boxes in cxcywh format (normalized coordinates [0, 1]) - # outputs[1]: class logits (need sigmoid activation) - # outputs[2]: instance segmentation masks (optional) - boxes_tensor = output.getTensor(layer_names[0], dequantize=True).astype( - np.float32 + if len(layer_names) < 2 or len(layer_names) > 3: + raise ValueError( + "Expected 2 or 3 output layers " + f"(boxes, logits, optional masks), got {len(layer_names)} layers." ) - logits_tensor = output.getTensor(layer_names[1], dequantize=True).astype( + + boxes_tensor = output.getTensor(layer_names[0], dequantize=True).astype( + np.float32 + ) + logits_tensor = output.getTensor(layer_names[1], dequantize=True).astype( + np.float32 + ) + masks_tensor = None + if len(layer_names) == 3: + masks_tensor = output.getTensor(layer_names[2], dequantize=True).astype( np.float32 ) + return boxes_tensor, logits_tensor, masks_tensor + + @staticmethod + def compute( + boxes_tensor: np.ndarray, + logits_tensor: np.ndarray, + *, + conf_threshold: float, + max_det: int, + label_names: list[str] | None, + mask_conf: float, + input_shape: tuple[int, int] | None, + masks_tensor: np.ndarray | None = None, + ) -> tuple[np.ndarray, np.ndarray, np.ndarray, list[str] | None, np.ndarray | None]: + return compute_rfdetr_detections( + boxes_tensor, + logits_tensor, + conf_threshold=conf_threshold, + max_det=max_det, + label_names=label_names, + mask_conf=mask_conf, + input_shape=input_shape, + masks_tensor=masks_tensor, + ) - masks_tensor = None - if len(layer_names) == 3: - masks_tensor = output.getTensor(layer_names[2], dequantize=True).astype( - np.float32 - ) - - prob = sigmoid(logits_tensor) - - scores = np.max(prob, axis=2).squeeze() # (num_queries,) - labels = np.argmax(prob, axis=2).squeeze() # (num_queries,) - - sorted_idx = np.argsort(scores)[::-1] - - effective_max_det = self.max_det - if masks_tensor is not None: - # SegmentationMask is uint8 and reserves 255 for background, - # so valid instance IDs are 0..254. - max_segmentation_instances = 255 - - num_valid_instances = int( - np.count_nonzero( - scores[sorted_idx][: self.max_det] > self.conf_threshold - ) - ) - if num_valid_instances > max_segmentation_instances: - self._logger.warning( - "RFDETRParser can encode at most 255 instances in " - "SegmentationMask; ignoring " - f"{num_valid_instances - max_segmentation_instances} " - "lowest-scoring instances." - ) - - effective_max_det = min(effective_max_det, max_segmentation_instances) - - scores = scores[sorted_idx][:effective_max_det] - labels = labels[sorted_idx][:effective_max_det] - boxes_cxcywh = boxes_tensor.squeeze()[sorted_idx][:effective_max_det] - - masks = None - if masks_tensor is not None: - masks = masks_tensor.squeeze()[sorted_idx][:effective_max_det] - - # Convert boxes from cxcywh (normalized) to xyxy (normalized) - boxes = np.clip(xywh_to_xyxy(boxes_cxcywh), 0, 1) - - # Filter detections by confidence threshold - confidence_mask = scores > self.conf_threshold - scores = scores[confidence_mask] - labels = labels[confidence_mask] - boxes = boxes[confidence_mask] - boxes_cxcywh = boxes_cxcywh[confidence_mask] - - if masks is not None: - masks = masks[confidence_mask] - - final_mask = None - mode = self._SEG_MODE if masks is not None else self._DET_MODE - - if mode == self._SEG_MODE: - if self.input_shape is None: - raise ValueError( - "RFDETRParser segmentation mode requires model input shape." - ) - - final_mask = np.full(self.input_shape, 255, dtype=np.uint8) - - for i, (mask_logits, bbox) in enumerate(zip(masks, boxes_cxcywh)): - resized_mask = process_single_mask_rfdetr( - mask_logits=mask_logits, - mask_conf=self.mask_conf, - bbox=bbox, - input_shape=self.input_shape, - ) - foreground = resized_mask > 0 - final_mask[(final_mask == 255) & foreground] = i - - boxes = xyxy_to_xywh(boxes) - - label_names_list = None - if self.label_names: - label_names_list = [ - ( - self.label_names[int(label)] - if int(label) < len(self.label_names) - else f"class_{int(label)}" - ) - for label in labels - ] - - # Create detection message - message = create_detection_message( - bboxes=boxes, - scores=scores, - labels=labels.astype(int), - label_names=label_names_list, - masks=final_mask, - ) + def emit( + self, + output: dai.NNData, + boxes: np.ndarray, + scores: np.ndarray, + labels: np.ndarray, + label_names_list: list[str] | None, + final_mask: np.ndarray | None, + ) -> None: + message = create_detection_message( + bboxes=boxes, + scores=scores, + labels=labels.astype(int), + label_names=label_names_list, + masks=final_mask, + ) + + transformation = output.getTransformation() + if transformation is not None: + message.setTransformation(transformation) + message.setTimestamp(output.getTimestamp()) + message.setSequenceNum(output.getSequenceNum()) + message.setTimestampDevice(output.getTimestampDevice()) - # Set message metadata - 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(boxes)} objects") - self.out.send(message) - self._logger.debug("Detections message sent successfully") + self._logger.debug(f"Created detections message with {len(boxes)} objects") + self.out.send(message) + self._logger.debug("Detections message sent successfully") diff --git a/depthai_nodes/node/parsers/scrfd.py b/depthai_nodes/node/parsers/scrfd.py index 95b09d96..8df7a8fa 100644 --- a/depthai_nodes/node/parsers/scrfd.py +++ b/depthai_nodes/node/parsers/scrfd.py @@ -6,7 +6,10 @@ from depthai_nodes.message.creators import create_detection_message from depthai_nodes.node.parsers.detection import DetectionParser from depthai_nodes.node.parsers.utils import xyxy_to_xywh -from depthai_nodes.node.parsers.utils.scrfd import compute_anchor_centers, decode_scrfd +from depthai_nodes.node.parsers.utils.scrfd import ( + compute_anchor_centers, + compute_scrfd_detections, +) class SCRFDParser(DetectionParser): @@ -176,55 +179,12 @@ def run(self): except dai.MessageQueue.QueueException: break # Pipeline was stopped - scores_concatenated = [] - bboxes_concatenated = [] - kps_concatenated = [] - - if len(self.output_layer_names) == 0: - self.output_layer_names = output.getAllLayerNames() - - self._logger.debug( - f"Processing input with layers: {output.getAllLayerNames()}" - ) - - for stride in self.feat_stride_fpn: - score_layer_name = f"score_{stride}" - bbox_layer_name = f"bbox_{stride}" - kps_layer_name = f"kps_{stride}" - if score_layer_name not in self.output_layer_names: - raise ValueError( - f"Layer {score_layer_name} not found in the model output." - ) - if bbox_layer_name not in self.output_layer_names: - raise ValueError( - f"Layer {bbox_layer_name} not found in the model output." - ) - if kps_layer_name not in self.output_layer_names: - raise ValueError( - f"Layer {kps_layer_name} not found in the model output." - ) - - score_tensor = ( - output.getTensor(score_layer_name, dequantize=True) - .flatten() - .astype(np.float32) - ) - bbox_tensor = ( - output.getTensor(bbox_layer_name, dequantize=True) - .reshape(len(score_tensor), 4) - .astype(np.float32) - ) - kps_tensor = ( - output.getTensor(kps_layer_name, dequantize=True) - .reshape(len(score_tensor), 10) - .astype(np.float32) - ) - - scores_concatenated.append(score_tensor) - bboxes_concatenated.append(bbox_tensor) - kps_concatenated.append(kps_tensor) - - bboxes, scores, keypoints = decode_scrfd( + ( + bboxes_concatenated, + scores_concatenated, + kps_concatenated, + ) = self.extract(output) + bboxes, scores, keypoints, labels, label_names = self.compute( bboxes_concatenated=bboxes_concatenated, scores_concatenated=scores_concatenated, kps_concatenated=kps_concatenated, @@ -234,35 +194,91 @@ def run(self): score_threshold=self.conf_threshold, nms_threshold=self.iou_threshold, anchors=self._cached_anchors, + label_names=self.label_names, ) - bboxes = xyxy_to_xywh(bboxes) - bboxes = np.clip(bboxes, 0, 1) - - labels = np.array([0] * len(bboxes)) + self.emit(output, bboxes, scores, keypoints, labels, label_names) + + def extract( + self, output: dai.NNData + ) -> tuple[list[np.ndarray], list[np.ndarray], list[np.ndarray]]: + scores_concatenated = [] + bboxes_concatenated = [] + kps_concatenated = [] + + if len(self.output_layer_names) == 0: + self.output_layer_names = output.getAllLayerNames() + + self._logger.debug(f"Processing input with layers: {output.getAllLayerNames()}") + + for stride in self.feat_stride_fpn: + score_layer_name = f"score_{stride}" + bbox_layer_name = f"bbox_{stride}" + kps_layer_name = f"kps_{stride}" + if score_layer_name not in self.output_layer_names: + raise ValueError( + f"Layer {score_layer_name} not found in the model output." + ) + if bbox_layer_name not in self.output_layer_names: + raise ValueError( + f"Layer {bbox_layer_name} not found in the model output." + ) + if kps_layer_name not in self.output_layer_names: + raise ValueError( + f"Layer {kps_layer_name} not found in the model output." + ) - label_names = ( - [self.label_names[label] for label in labels] - if self.label_names - else None + score_tensor = ( + output.getTensor(score_layer_name, dequantize=True) + .flatten() + .astype(np.float32) ) - message = create_detection_message( - bboxes=bboxes, - scores=scores, - labels=labels, - label_names=label_names, - keypoints=keypoints, + bbox_tensor = ( + output.getTensor(bbox_layer_name, dequantize=True) + .reshape(len(score_tensor), 4) + .astype(np.float32) ) - message.setTimestamp(output.getTimestamp()) - message.setSequenceNum(output.getSequenceNum()) - message.setTimestampDevice(output.getTimestampDevice()) - transformation = output.getTransformation() - if transformation is not None: - message.setTransformation(transformation) - - self._logger.debug( - f"Created detection message with {len(bboxes)} detections" + kps_tensor = ( + output.getTensor(kps_layer_name, dequantize=True) + .reshape(len(score_tensor), 10) + .astype(np.float32) ) - self.out.send(message) + scores_concatenated.append(score_tensor) + bboxes_concatenated.append(bbox_tensor) + kps_concatenated.append(kps_tensor) + + return bboxes_concatenated, scores_concatenated, kps_concatenated - self._logger.debug("Detection message sent successfully") + @staticmethod + def compute(**kwargs): + return compute_scrfd_detections(**kwargs) + + def emit( + self, + output: dai.NNData, + bboxes: np.ndarray, + scores: np.ndarray, + keypoints: np.ndarray, + labels: np.ndarray, + label_names: list[str] | None, + ) -> None: + bboxes = xyxy_to_xywh(bboxes) + bboxes = np.clip(bboxes, 0, 1) + + message = create_detection_message( + bboxes=bboxes, + scores=scores, + labels=labels, + label_names=label_names, + keypoints=keypoints, + ) + message.setTimestamp(output.getTimestamp()) + message.setSequenceNum(output.getSequenceNum()) + message.setTimestampDevice(output.getTimestampDevice()) + transformation = output.getTransformation() + if transformation is not None: + message.setTransformation(transformation) + + self._logger.debug(f"Created detection message with {len(bboxes)} detections") + self.out.send(message) + self._logger.debug("Detection message sent successfully") diff --git a/depthai_nodes/node/parsers/segmentation.py b/depthai_nodes/node/parsers/segmentation.py index 33239bd4..4b4080e9 100644 --- a/depthai_nodes/node/parsers/segmentation.py +++ b/depthai_nodes/node/parsers/segmentation.py @@ -5,6 +5,9 @@ from depthai_nodes.message.creators import create_segmentation_message from depthai_nodes.node.parsers.base_parser import BaseParser +from depthai_nodes.node.parsers.utils.segmentation import ( + compute_segmentation_class_map, +) class SegmentationParser(BaseParser): @@ -107,73 +110,47 @@ def run(self): except dai.MessageQueue.QueueException: break # Pipeline was stopped - 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." - ) - - segmentation_mask = output.getTensor( - self.output_layer_name, dequantize=True + segmentation_mask = self.extract(output) + class_map = self.compute( + segmentation_mask, + classes_in_one_layer=self.classes_in_one_layer, ) - if len(segmentation_mask.shape) == 4: - segmentation_mask = segmentation_mask[0] - - if len(segmentation_mask.shape) != 3: - raise ValueError( - f"Expected 3D output tensor, got {len(segmentation_mask.shape)}D." - ) - - np_function = np.argmax - mask_shape = segmentation_mask.shape - min_dim = np.argmin(mask_shape) - if min_dim == len(mask_shape) - 1: - segmentation_mask = segmentation_mask.transpose(2, 0, 1) - adding_unassigned_class = False - if segmentation_mask.shape[0] == 1: # shape is (1, H, W) - if self.classes_in_one_layer: - np_function = np.max - else: - # If there is only one class, add an unassigned class - adding_unassigned_class = True - segmentation_mask = np.vstack( - ( - np.zeros( - ( - 1, - segmentation_mask.shape[1], - segmentation_mask.shape[2], - ), - dtype=np.float32, - ), - segmentation_mask, - ) - ) - - class_map = ( - np_function(segmentation_mask, axis=0) - .reshape(segmentation_mask.shape[1], segmentation_mask.shape[2]) - .astype(np.int16) + self.emit(output, class_map) + + def extract(self, output: dai.NNData): + 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 adding_unassigned_class: - class_map = class_map - 1 - - mask_message = create_segmentation_message(class_map) - mask_message.setTimestamp(output.getTimestamp()) - mask_message.setSequenceNum(output.getSequenceNum()) - mask_message.setTimestampDevice(output.getTimestampDevice()) - transformation = output.getTransformation() - if transformation is not None: - mask_message.setTransformation(transformation) - - self._logger.debug( - f"Created segmentation message with {class_map.shape[0]} classes" - ) + return output.getTensor(self.output_layer_name, dequantize=True) + + @staticmethod + def compute( + segmentation_mask: np.ndarray, + *, + classes_in_one_layer: bool = False, + ) -> np.ndarray: + return compute_segmentation_class_map( + segmentation_mask, + classes_in_one_layer=classes_in_one_layer, + ) - self.out.send(mask_message) + def emit(self, output: dai.NNData, class_map: np.ndarray) -> None: + mask_message = create_segmentation_message(class_map) + mask_message.setTimestamp(output.getTimestamp()) + mask_message.setSequenceNum(output.getSequenceNum()) + mask_message.setTimestampDevice(output.getTimestampDevice()) + transformation = output.getTransformation() + if transformation is not None: + mask_message.setTransformation(transformation) - self._logger.debug("Segmentation message sent successfully") + self._logger.debug( + f"Created segmentation message with {class_map.shape[0]} classes" + ) + self.out.send(mask_message) + self._logger.debug("Segmentation message sent successfully") diff --git a/depthai_nodes/node/parsers/superanimal_landmarker.py b/depthai_nodes/node/parsers/superanimal_landmarker.py index c924426a..ae52fad8 100644 --- a/depthai_nodes/node/parsers/superanimal_landmarker.py +++ b/depthai_nodes/node/parsers/superanimal_landmarker.py @@ -5,7 +5,9 @@ from depthai_nodes.message.creators import create_keypoints_message from depthai_nodes.node.parsers.keypoints import KeypointParser -from depthai_nodes.node.parsers.utils.superanimal import get_pose_prediction +from depthai_nodes.node.parsers.utils.superanimal import ( + compute_superanimal_keypoints, +) class SuperAnimalParser(KeypointParser): @@ -97,45 +99,52 @@ def run(self): except dai.MessageQueue.QueueException: break # Pipeline was stopped - 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." - ) - heatmaps = output.getTensor(self.output_layer_name, dequantize=True).astype( - np.float32 + heatmaps = self.extract(output) + keypoints, scores = self.compute(heatmaps, scale_factor=self.scale_factor) + self.emit(output, keypoints, 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." ) - heatmaps_scale_factor = ( - self.scale_factor / heatmaps.shape[1], - self.scale_factor / heatmaps.shape[2], - ) - - keypoints = get_pose_prediction(heatmaps, None, heatmaps_scale_factor)[0][0] - scores = keypoints[:, 2] - keypoints = keypoints[:, :2] / self.scale_factor - - msg = create_keypoints_message( - keypoints, - scores, - self.score_threshold, - label_names=self.label_names, - edges=self.edges, - ) - 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 keypoints message with {len(keypoints)} points" - ) + return output.getTensor(self.output_layer_name, dequantize=True).astype( + np.float32 + ) - self.out.send(msg) + @staticmethod + def compute( + heatmaps: np.ndarray, + *, + scale_factor: float, + ) -> tuple[np.ndarray, np.ndarray]: + return compute_superanimal_keypoints( + heatmaps, + scale_factor=scale_factor, + ) - self._logger.debug("Keypoint output sent successfully") + def emit( + self, output: dai.NNData, keypoints: np.ndarray, scores: np.ndarray + ) -> None: + msg = create_keypoints_message( + keypoints, + scores, + self.score_threshold, + label_names=self.label_names, + edges=self.edges, + ) + 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 keypoints message with {len(keypoints)} points") + self.out.send(msg) + self._logger.debug("Keypoint output sent successfully") diff --git a/depthai_nodes/node/parsers/utils/classification_sequence.py b/depthai_nodes/node/parsers/utils/classification_sequence.py new file mode 100644 index 00000000..2fce37b6 --- /dev/null +++ b/depthai_nodes/node/parsers/utils/classification_sequence.py @@ -0,0 +1,34 @@ +import numpy as np + +from .activations import softmax + + +def compute_classification_sequence_scores( + scores: np.ndarray, + *, + is_softmax: bool = True, +) -> np.ndarray: + """Return per-step classification scores, applying softmax when needed.""" + computed_scores = np.asarray(scores, dtype=np.float32) + + if computed_scores.ndim not in (2, 3): + raise ValueError( + f"Scores should be a 3D or 2D array, got shape {computed_scores.shape}." + ) + + if computed_scores.ndim == 3: + if computed_scores.shape[0] == 1: + computed_scores = computed_scores[0] + elif computed_scores.shape[2] == 1: + computed_scores = computed_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 is_softmax: + computed_scores = softmax(computed_scores, axis=1, keep_dims=True) + + return computed_scores diff --git a/depthai_nodes/node/parsers/utils/detection.py b/depthai_nodes/node/parsers/utils/detection.py new file mode 100644 index 00000000..b019132d --- /dev/null +++ b/depthai_nodes/node/parsers/utils/detection.py @@ -0,0 +1,23 @@ +import numpy as np + +from .bbox_format_converters import xyxy_to_xywh +from .nms import nms_cv2 + + +def compute_detection_outputs( + bboxes: np.ndarray, + scores: np.ndarray, + *, + conf_threshold: float, + iou_threshold: float, + 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) + + if len(indices) == 0: + return np.array([]), np.array([]) + + filtered_bboxes = xyxy_to_xywh(bboxes[indices]) + filtered_scores = scores[indices] + return filtered_bboxes, filtered_scores diff --git a/depthai_nodes/node/parsers/utils/embeddings.py b/depthai_nodes/node/parsers/utils/embeddings.py new file mode 100644 index 00000000..074c5f03 --- /dev/null +++ b/depthai_nodes/node/parsers/utils/embeddings.py @@ -0,0 +1,3 @@ +def compute_embeddings_output(output): + """Return the embeddings payload without modification.""" + return output diff --git a/depthai_nodes/node/parsers/utils/fastsam.py b/depthai_nodes/node/parsers/utils/fastsam.py index e27eed8a..20162920 100644 --- a/depthai_nodes/node/parsers/utils/fastsam.py +++ b/depthai_nodes/node/parsers/utils/fastsam.py @@ -391,3 +391,76 @@ def merge_masks(masks: np.ndarray) -> np.ndarray: merged_masks[masks[i] > 0] = i return merged_masks + + +def compute_fastsam_mask( + outputs_values: list[np.ndarray], + masks_outputs_values: list[np.ndarray], + protos_output: np.ndarray, + protos_len: int, + *, + conf_threshold: float, + n_classes: int, + iou_threshold: float, + mask_conf: float, + prompt: str, + 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.""" + width = outputs_values[0].shape[3] * 8 + height = outputs_values[0].shape[2] * 8 + input_shape = (width, height) + + results = decode_fastsam_output( + outputs_values, + [8, 16, 32], + [None, None, None], + img_shape=input_shape[::-1], + conf_thres=conf_threshold, + iou_thres=iou_threshold, + num_classes=n_classes, + ) + + if results.shape[0] == 0: + results_masks = np.array([]) + else: + results_bboxes = np.concatenate( + [ + results[:, :4].astype(int), + results[:, 4:5], + results[:, 5:6].astype(int), + ], + axis=1, + ) + mask_coeffs = build_mask_coeffs( + parsed_results=results, + masks_outputs_values=masks_outputs_values, + protos_len=protos_len, + ) + results_masks = process_masks( + parsed_results=results, + mask_coeffs=mask_coeffs, + protos=protos_output[0], + orig_shape=input_shape, + mask_conf=mask_conf, + ) + + if prompt == "bbox": + results_masks = box_prompt( + results_masks, bbox=bbox, orig_shape=input_shape[::-1] + ) + elif prompt == "point": + results_masks = point_prompt( + results_bboxes, + results_masks, + points=points, + pointlabel=point_label, + orig_shape=input_shape[::-1], + ) + + if len(results_masks) == 0: + results_masks = np.full((1, height, width), -1, dtype=np.int16) + + return merge_masks(results_masks) diff --git a/depthai_nodes/node/parsers/utils/hrnet.py b/depthai_nodes/node/parsers/utils/hrnet.py new file mode 100644 index 00000000..1ab6d24e --- /dev/null +++ b/depthai_nodes/node/parsers/utils/hrnet.py @@ -0,0 +1,27 @@ +import numpy as np + + +def compute_hrnet_keypoints( + heatmaps: np.ndarray, +) -> tuple[np.ndarray, np.ndarray]: + """Extract normalized keypoints and scores from HRNet heatmaps.""" + maps = np.asarray(heatmaps) + + if maps.shape[0] == 1: + maps = maps[0] + + if maps.ndim != 3: + raise ValueError(f"Expected 3D output tensor, got {maps.ndim}D.") + + _, map_h, map_w = maps.shape + + scores = np.array([np.max(heatmap) for heatmap in maps], dtype=np.float32) + scores = np.clip(scores, 0, 1) + + keypoints = np.array( + [np.unravel_index(heatmap.argmax(), heatmap.shape) for heatmap in maps], + dtype=np.float32, + ) + keypoints = keypoints[:, ::-1] / np.array([map_w, map_h], dtype=np.float32) + + return keypoints, scores diff --git a/depthai_nodes/node/parsers/utils/image_output.py b/depthai_nodes/node/parsers/utils/image_output.py new file mode 100644 index 00000000..bd944ede --- /dev/null +++ b/depthai_nodes/node/parsers/utils/image_output.py @@ -0,0 +1,16 @@ +import numpy as np + +from .denormalize import unnormalize_image + + +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: + image = image[0] + + if image.ndim != 3: + raise ValueError(f"Expected 3D output tensor, got {image.ndim}D.") + + return unnormalize_image(image) diff --git a/depthai_nodes/node/parsers/utils/keypoints.py b/depthai_nodes/node/parsers/utils/keypoints.py index dc8aaae8..731cc68d 100644 --- a/depthai_nodes/node/parsers/utils/keypoints.py +++ b/depthai_nodes/node/parsers/utils/keypoints.py @@ -37,3 +37,23 @@ def normalize_keypoints(keypoints: np.ndarray, height: int, width: int) -> np.nd keypoints[:, 1] = keypoints[:, 1] / height return keypoints + + +def compute_keypoints( + keypoints: np.ndarray, + *, + n_keypoints: int, + scale_factor: float = 1.0, +) -> np.ndarray: + """Reshape and normalize a keypoint tensor.""" + parsed_keypoints = np.asarray(keypoints, dtype=np.float32) + num_coords = int(np.prod(parsed_keypoints.shape) / n_keypoints) + + if num_coords not in [2, 3]: + raise ValueError( + f"Expected 2 or 3 coordinates per keypoint, got {num_coords}." + ) + + parsed_keypoints = parsed_keypoints.reshape(n_keypoints, num_coords) + parsed_keypoints /= scale_factor + return np.clip(parsed_keypoints, 0, 1) diff --git a/depthai_nodes/node/parsers/utils/lane_detection.py b/depthai_nodes/node/parsers/utils/lane_detection.py new file mode 100644 index 00000000..ae9533b7 --- /dev/null +++ b/depthai_nodes/node/parsers/utils/lane_detection.py @@ -0,0 +1,22 @@ +import numpy as np + +from .ufld import decode_ufld + + +def compute_lane_detection_points( + tensor: np.ndarray, + *, + row_anchors: list[int], + griding_num: int, + cls_num_per_lane: int, + input_size: tuple[int, int], +) -> list[list[tuple[int, int]]]: + """Decode lane points from the UFLD output tensor.""" + return decode_ufld( + anchors=row_anchors, + griding_num=griding_num, + cls_num_per_lane=cls_num_per_lane, + input_width=input_size[0], + input_height=input_size[1], + y=tensor[0], + ) diff --git a/depthai_nodes/node/parsers/utils/map_output.py b/depthai_nodes/node/parsers/utils/map_output.py new file mode 100644 index 00000000..e04ee141 --- /dev/null +++ b/depthai_nodes/node/parsers/utils/map_output.py @@ -0,0 +1,9 @@ +import numpy as np + + +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: + map_output = map_output[0] + return map_output diff --git a/depthai_nodes/node/parsers/utils/medipipe.py b/depthai_nodes/node/parsers/utils/medipipe.py index 0e951bf7..38cfd268 100644 --- a/depthai_nodes/node/parsers/utils/medipipe.py +++ b/depthai_nodes/node/parsers/utils/medipipe.py @@ -15,6 +15,7 @@ import math from collections import namedtuple +import cv2 import numpy as np @@ -384,3 +385,85 @@ def decode(bboxes, scores, anchors, threshold=0.5, scale=192): detections_to_rect(decoded_bboxes) rect_transformation(decoded_bboxes, scale, scale, no_shift=True) return decoded_bboxes + + +def compute_mediapipe_palm_detections( + bboxes: np.ndarray, + scores: np.ndarray, + *, + anchors: np.ndarray, + conf_threshold: float, + iou_threshold: float, + max_det: int, + scale: int, + label_names: list[str] | None = None, +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, list[str] | None]: + """Decode Mediapipe palm detections and apply NMS.""" + decoded_bboxes = decode( + bboxes=bboxes, + scores=scores, + anchors=anchors, + threshold=conf_threshold, + scale=scale, + ) + + bbox_list = [] + score_list = [] + angle_list = [] + for hand in decoded_bboxes: + extended_points = np.array(hand.rect_points) + + x_dist = extended_points[3][0] - extended_points[0][0] + y_dist = extended_points[3][1] - extended_points[0][1] + + angle = np.degrees(np.arctan2(y_dist, x_dist)) + x_center, y_center = np.mean(extended_points, axis=0) + width = np.linalg.norm(extended_points[0] - extended_points[3]) + height = np.linalg.norm(extended_points[0] - extended_points[1]) + + bbox_list.append([x_center, y_center, width, height]) + angle_list.append(angle) + score_list.append(hand.pd_score) + + if len(bbox_list) == 0: + return ( + np.array([]), + np.array([]), + np.array([]), + np.array([], dtype=int), + [] if label_names is not None else None, + ) + + indices = cv2.dnn.NMSBoxes( + bbox_list, + score_list, + conf_threshold, + iou_threshold, + top_k=max_det, + ) + indices = np.array(indices).reshape(-1) + if indices.size == 0: + return ( + np.array([]), + np.array([]), + np.array([]), + np.array([], dtype=int), + [] if label_names is not None else None, + ) + + filtered_bboxes = np.array(bbox_list)[indices].astype(np.float32) / scale + filtered_scores = np.array(score_list)[indices].astype(np.float32) + filtered_angles = np.round(np.array(angle_list)[indices], 0) + filtered_bboxes = np.clip(filtered_bboxes, 0, 1) + + labels = np.zeros(len(filtered_bboxes), dtype=int) + mapped_label_names = ( + [label_names[label] for label in labels] if label_names is not None else None + ) + return ( + filtered_bboxes, + filtered_scores, + filtered_angles, + labels, + mapped_label_names, + ) diff --git a/depthai_nodes/node/parsers/utils/mlsd.py b/depthai_nodes/node/parsers/utils/mlsd.py index fc3ec058..7d47a9fa 100644 --- a/depthai_nodes/node/parsers/utils/mlsd.py +++ b/depthai_nodes/node/parsers/utils/mlsd.py @@ -83,3 +83,20 @@ def get_lines( lines = 2 * lines / input_size # scale: 256→512 return lines, pts_score[keep].tolist() + + +def compute_mlsd_lines( + tpMap: np.ndarray, + heat: np.ndarray, + *, + topk_n: int, + score_thr: float, + dist_thr: float, +) -> tuple[np.ndarray, np.ndarray]: + """Decode MLSD outputs into lines and scores.""" + if len(tpMap.shape) != 4: + raise ValueError("Invalid shape of the tpMap tensor. Should be 4D.") + + pts, pts_score, vmap = decode_scores_and_points(tpMap, heat, topk_n) + lines, scores = get_lines(pts, pts_score, vmap, score_thr, dist_thr) + return lines, np.array(scores, dtype=np.float32) diff --git a/depthai_nodes/node/parsers/utils/ppdet.py b/depthai_nodes/node/parsers/utils/ppdet.py index 3d06d3b9..31e9c29c 100644 --- a/depthai_nodes/node/parsers/utils/ppdet.py +++ b/depthai_nodes/node/parsers/utils/ppdet.py @@ -173,3 +173,22 @@ def parse_paddle_detection_outputs( if boxes.size > 0: boxes = np.clip(boxes, 0.0, 1.0) return boxes, angles, scores + + +def compute_pp_text_detections( + predictions: np.ndarray, + *, + mask_threshold: float, + conf_threshold: float, + max_det: int, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Decode PaddleOCR text detection output into boxes, angles, and scores.""" + _, _, height, width = predictions.shape + return parse_paddle_detection_outputs( + predictions, + mask_threshold, + conf_threshold, + max_det, + width=width, + height=height, + ) diff --git a/depthai_nodes/node/parsers/utils/regression.py b/depthai_nodes/node/parsers/utils/regression.py new file mode 100644 index 00000000..3ca44001 --- /dev/null +++ b/depthai_nodes/node/parsers/utils/regression.py @@ -0,0 +1,7 @@ +import numpy as np + + +def compute_regression_predictions(predictions: np.ndarray) -> list[float]: + """Convert a regression tensor into a flat Python list.""" + squeezed = np.asarray(predictions).squeeze() + return np.atleast_1d(squeezed).tolist() diff --git a/depthai_nodes/node/parsers/utils/rf_detr.py b/depthai_nodes/node/parsers/utils/rf_detr.py new file mode 100644 index 00000000..bb9db5ac --- /dev/null +++ b/depthai_nodes/node/parsers/utils/rf_detr.py @@ -0,0 +1,86 @@ +import numpy as np + +from depthai_nodes.node.parsers.utils import sigmoid, xyxy_to_xywh +from depthai_nodes.node.parsers.utils.bbox_format_converters import xywh_to_xyxy +from depthai_nodes.node.parsers.utils.masks_utils import process_single_mask_rfdetr + + +def compute_rfdetr_detections( + boxes_tensor: np.ndarray, + logits_tensor: np.ndarray, + *, + conf_threshold: float, + max_det: int, + label_names: list[str] | None, + mask_conf: float, + input_shape: tuple[int, int] | None, + masks_tensor: np.ndarray | None = None, +) -> tuple[np.ndarray, np.ndarray, np.ndarray, list[str] | None, np.ndarray | None]: + """Decode RF-DETR detections and optional masks.""" + prob = sigmoid(logits_tensor) + + scores = np.max(prob, axis=2).squeeze() + labels = np.argmax(prob, axis=2).squeeze() + + sorted_idx = np.argsort(scores)[::-1] + + effective_max_det = max_det + if masks_tensor is not None: + max_segmentation_instances = 255 + num_valid_instances = int( + np.count_nonzero(scores[sorted_idx][:max_det] > conf_threshold) + ) + effective_max_det = min(effective_max_det, max_segmentation_instances) + if num_valid_instances > max_segmentation_instances: + num_valid_instances = max_segmentation_instances + + scores = scores[sorted_idx][:effective_max_det] + labels = labels[sorted_idx][:effective_max_det] + boxes_cxcywh = boxes_tensor.squeeze()[sorted_idx][:effective_max_det] + + masks = None + if masks_tensor is not None: + masks = masks_tensor.squeeze()[sorted_idx][:effective_max_det] + + boxes = np.clip(xywh_to_xyxy(boxes_cxcywh), 0, 1) + + confidence_mask = scores > conf_threshold + scores = scores[confidence_mask] + labels = labels[confidence_mask] + boxes = boxes[confidence_mask] + boxes_cxcywh = boxes_cxcywh[confidence_mask] + if masks is not None: + masks = masks[confidence_mask] + + final_mask = None + if masks is not None: + if input_shape is None: + raise ValueError( + "RFDETRParser segmentation mode requires model input shape." + ) + + final_mask = np.full(input_shape, 255, dtype=np.uint8) + for i, (mask_logits, bbox) in enumerate(zip(masks, boxes_cxcywh)): + resized_mask = process_single_mask_rfdetr( + mask_logits=mask_logits, + mask_conf=mask_conf, + bbox=bbox, + input_shape=input_shape, + ) + foreground = resized_mask > 0 + final_mask[(final_mask == 255) & foreground] = i + + boxes = xyxy_to_xywh(boxes) + + label_names_list = None + if label_names: + label_names_list = [ + ( + label_names[int(label)] + if int(label) < len(label_names) + else f"class_{int(label)}" + ) + for label in labels + ] + + return boxes, scores, labels.astype(int), label_names_list, final_mask diff --git a/depthai_nodes/node/parsers/utils/scrfd.py b/depthai_nodes/node/parsers/utils/scrfd.py index cd5afbcc..c4e5e556 100644 --- a/depthai_nodes/node/parsers/utils/scrfd.py +++ b/depthai_nodes/node/parsers/utils/scrfd.py @@ -166,3 +166,36 @@ def decode_scrfd( keypoints = np.clip(keypoints, 0, 1) return bboxes, scores, keypoints + + +def compute_scrfd_detections( + *, + bboxes_concatenated: list[np.ndarray], + scores_concatenated: list[np.ndarray], + kps_concatenated: list[np.ndarray], + feat_stride_fpn: tuple[int, ...] | list[int], + input_size: tuple[int, int], + num_anchors: int, + score_threshold: float, + nms_threshold: float, + anchors: dict[int, np.ndarray], + label_names: list[str] | None = None, +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, list[str] | None]: + """Decode SCRFD outputs into final detection payloads.""" + bboxes, scores, keypoints = decode_scrfd( + bboxes_concatenated=bboxes_concatenated, + scores_concatenated=scores_concatenated, + kps_concatenated=kps_concatenated, + feat_stride_fpn=feat_stride_fpn, + input_size=input_size, + num_anchors=num_anchors, + score_threshold=score_threshold, + nms_threshold=nms_threshold, + anchors=anchors, + ) + + labels = np.zeros(len(bboxes), dtype=int) + mapped_label_names = ( + [label_names[label] for label in labels] if label_names else None + ) + return bboxes, scores, keypoints, labels, mapped_label_names diff --git a/depthai_nodes/node/parsers/utils/segmentation.py b/depthai_nodes/node/parsers/utils/segmentation.py new file mode 100644 index 00000000..7109e4f4 --- /dev/null +++ b/depthai_nodes/node/parsers/utils/segmentation.py @@ -0,0 +1,46 @@ +import numpy as np + + +def compute_segmentation_class_map( + segmentation_mask: np.ndarray, + *, + classes_in_one_layer: bool = False, +) -> np.ndarray: + """Convert segmentation logits into a class map.""" + mask = np.asarray(segmentation_mask) + + if mask.ndim == 4: + mask = mask[0] + + if mask.ndim != 3: + raise ValueError(f"Expected 3D output tensor, got {mask.ndim}D.") + + reduce_fn = np.argmax + mask_shape = mask.shape + min_dim = np.argmin(mask_shape) + if min_dim == len(mask_shape) - 1: + mask = mask.transpose(2, 0, 1) + + adding_unassigned_class = False + if mask.shape[0] == 1: + if classes_in_one_layer: + reduce_fn = np.max + else: + adding_unassigned_class = True + mask = np.vstack( + ( + np.zeros((1, mask.shape[1], mask.shape[2]), dtype=np.float32), + mask, + ) + ) + + class_map = ( + reduce_fn(mask, axis=0) + .reshape(mask.shape[1], mask.shape[2]) + .astype(np.int16) + ) + + if adding_unassigned_class: + class_map = class_map - 1 + + return class_map diff --git a/depthai_nodes/node/parsers/utils/superanimal.py b/depthai_nodes/node/parsers/utils/superanimal.py index 5ccf02f8..ae91f63e 100644 --- a/depthai_nodes/node/parsers/utils/superanimal.py +++ b/depthai_nodes/node/parsers/utils/superanimal.py @@ -60,3 +60,20 @@ def get_pose_prediction(heatmap, locref, scale_factors): ) return pose + + +def compute_superanimal_keypoints( + heatmaps: np.ndarray, + *, + scale_factor: float, +) -> tuple[np.ndarray, np.ndarray]: + """Extract normalized keypoints and scores from SuperAnimal heatmaps.""" + heatmaps_scale_factor = ( + scale_factor / heatmaps.shape[1], + scale_factor / heatmaps.shape[2], + ) + + keypoints = get_pose_prediction(heatmaps, None, heatmaps_scale_factor)[0][0] + scores = keypoints[:, 2] + keypoints = keypoints[:, :2] / scale_factor + return keypoints, scores diff --git a/depthai_nodes/node/parsers/utils/xfeat.py b/depthai_nodes/node/parsers/utils/xfeat.py index 1972399d..608cbef0 100644 --- a/depthai_nodes/node/parsers/utils/xfeat.py +++ b/depthai_nodes/node/parsers/utils/xfeat.py @@ -285,6 +285,28 @@ def detect_and_compute( return result +def compute_xfeat_result( + feats: np.ndarray, + keypoints: np.ndarray, + heatmaps: np.ndarray, + *, + resize_rate_w: float, + resize_rate_h: float, + input_size: tuple[int, int], + max_keypoints: int, +) -> list[dict[str, Any]] | None: + """Compute XFeat keypoints, scores, and descriptors.""" + return detect_and_compute( + feats, + keypoints, + heatmaps, + resize_rate_w, + resize_rate_h, + input_size, + max_keypoints, + ) + + def _match_mkpts( feats1: np.ndarray, feats2: np.ndarray, min_cossim: float = 0.62 ) -> tuple[np.ndarray, np.ndarray]: @@ -343,3 +365,10 @@ def match( mkpts1 = result2["keypoints"][indexes2] return mkpts0, mkpts1 + + +def compute_xfeat_matches( + result1: dict[str, Any], result2: dict[str, Any], min_cossim: float = -1 +) -> tuple[np.ndarray, np.ndarray]: + """Compute matched XFeat keypoints.""" + return match(result1, result2, min_cossim=min_cossim) diff --git a/depthai_nodes/node/parsers/utils/yolo.py b/depthai_nodes/node/parsers/utils/yolo.py index d8552559..99c8b3c9 100644 --- a/depthai_nodes/node/parsers/utils/yolo.py +++ b/depthai_nodes/node/parsers/utils/yolo.py @@ -1,10 +1,17 @@ import time from enum import Enum +import cv2 import numpy as np from depthai_nodes.logging import get_logger -from depthai_nodes.node.parsers.utils import sigmoid, xywh_to_xyxy +from depthai_nodes.node.parsers.utils import ( + normalize_bboxes, + sigmoid, + xywh_to_xyxy, + xyxy_to_xywh, +) +from depthai_nodes.node.parsers.utils.masks_utils import process_single_mask from depthai_nodes.node.parsers.utils.nms import nms logger = get_logger(__name__) @@ -512,3 +519,193 @@ def decode_yolo_output( )[0] return output_nms + + +def compute_yolo_detections( + *, + subtype: YOLOSubtype, + layer_names: list[str], + outputs_values: list[np.ndarray], + conf_threshold: float, + n_classes: int, + iou_threshold: float, + max_det: int, + anchors: np.ndarray | None, + n_keypoints: int, + label_names: list[str] | None, + keypoint_label_names: list[str] | None, + keypoint_edges: list[tuple[int, int]] | None, + input_shape: tuple[int, int] | None = None, + kpts_outputs: list[np.ndarray] | None = None, + masks_outputs_values: list[np.ndarray] | None = None, + protos_output: np.ndarray | None = None, + protos_len: int | None = None, + mask_conf: float = 0.5, + v26_mask_coeffs: np.ndarray | None = None, + v26_protos: np.ndarray | None = None, + v26_pose_kpts: np.ndarray | None = None, +) -> dict[str, np.ndarray | list[str] | int | None]: + """Decode YOLO detection, pose, or segmentation outputs into message payloads.""" + det_mode = 0 + kpts_mode = 1 + seg_mode = 2 + + if subtype == YOLOSubtype.V26: + if any("kpt_output" in name for name in layer_names): + mode = kpts_mode + elif any("output_masks" in name for name in layer_names): + mode = seg_mode + else: + mode = det_mode + elif any("kpt_output" in name for name in layer_names) and subtype != YOLOSubtype.P: + mode = kpts_mode + elif any("_masks" in name for name in layer_names) and subtype != YOLOSubtype.P: + mode = seg_mode + else: + mode = det_mode + + if subtype == YOLOSubtype.V26: + if input_shape is None: + raise ValueError("YOLO26 parsing requires model input shape in head_config.") + + if len(outputs_values) != 1: + raise ValueError("YOLO26 requires detection output layer.") + + if mode == seg_mode: + extra_raw = v26_mask_coeffs + elif mode == kpts_mode: + extra_raw = v26_pose_kpts + else: + extra_raw = None + + results, extra = decode_yolo26( + outputs_values[0], + conf_threshold, + max_det, + extra_raw=extra_raw, + ) + + if mode == seg_mode: + v26_seg_mask_coeffs = extra + v26_seg_protos = v26_protos + elif mode == kpts_mode: + v26_pose_kpts = extra + else: + strides = ( + [8, 16, 32] + if subtype not in [YOLOSubtype.V3UT, YOLOSubtype.V3T, YOLOSubtype.V4T] + else [16, 32] + ) + + if input_shape is None: + input_shape = tuple( + dim * strides[0] for dim in outputs_values[0].shape[2:4] + ) + + if anchors is not None: + anchors = np.array(anchors).reshape(len(strides), -1) + + num_classes_check = ( + outputs_values[0].shape[1] - 5 + if anchors is None + else (outputs_values[0].shape[1] // anchors.shape[0]) - 5 + ) + if num_classes_check != n_classes: + raise ValueError( + f"The provided number of classes {n_classes} does not match the model's {num_classes_check}." + ) + + if mode == kpts_mode: + num_keypoints_check = kpts_outputs[0].shape[1] // 3 + if num_keypoints_check != n_keypoints: + raise ValueError( + f"The provided number of keypoints {n_keypoints} does not match the model's {num_keypoints_check}." + ) + + results = decode_yolo_output( + outputs_values, + strides, + anchors, + kpts=kpts_outputs if mode == kpts_mode else None, + conf_thres=conf_threshold, + iou_thres=iou_threshold, + num_classes=n_classes, + det_mode=mode == det_mode, + subtype=subtype, + ) + + bboxes = [] + labels = [] + mapped_label_names = [] + scores = [] + additional_output = [] + final_mask = np.full(input_shape, 255, dtype=np.uint8) + + for i in range(results.shape[0]): + bbox, conf, label, other = ( + results[i, :4], + results[i, 4], + results[i, 5].astype(int), + results[i, 6:], + ) + + bbox = xyxy_to_xywh(bbox.reshape(1, 4)) + bbox = normalize_bboxes(bbox, height=input_shape[0], width=input_shape[1])[0] + bboxes.append(bbox) + labels.append(int(label)) + if label_names: + mapped_label_names.append(label_names[int(label)]) + scores.append(conf) + + if mode == kpts_mode: + if subtype == YOLOSubtype.V26: + kpts = parse_kpts(v26_pose_kpts[i], n_keypoints, input_shape) + else: + kpts = parse_kpts(other, n_keypoints, input_shape) + additional_output.append(kpts) + elif mode == seg_mode: + if subtype == YOLOSubtype.V26: + mask_coeff = v26_seg_mask_coeffs[i] + mask = process_single_mask( + v26_seg_protos, mask_coeff, mask_conf, bbox + ) + else: + seg_coeff = other.astype(int) + hi, ai, xi, yi = seg_coeff + mask_coeff = masks_outputs_values[hi][ + 0, ai * protos_len : (ai + 1) * protos_len, yi, xi + ] + mask = process_single_mask( + protos_output[0], mask_coeff, mask_conf, bbox + ) + resized_mask = cv2.resize( + mask, + (input_shape[1], input_shape[0]), + interpolation=cv2.INTER_NEAREST, + ) + final_mask[resized_mask > 0] = i + + bboxes = np.array(bboxes) + bboxes = np.clip(bboxes, 0, 1) + + keypoints = np.array([]) + keypoints_scores = np.array([]) + if mode == kpts_mode: + additional_output = np.array(additional_output) + if additional_output.size > 0: + keypoints = additional_output[:, :, :2] + keypoints_scores = additional_output[:, :, 2] + keypoints = np.clip(keypoints, 0, 1) + + return { + "mode": mode, + "bboxes": bboxes, + "scores": np.array(scores), + "labels": np.array(labels), + "label_names": mapped_label_names, + "keypoints": keypoints, + "keypoints_scores": keypoints_scores, + "keypoint_label_names": keypoint_label_names, + "keypoint_edges": keypoint_edges, + "masks": final_mask if mode == seg_mode else None, + } diff --git a/depthai_nodes/node/parsers/utils/yunet.py b/depthai_nodes/node/parsers/utils/yunet.py index 14ff2d37..3b279a1b 100644 --- a/depthai_nodes/node/parsers/utils/yunet.py +++ b/depthai_nodes/node/parsers/utils/yunet.py @@ -274,3 +274,65 @@ def decode_and_prune_detections( scores_filtered = scores_filtered[:, np.newaxis] return bboxes, keypoints, scores_filtered + + +def compute_yunet_detections( + *, + input_size: tuple[int, int], + loc: np.ndarray, + conf: np.ndarray, + iou: np.ndarray, + conf_threshold: float, + iou_threshold: float, + max_det: int, + anchors: np.ndarray, + label_names: list[str] | None = None, + nms_fn=None, + top_left_wh_to_xywh_fn=None, +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, list[str] | None]: + """Decode YuNet outputs into final detection payloads.""" + bboxes, keypoints, scores = decode_and_prune_detections( + input_size=input_size, + loc=loc, + conf=conf, + iou=iou, + conf_threshold=conf_threshold, + anchors=anchors, + ) + + if len(bboxes) == 0: + return ( + np.array([]), + np.array([]), + np.array([]), + np.array([], dtype=int), + [] if label_names is not None else None, + ) + + bboxes, keypoints, scores = format_detections( + bboxes=bboxes, + keypoints=keypoints, + scores=scores, + input_size=input_size, + ) + + keep_indices = nms_fn( + bboxes=bboxes, + scores=scores, + conf_threshold=conf_threshold, + iou_threshold=iou_threshold, + max_det=max_det, + ) + + bboxes = top_left_wh_to_xywh_fn(bboxes[keep_indices]) + keypoints = keypoints[keep_indices] + scores = scores[keep_indices] + + bboxes = np.clip(bboxes, 0, 1) + keypoints = np.clip(keypoints, 0, 1) + + labels = np.zeros(len(bboxes), dtype=int) + mapped_label_names = ( + [label_names[label] for label in labels] if label_names else None + ) + return bboxes, keypoints, scores, labels, mapped_label_names diff --git a/depthai_nodes/node/parsers/xfeat.py b/depthai_nodes/node/parsers/xfeat.py index 0ff5b8f1..d27fd8a8 100644 --- a/depthai_nodes/node/parsers/xfeat.py +++ b/depthai_nodes/node/parsers/xfeat.py @@ -5,7 +5,10 @@ from depthai_nodes.message.creators import create_tracked_features_message from depthai_nodes.node.parsers.base_parser import BaseParser -from depthai_nodes.node.parsers.utils.xfeat import detect_and_compute, match +from depthai_nodes.node.parsers.utils.xfeat import ( + compute_xfeat_matches, + compute_xfeat_result, +) class XFeatBaseParser(BaseParser): @@ -344,55 +347,71 @@ def run(self): except dai.MessageQueue.QueueException: break # Pipeline was stopped - self._logger.debug( - f"Processing input with layers: {output.getAllLayerNames()}" - ) - feats, keypoints, heatmaps = self.extractTensors(output) - - result = detect_and_compute( + feats, keypoints, heatmaps = self.extract(output) + result = self.compute( feats, keypoints, heatmaps, - resize_rate_w, - resize_rate_h, - self.input_size, - self.max_keypoints, + resize_rate_w=resize_rate_w, + resize_rate_h=resize_rate_h, ) + self.emit(output, result) - if result is not None: - result = result[0] - else: - matched_points = dai.TrackedFeatures() - matched_points.setTimestamp(output.getTimestamp()) - matched_points.setSequenceNum(output.getSequenceNum()) - self._logger.debug( - "No keypoints found, sending TrackedFeatures message" - ) - self.out.send(matched_points) - self._logger.debug("TrackedFeatures message sent") - continue - - if self.previous_results is not None: - mkpts0, mkpts1 = match(self.previous_results, result) - matched_points = create_tracked_features_message(mkpts0, mkpts1) - matched_points.setTimestamp(output.getTimestamp()) - matched_points.setSequenceNum(output.getSequenceNum()) - self._logger.debug("Keypoints found, sending TrackedFeatures message") - self.out.send(matched_points) - self._logger.debug("TrackedFeatures message sent") - else: - matched_points = dai.TrackedFeatures() - matched_points.setTimestamp(output.getTimestamp()) - matched_points.setSequenceNum(output.getSequenceNum()) - self._logger.debug( - "No previous results, sending TrackedFeatures message" - ) - self.out.send(matched_points) - self._logger.debug("TrackedFeatures message sent") - - if self.trigger: - self.previous_results = result - self.trigger = False + def extract(self, output: dai.NNData) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + self._logger.debug(f"Processing input with layers: {output.getAllLayerNames()}") + return self.extractTensors(output) + + def compute( + self, + feats: np.ndarray, + keypoints: np.ndarray, + heatmaps: np.ndarray, + *, + resize_rate_w: float, + resize_rate_h: float, + ) -> dict[str, Any] | None: + result = compute_xfeat_result( + feats, + keypoints, + heatmaps, + resize_rate_w=resize_rate_w, + resize_rate_h=resize_rate_h, + input_size=self.input_size, + max_keypoints=self.max_keypoints, + ) + if result is not None: + return result[0] + return None + + def emit(self, output: dai.NNData, result: dict[str, Any] | None) -> None: + if result is None: + matched_points = dai.TrackedFeatures() + matched_points.setTimestamp(output.getTimestamp()) + matched_points.setSequenceNum(output.getSequenceNum()) + self._logger.debug("No keypoints found, sending TrackedFeatures message") + self.out.send(matched_points) + self._logger.debug("TrackedFeatures message sent") + return + + if self.previous_results is not None: + mkpts0, mkpts1 = compute_xfeat_matches(self.previous_results, result) + matched_points = create_tracked_features_message(mkpts0, mkpts1) + matched_points.setTimestamp(output.getTimestamp()) + matched_points.setSequenceNum(output.getSequenceNum()) + self._logger.debug("Keypoints found, sending TrackedFeatures message") + self.out.send(matched_points) + self._logger.debug("TrackedFeatures message sent") + else: + matched_points = dai.TrackedFeatures() + matched_points.setTimestamp(output.getTimestamp()) + matched_points.setSequenceNum(output.getSequenceNum()) + self._logger.debug("No previous results, sending TrackedFeatures message") + self.out.send(matched_points) + self._logger.debug("TrackedFeatures message sent") + + if self.trigger: + self.previous_results = result + self.trigger = False class XFeatStereoParser(XFeatBaseParser): @@ -486,73 +505,81 @@ def run(self): except dai.MessageQueue.QueueException: break # Pipeline was stopped - self._logger.debug( - f"Processing reference input with layers: {reference_output.getAllLayerNames()}" + reference_tensors, target_tensors = self.extract( + reference_output, target_output ) - self._logger.debug( - f"Processing target input with layers: {target_output.getAllLayerNames()}" + match_result = self.compute( + reference_tensors, + target_tensors, + resize_rate_w=resize_rate_w, + resize_rate_h=resize_rate_h, ) + self.emit(reference_output, target_output, match_result) + + def extract( + self, reference_output: dai.NNData, target_output: dai.NNData + ) -> tuple[ + tuple[np.ndarray, np.ndarray, np.ndarray], + tuple[np.ndarray, np.ndarray, np.ndarray], + ]: + self._logger.debug( + f"Processing reference input with layers: {reference_output.getAllLayerNames()}" + ) + self._logger.debug( + f"Processing target input with layers: {target_output.getAllLayerNames()}" + ) + return self.extractTensors(reference_output), self.extractTensors(target_output) - ( - reference_feats, - reference_keypoints, - reference_heatmaps, - ) = self.extractTensors(reference_output) - target_feats, target_keypoints, target_heatmaps = self.extractTensors( - target_output - ) + def compute( + self, + reference_tensors: tuple[np.ndarray, np.ndarray, np.ndarray], + target_tensors: tuple[np.ndarray, np.ndarray, np.ndarray], + *, + resize_rate_w: float, + resize_rate_h: float, + ) -> tuple[np.ndarray, np.ndarray] | None: + reference_result = compute_xfeat_result( + *reference_tensors, + resize_rate_w=resize_rate_w, + resize_rate_h=resize_rate_h, + input_size=self.input_size, + max_keypoints=self.max_keypoints, + ) + target_result = compute_xfeat_result( + *target_tensors, + resize_rate_w=resize_rate_w, + resize_rate_h=resize_rate_h, + input_size=self.input_size, + max_keypoints=self.max_keypoints, + ) - reference_result = detect_and_compute( - reference_feats, - reference_keypoints, - reference_heatmaps, - resize_rate_w, - resize_rate_h, - self.input_size, - self.max_keypoints, - ) + if reference_result is None or target_result is None: + return None - target_result = detect_and_compute( - target_feats, - target_keypoints, - target_heatmaps, - resize_rate_w, - resize_rate_h, - self.input_size, - self.max_keypoints, - ) + return compute_xfeat_matches(reference_result[0], target_result[0]) - if reference_result is not None: - reference_result = reference_result[0] - else: - matched_points = dai.TrackedFeatures() - matched_points.setTimestamp(reference_output.getTimestamp()) - matched_points.setSequenceNum(reference_output.getSequenceNum()) - self._logger.debug( - "No reference keypoints found, sending TrackedFeatures message" - ) - self.out.send(matched_points) - self._logger.debug("TrackedFeatures message sent") - continue - - if target_result is not None: - target_result = target_result[0] - else: - matched_points = dai.TrackedFeatures() - matched_points.setTimestamp(target_output.getTimestamp()) - matched_points.setSequenceNum(reference_output.getSequenceNum()) - self._logger.debug( - "No target keypoints found, sending TrackedFeatures message" - ) - self.out.send(matched_points) - self._logger.debug("TrackedFeatures message sent") - continue - - mkpts0, mkpts1 = match(reference_result, target_result) - matched_points = create_tracked_features_message(mkpts0, mkpts1) + def emit( + self, + reference_output: dai.NNData, + target_output: dai.NNData, + match_result: tuple[np.ndarray, np.ndarray] | None, + ) -> None: + if match_result is None: + matched_points = dai.TrackedFeatures() matched_points.setTimestamp(target_output.getTimestamp()) matched_points.setSequenceNum(reference_output.getSequenceNum()) - matched_points.setTimestampDevice(target_output.getTimestampDevice()) - self._logger.debug("Keypoints found, sending TrackedFeatures message") + self._logger.debug( + "No stereo keypoints found, sending TrackedFeatures message" + ) self.out.send(matched_points) self._logger.debug("TrackedFeatures message sent") + return + + mkpts0, mkpts1 = match_result + matched_points = create_tracked_features_message(mkpts0, mkpts1) + matched_points.setTimestamp(target_output.getTimestamp()) + matched_points.setSequenceNum(reference_output.getSequenceNum()) + matched_points.setTimestampDevice(target_output.getTimestampDevice()) + self._logger.debug("Keypoints found, sending TrackedFeatures message") + self.out.send(matched_points) + self._logger.debug("TrackedFeatures message sent") diff --git a/depthai_nodes/node/parsers/yolo.py b/depthai_nodes/node/parsers/yolo.py index 01ec1d81..8b176e41 100644 --- a/depthai_nodes/node/parsers/yolo.py +++ b/depthai_nodes/node/parsers/yolo.py @@ -16,6 +16,7 @@ ) from depthai_nodes.node.parsers.utils.yolo import ( YOLOSubtype, + compute_yolo_detections, decode_yolo26, decode_yolo_output, parse_kpts, @@ -421,67 +422,61 @@ def run(self): output: dai.NNData = self.input.get() except dai.MessageQueue.QueueException: break # Pipeline was stopped, no more data - # Get all the layer names - layer_names = self.output_layer_names or output.getAllLayerNames() - self._logger.debug(f"Processing input with layers: {layer_names}") - - if self.subtype == YOLOSubtype.V26: - # YOLO26 end2end: infer task type from output layer names - if any("output_masks" in name for name in layer_names): - outputs_values = [ - output.getTensor("output_yolo26", dequantize=True).astype( - np.float32 - ) - ] - self._mask_coeffs = output.getTensor( - "output_masks", dequantize=True - ).astype(np.float32) - self._protos = output.getTensor( - self._protos_layer_name, - dequantize=True, - storageOrder=dai.TensorInfo.StorageOrder.NCHW, - ).astype(np.float32) - elif any("kpt_output" in name for name in layer_names): - outputs_values = [ - output.getTensor("output_yolo26", dequantize=True).astype( - np.float32 - ) - ] - self._kpts_output = output.getTensor( - "kpt_output", dequantize=True - ).astype(np.float32) - else: - outputs_names = list(layer_names) - outputs_values = [ - output.getTensor(o, dequantize=True).astype(np.float32) - for o in outputs_names - ] + extracted = self.extract(output) + payload = self.compute(**extracted) + self.emit(output, payload) + + def extract(self, output: dai.NNData) -> dict[str, Any]: + layer_names = self.output_layer_names or output.getAllLayerNames() + self._logger.debug(f"Processing input with layers: {layer_names}") + + kpts_outputs = None + masks_outputs_values = None + protos_output = None + protos_len = None + v26_mask_coeffs = None + v26_protos = None + v26_pose_kpts = None + + if self.subtype == YOLOSubtype.V26: + if any("output_masks" in name for name in layer_names): + outputs_values = [ + output.getTensor("output_yolo26", dequantize=True).astype(np.float32) + ] + v26_mask_coeffs = output.getTensor( + "output_masks", dequantize=True + ).astype(np.float32) + v26_protos = output.getTensor( + self._protos_layer_name, + dequantize=True, + storageOrder=dai.TensorInfo.StorageOrder.NCHW, + ).astype(np.float32)[0] + elif any("kpt_output" in name for name in layer_names): + outputs_values = [ + output.getTensor("output_yolo26", dequantize=True).astype(np.float32) + ] + v26_pose_kpts = output.getTensor( + "kpt_output", dequantize=True + ).astype(np.float32) else: - outputs_names = sorted( - [name for name in layer_names if "_yolo" in name or "yolo-" in name] - ) outputs_values = [ - output.getTensor( - o, - dequantize=True, - storageOrder=dai.TensorInfo.StorageOrder.NCHW, - ).astype(np.float32) - for o in outputs_names + output.getTensor(o, dequantize=True).astype(np.float32) + for o in list(layer_names) ] + else: + outputs_names = sorted( + [name for name in layer_names if "_yolo" in name or "yolo-" in name] + ) + outputs_values = [ + output.getTensor( + o, + dequantize=True, + storageOrder=dai.TensorInfo.StorageOrder.NCHW, + ).astype(np.float32) + for o in outputs_names + ] - if self.subtype == YOLOSubtype.V26: - if any("kpt_output" in name for name in layer_names): - mode = self._KPTS_MODE - elif any("output_masks" in name for name in layer_names): - mode = self._SEG_MODE - else: - mode = self._DET_MODE - elif ( - any("kpt_output" in name for name in layer_names) - and self.subtype != YOLOSubtype.P - ): - mode = self._KPTS_MODE - # Get the keypoint outputs + if any("kpt_output" in name for name in layer_names) and self.subtype != YOLOSubtype.P: kpts_output_names = sorted( [name for name in layer_names if "kpt_output" in name] ) @@ -489,201 +484,79 @@ def run(self): output.getTensor(o, dequantize=True).astype(np.float32) for o in kpts_output_names ] - elif ( - any("_masks" in name for name in layer_names) - and self.subtype != YOLOSubtype.P - ): - mode = self._SEG_MODE - # Get the segmentation outputs + elif any("_masks" in name for name in layer_names) and self.subtype != YOLOSubtype.P: ( masks_outputs_values, protos_output, protos_len, ) = get_segmentation_outputs(output) - else: - mode = self._DET_MODE - - # Get the model's input shape - if self.subtype == YOLOSubtype.V26: - if self.input_shape is None: - raise ValueError( - "YOLO26 parsing requires model input shape in head_config." - ) - input_shape = self.input_shape - else: - strides = ( - [8, 16, 32] - if self.subtype - not in [YOLOSubtype.V3UT, YOLOSubtype.V3T, YOLOSubtype.V4T] - else [16, 32] - ) - input_shape = tuple( - dim * strides[0] for dim in outputs_values[0].shape[2:4] - ) - - # Reshape the anchors based on the model's output heads - if self.subtype != YOLOSubtype.V26 and self.anchors is not None: - self.anchors = np.array(self.anchors).reshape(len(strides), -1) - - # Ensure the number of classes is correct - if self.subtype != YOLOSubtype.V26: - num_classes_check = ( - outputs_values[0].shape[1] - 5 - if self.anchors is None - else (outputs_values[0].shape[1] // self.anchors.shape[0]) - 5 - ) - if num_classes_check != self.n_classes: - raise ValueError( - f"The provided number of classes {self.n_classes} does not match the model's {num_classes_check}." - ) - # Ensure the number of keypoints is correct (skip for V26 pose which handles kpts differently) - if mode == self._KPTS_MODE and self.subtype != YOLOSubtype.V26: - num_keypoints_check = kpts_outputs[0].shape[1] // 3 - if num_keypoints_check != self.n_keypoints: - raise ValueError( - f"The provided number of keypoints {self.n_keypoints} does not match the model's {num_keypoints_check}." - ) - - # Decode the outputs - if self.subtype == YOLOSubtype.V26: - if len(outputs_values) != 1: - raise ValueError("YOLO26 requires detection output layer.") - if mode == self._SEG_MODE: - extra_raw = self._mask_coeffs - elif mode == self._KPTS_MODE: - extra_raw = self._kpts_output - else: - extra_raw = None - - results, extra = decode_yolo26( - outputs_values[0], - self.conf_threshold, - self.max_det, - extra_raw=extra_raw, - ) - - if mode == self._SEG_MODE: - self._v26_seg_mask_coeffs = extra - self._v26_seg_protos = self._protos[0] # (nm, H, W) - elif mode == self._KPTS_MODE: - self._v26_pose_kpts = extra - else: - results = decode_yolo_output( - outputs_values, - strides, - self.anchors, - kpts=kpts_outputs if mode == self._KPTS_MODE else None, - conf_thres=self.conf_threshold, - iou_thres=self.iou_threshold, - num_classes=self.n_classes, - det_mode=mode == self._DET_MODE, - subtype=self.subtype, - ) - - bboxes, labels, label_names, scores, additional_output = [], [], [], [], [] - final_mask = np.full(input_shape, 255, dtype=np.uint8) - for i in range(results.shape[0]): - bbox, conf, label, other = ( - results[i, :4], - results[i, 4], - results[i, 5].astype(int), - results[i, 6:], - ) - - bbox = xyxy_to_xywh(bbox.reshape(1, 4)) - bbox = normalize_bboxes( - bbox, height=input_shape[0], width=input_shape[1] - )[0] - bboxes.append(bbox) - labels.append(int(label)) - if self.label_names: - label_names.append(self.label_names[int(label)]) - scores.append(conf) - - if mode == self._KPTS_MODE: - if self.subtype == YOLOSubtype.V26: - kpts = parse_kpts( - self._v26_pose_kpts[i], self.n_keypoints, input_shape - ) - else: - kpts = parse_kpts(other, self.n_keypoints, input_shape) - additional_output.append(kpts) - elif mode == self._SEG_MODE: - if self.subtype == YOLOSubtype.V26: - # V26 segmentation: mask coefficients are directly available - mask_coeff = self._v26_seg_mask_coeffs[i] - mask = process_single_mask( - self._v26_seg_protos, mask_coeff, self.mask_conf, bbox - ) - else: - # Other YOLO versions: extract mask coefficients from indexed outputs - seg_coeff = other.astype(int) - hi, ai, xi, yi = seg_coeff - mask_coeff = masks_outputs_values[hi][ - 0, ai * protos_len : (ai + 1) * protos_len, yi, xi - ] - mask = process_single_mask( - protos_output[0], mask_coeff, self.mask_conf, bbox - ) - # Resize mask to input shape - resized_mask = cv2.resize( - mask, - (input_shape[1], input_shape[0]), - interpolation=cv2.INTER_NEAREST, - ) - # Fill the final mask with the instance values - final_mask[resized_mask > 0] = i - - bboxes = np.array(bboxes) - bboxes = np.clip(bboxes, 0, 1) - - if mode == self._KPTS_MODE: - additional_output = np.array(additional_output) - keypoints = np.array([]) - keypoints_scores = np.array([]) - if additional_output.size > 0: - keypoints = additional_output[:, :, :2] - keypoints_scores = additional_output[:, :, 2] - - keypoints = np.clip(keypoints, 0, 1) - detections_message = create_detection_message( - bboxes=bboxes, - scores=np.array(scores), - labels=np.array(labels), - label_names=label_names, - keypoints=keypoints, - keypoints_scores=keypoints_scores, - keypoint_label_names=self.keypoint_label_names, - keypoint_edges=self.keypoint_edges, - ) - elif mode == self._SEG_MODE: - detections_message = create_detection_message( - bboxes=bboxes, - scores=np.array(scores), - labels=np.array(labels), - label_names=label_names, - masks=final_mask, - ) - else: - detections_message = create_detection_message( - bboxes=bboxes, - scores=np.array(scores), - labels=np.array(labels), - label_names=label_names, - ) - - detections_message.setTimestamp(output.getTimestamp()) - detections_message.setTimestampDevice(output.getTimestampDevice()) - detections_message.setSequenceNum(output.getSequenceNum()) - transformation = output.getTransformation() - if transformation is not None: - detections_message.setTransformation(transformation) - - self._logger.debug( - f"Created detection message with {len(bboxes)} detections" + return { + "subtype": self.subtype, + "layer_names": list(layer_names), + "outputs_values": outputs_values, + "conf_threshold": self.conf_threshold, + "n_classes": self.n_classes, + "iou_threshold": self.iou_threshold, + "max_det": self.max_det, + "anchors": self.anchors, + "n_keypoints": self.n_keypoints, + "label_names": self.label_names, + "keypoint_label_names": self.keypoint_label_names, + "keypoint_edges": self.keypoint_edges, + "input_shape": self.input_shape, + "kpts_outputs": kpts_outputs, + "masks_outputs_values": masks_outputs_values, + "protos_output": protos_output, + "protos_len": protos_len, + "mask_conf": self.mask_conf, + "v26_mask_coeffs": v26_mask_coeffs, + "v26_protos": v26_protos, + "v26_pose_kpts": v26_pose_kpts, + } + + @staticmethod + def compute(**kwargs) -> dict[str, Any]: + return compute_yolo_detections(**kwargs) + + def emit(self, output: dai.NNData, payload: dict[str, Any]) -> None: + mode = payload["mode"] + if mode == self._KPTS_MODE: + detections_message = create_detection_message( + bboxes=payload["bboxes"], + scores=payload["scores"], + labels=payload["labels"], + label_names=payload["label_names"], + keypoints=payload["keypoints"], + keypoints_scores=payload["keypoints_scores"], + keypoint_label_names=payload["keypoint_label_names"], + keypoint_edges=payload["keypoint_edges"], + ) + elif mode == self._SEG_MODE: + detections_message = create_detection_message( + bboxes=payload["bboxes"], + scores=payload["scores"], + labels=payload["labels"], + label_names=payload["label_names"], + masks=payload["masks"], + ) + else: + detections_message = create_detection_message( + bboxes=payload["bboxes"], + scores=payload["scores"], + labels=payload["labels"], + label_names=payload["label_names"], ) - self.out.send(detections_message) + detections_message.setTimestamp(output.getTimestamp()) + detections_message.setTimestampDevice(output.getTimestampDevice()) + detections_message.setSequenceNum(output.getSequenceNum()) + transformation = output.getTransformation() + if transformation is not None: + detections_message.setTransformation(transformation) - self._logger.debug("Detection message sent successfully") + self._logger.debug( + f"Created detection message with {len(payload['bboxes'])} detections" + ) + self.out.send(detections_message) + self._logger.debug("Detection message sent successfully") diff --git a/depthai_nodes/node/parsers/yunet.py b/depthai_nodes/node/parsers/yunet.py index 07e18748..914bf45f 100644 --- a/depthai_nodes/node/parsers/yunet.py +++ b/depthai_nodes/node/parsers/yunet.py @@ -8,6 +8,7 @@ from depthai_nodes.node.parsers.utils import top_left_wh_to_xywh from depthai_nodes.node.parsers.utils.nms import nms_cv2 from depthai_nodes.node.parsers.utils.yunet import ( + compute_yunet_detections, decode_and_prune_detections, format_detections, generate_anchors, @@ -202,169 +203,122 @@ def run(self): except dai.MessageQueue.QueueException: break # Pipeline was stopped - output_layer_names = output.getAllLayerNames() - self._logger.debug(f"Processing input with layers: {output_layer_names}") - - # get loc - if self.loc_output_layer_name: - try: - loc = output.getTensor(self.loc_output_layer_name, dequantize=True) - except KeyError as err: - raise ValueError( - f"Layer {self.loc_output_layer_name} not found in the model output." - ) from err - else: - loc_output_layer_name_candidates = [ - layer_name - for layer_name in output_layer_names - if layer_name.startswith("loc") - ] - if len(loc_output_layer_name_candidates) == 0: - raise ValueError( - "No loc layer candidates found in the model output." - ) - elif len(loc_output_layer_name_candidates) > 1: - raise ValueError( - "Multiple loc layer candidates found in the model output." - ) - else: - self.loc_output_layer_name = loc_output_layer_name_candidates[0] - - # get conf - if self.conf_output_layer_name: - try: - conf = output.getTensor( - self.conf_output_layer_name, dequantize=True - ) - except KeyError as err: - raise ValueError( - f"Layer {self.conf_output_layer_name} not found in the model output." - ) from err - else: - conf_output_layer_name_candidates = [ - layer_name - for layer_name in output_layer_names - if layer_name.startswith("conf") - ] - if len(conf_output_layer_name_candidates) == 0: - raise ValueError( - "No conf layer candidates found in the model output." - ) - elif len(conf_output_layer_name_candidates) > 1: - raise ValueError( - "Multiple conf layer candidates found in the model output." - ) - else: - self.conf_output_layer_name = conf_output_layer_name_candidates[0] - - # get iou - if self.iou_output_layer_name: - try: - iou = output.getTensor(self.iou_output_layer_name, dequantize=True) - except KeyError as err: - raise ValueError( - f"Layer {self.iou_output_layer_name} not found in the model output." - ) from err - else: - iou_output_layer_name_candidates = [ - layer_name - for layer_name in output_layer_names - if layer_name.startswith("iou") - ] - if len(iou_output_layer_name_candidates) == 0: - raise ValueError( - "No iou layer candidates found in the model output." - ) - elif len(iou_output_layer_name_candidates) > 1: - raise ValueError( - "Multiple iou layer candidates found in the model output." - ) - else: - self.iou_output_layer_name = iou_output_layer_name_candidates[0] - - loc = output.getTensor(self.loc_output_layer_name, dequantize=True) - conf = output.getTensor(self.conf_output_layer_name, dequantize=True) - iou = output.getTensor(self.iou_output_layer_name, dequantize=True) - - # Use optimized combined decode and prune function - anchors = self._get_cached_anchors() - bboxes, keypoints, scores = decode_and_prune_detections( + loc, conf, iou = self.extract(output) + bboxes, keypoints, scores, labels, label_names = self.compute( input_size=self.input_size, loc=loc, conf=conf, iou=iou, conf_threshold=self.conf_threshold, - anchors=anchors, - ) - - # Skip further processing if no detections - if len(bboxes) == 0: - detections_message = create_detection_message( - bboxes=np.array([]), - scores=np.array([]), - keypoints=np.array([]), - labels=np.array([]), - label_names=[], - ) - detections_message.setTimestamp(output.getTimestamp()) - detections_message.setTimestampDevice(output.getTimestampDevice()) - transformation = output.getTransformation() - if transformation is not None: - detections_message.setTransformation(transformation) - detections_message.setSequenceNum(output.getSequenceNum()) - self.out.send(detections_message) - continue - - # format detections - bboxes, keypoints, scores = format_detections( - bboxes=bboxes, - keypoints=keypoints, - scores=scores, - input_size=self.input_size, - ) - - # run nms - keep_indices = nms_cv2( - bboxes=bboxes, - scores=scores, - conf_threshold=self.conf_threshold, iou_threshold=self.iou_threshold, max_det=self.max_det, + anchors=self._get_cached_anchors(), + label_names=self.label_names, + nms_fn=nms_cv2, + top_left_wh_to_xywh_fn=top_left_wh_to_xywh, ) + self.emit(output, bboxes, keypoints, scores, labels, label_names) - bboxes = bboxes[keep_indices] - bboxes = top_left_wh_to_xywh(bboxes) - keypoints = keypoints[keep_indices] - scores = scores[keep_indices] + def extract(self, output: dai.NNData) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + output_layer_names = output.getAllLayerNames() + self._logger.debug(f"Processing input with layers: {output_layer_names}") - bboxes = np.clip(bboxes, 0, 1) - keypoints = np.clip(keypoints, 0, 1) + if self.loc_output_layer_name: + try: + loc = output.getTensor(self.loc_output_layer_name, dequantize=True) + except KeyError as err: + raise ValueError( + f"Layer {self.loc_output_layer_name} not found in the model output." + ) from err + else: + loc_output_layer_name_candidates = [ + layer_name + for layer_name in output_layer_names + if layer_name.startswith("loc") + ] + if len(loc_output_layer_name_candidates) == 0: + raise ValueError("No loc layer candidates found in the model output.") + if len(loc_output_layer_name_candidates) > 1: + raise ValueError( + "Multiple loc layer candidates found in the model output." + ) + self.loc_output_layer_name = loc_output_layer_name_candidates[0] + loc = output.getTensor(self.loc_output_layer_name, dequantize=True) - labels = np.array([0] * len(bboxes)) + if self.conf_output_layer_name: + try: + conf = output.getTensor(self.conf_output_layer_name, dequantize=True) + except KeyError as err: + raise ValueError( + f"Layer {self.conf_output_layer_name} not found in the model output." + ) from err + else: + conf_output_layer_name_candidates = [ + layer_name + for layer_name in output_layer_names + if layer_name.startswith("conf") + ] + if len(conf_output_layer_name_candidates) == 0: + raise ValueError("No conf layer candidates found in the model output.") + if len(conf_output_layer_name_candidates) > 1: + raise ValueError( + "Multiple conf layer candidates found in the model output." + ) + self.conf_output_layer_name = conf_output_layer_name_candidates[0] + conf = output.getTensor(self.conf_output_layer_name, dequantize=True) - label_names = ( - [self.label_names[label] for label in labels] - if self.label_names - else None - ) + if self.iou_output_layer_name: + try: + iou = output.getTensor(self.iou_output_layer_name, dequantize=True) + except KeyError as err: + raise ValueError( + f"Layer {self.iou_output_layer_name} not found in the model output." + ) from err + else: + iou_output_layer_name_candidates = [ + layer_name + for layer_name in output_layer_names + if layer_name.startswith("iou") + ] + if len(iou_output_layer_name_candidates) == 0: + raise ValueError("No iou layer candidates found in the model output.") + if len(iou_output_layer_name_candidates) > 1: + raise ValueError( + "Multiple iou layer candidates found in the model output." + ) + self.iou_output_layer_name = iou_output_layer_name_candidates[0] + iou = output.getTensor(self.iou_output_layer_name, dequantize=True) - detections_message = create_detection_message( - bboxes=bboxes, - scores=scores, - keypoints=keypoints, - labels=labels, - label_names=label_names, - ) + return loc, conf, iou - detections_message.setTimestamp(output.getTimestamp()) - detections_message.setSequenceNum(output.getSequenceNum()) - detections_message.setTimestampDevice(output.getTimestampDevice()) - transformation = output.getTransformation() - if transformation is not None: - detections_message.setTransformation(transformation) + @staticmethod + def compute(**kwargs): + return compute_yunet_detections(**kwargs) - self._logger.debug(f"Created detections message with {len(bboxes)} faces") + def emit( + self, + output: dai.NNData, + bboxes: np.ndarray, + keypoints: np.ndarray, + scores: np.ndarray, + labels: np.ndarray, + label_names: list[str] | None, + ) -> None: + detections_message = create_detection_message( + bboxes=bboxes, + scores=scores, + keypoints=keypoints, + labels=labels, + label_names=label_names, + ) - self.out.send(detections_message) + detections_message.setTimestamp(output.getTimestamp()) + detections_message.setSequenceNum(output.getSequenceNum()) + detections_message.setTimestampDevice(output.getTimestampDevice()) + transformation = output.getTransformation() + if transformation is not None: + detections_message.setTransformation(transformation) - self._logger.debug("Detections message sent successfully") + self._logger.debug(f"Created detections message with {len(bboxes)} faces") + self.out.send(detections_message) + self._logger.debug("Detections message sent successfully") From a34c3f0ec42db039434298ee35ebc255ad0c6f21 Mon Sep 17 00:00:00 2001 From: Dani Rogmans Date: Wed, 1 Jul 2026 10:45:24 +0200 Subject: [PATCH 03/11] XFeatStereoParser preserve old fallback behavior, restore the RF-DETR truncation warning --- depthai_nodes/node/parsers/rf_detr.py | 3 +- depthai_nodes/node/parsers/utils/rf_detr.py | 10 +++++-- depthai_nodes/node/parsers/xfeat.py | 31 +++++++++++++++------ 3 files changed, 32 insertions(+), 12 deletions(-) diff --git a/depthai_nodes/node/parsers/rf_detr.py b/depthai_nodes/node/parsers/rf_detr.py index 7dbe533b..ffd62707 100644 --- a/depthai_nodes/node/parsers/rf_detr.py +++ b/depthai_nodes/node/parsers/rf_detr.py @@ -237,8 +237,8 @@ def extract( ) return boxes_tensor, logits_tensor, masks_tensor - @staticmethod def compute( + self, boxes_tensor: np.ndarray, logits_tensor: np.ndarray, *, @@ -258,6 +258,7 @@ def compute( mask_conf=mask_conf, input_shape=input_shape, masks_tensor=masks_tensor, + logger=self._logger, ) def emit( diff --git a/depthai_nodes/node/parsers/utils/rf_detr.py b/depthai_nodes/node/parsers/utils/rf_detr.py index bb9db5ac..01e9ea29 100644 --- a/depthai_nodes/node/parsers/utils/rf_detr.py +++ b/depthai_nodes/node/parsers/utils/rf_detr.py @@ -15,6 +15,7 @@ def compute_rfdetr_detections( mask_conf: float, input_shape: tuple[int, int] | None, masks_tensor: np.ndarray | None = None, + logger=None, ) -> tuple[np.ndarray, np.ndarray, np.ndarray, list[str] | None, np.ndarray | None]: """Decode RF-DETR detections and optional masks.""" prob = sigmoid(logits_tensor) @@ -30,9 +31,14 @@ def compute_rfdetr_detections( num_valid_instances = int( np.count_nonzero(scores[sorted_idx][:max_det] > conf_threshold) ) + if num_valid_instances > max_segmentation_instances and logger is not None: + logger.warning( + "RFDETRParser can encode at most 255 instances in " + "SegmentationMask; ignoring " + f"{num_valid_instances - max_segmentation_instances} " + "lowest-scoring instances." + ) effective_max_det = min(effective_max_det, max_segmentation_instances) - if num_valid_instances > max_segmentation_instances: - num_valid_instances = max_segmentation_instances scores = scores[sorted_idx][:effective_max_det] labels = labels[sorted_idx][:effective_max_det] diff --git a/depthai_nodes/node/parsers/xfeat.py b/depthai_nodes/node/parsers/xfeat.py index d27fd8a8..08b0e330 100644 --- a/depthai_nodes/node/parsers/xfeat.py +++ b/depthai_nodes/node/parsers/xfeat.py @@ -537,7 +537,7 @@ def compute( *, resize_rate_w: float, resize_rate_h: float, - ) -> tuple[np.ndarray, np.ndarray] | None: + ) -> dict[str, tuple[np.ndarray, np.ndarray] | str | None]: reference_result = compute_xfeat_result( *reference_tensors, resize_rate_w=resize_rate_w, @@ -553,24 +553,37 @@ def compute( max_keypoints=self.max_keypoints, ) - if reference_result is None or target_result is None: - return None + if reference_result is None: + return {"status": "reference_missing", "match_result": None} + if target_result is None: + return {"status": "target_missing", "match_result": None} - return compute_xfeat_matches(reference_result[0], target_result[0]) + return { + "status": "matched", + "match_result": compute_xfeat_matches(reference_result[0], target_result[0]), + } def emit( self, reference_output: dai.NNData, target_output: dai.NNData, - match_result: tuple[np.ndarray, np.ndarray] | None, + result: dict[str, tuple[np.ndarray, np.ndarray] | str | None], ) -> None: + status = result["status"] + match_result = result["match_result"] if match_result is None: matched_points = dai.TrackedFeatures() - matched_points.setTimestamp(target_output.getTimestamp()) + if status == "reference_missing": + matched_points.setTimestamp(reference_output.getTimestamp()) + self._logger.debug( + "No reference keypoints found, sending TrackedFeatures message" + ) + else: + matched_points.setTimestamp(target_output.getTimestamp()) + self._logger.debug( + "No target keypoints found, sending TrackedFeatures message" + ) matched_points.setSequenceNum(reference_output.getSequenceNum()) - self._logger.debug( - "No stereo keypoints found, sending TrackedFeatures message" - ) self.out.send(matched_points) self._logger.debug("TrackedFeatures message sent") return From 0ab544b2ee46622cbebfb93ea3f545ce18c8b8ea Mon Sep 17 00:00:00 2001 From: Dani Rogmans Date: Wed, 1 Jul 2026 11:10:37 +0200 Subject: [PATCH 04/11] pre-commit fix --- .../node/parsers/classification_sequence.py | 4 +-- depthai_nodes/node/parsers/fastsam.py | 6 ---- .../node/parsers/mediapipe_palm_detection.py | 2 -- depthai_nodes/node/parsers/utils/keypoints.py | 4 +-- .../node/parsers/utils/segmentation.py | 4 +-- depthai_nodes/node/parsers/utils/yolo.py | 8 ++--- depthai_nodes/node/parsers/xfeat.py | 4 ++- depthai_nodes/node/parsers/yolo.py | 33 ++++++++++--------- depthai_nodes/node/parsers/yunet.py | 2 -- 9 files changed, 27 insertions(+), 40 deletions(-) diff --git a/depthai_nodes/node/parsers/classification_sequence.py b/depthai_nodes/node/parsers/classification_sequence.py index e5697f1f..5290a931 100644 --- a/depthai_nodes/node/parsers/classification_sequence.py +++ b/depthai_nodes/node/parsers/classification_sequence.py @@ -183,9 +183,7 @@ def extract(self, output: dai.NNData) -> np.ndarray: @staticmethod def compute(scores: np.ndarray, *, is_softmax: bool = True) -> np.ndarray: - return compute_classification_sequence_scores( - scores, is_softmax=is_softmax - ) + 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( diff --git a/depthai_nodes/node/parsers/fastsam.py b/depthai_nodes/node/parsers/fastsam.py index c453ec89..409f3e98 100644 --- a/depthai_nodes/node/parsers/fastsam.py +++ b/depthai_nodes/node/parsers/fastsam.py @@ -6,13 +6,7 @@ from depthai_nodes.message.creators import create_segmentation_message from depthai_nodes.node.parsers.base_parser import BaseParser from depthai_nodes.node.parsers.utils.fastsam import ( - box_prompt, - build_mask_coeffs, compute_fastsam_mask, - decode_fastsam_output, - merge_masks, - point_prompt, - process_masks, ) from depthai_nodes.node.parsers.utils.masks_utils import get_segmentation_outputs diff --git a/depthai_nodes/node/parsers/mediapipe_palm_detection.py b/depthai_nodes/node/parsers/mediapipe_palm_detection.py index 4b8b570d..f69e65fe 100644 --- a/depthai_nodes/node/parsers/mediapipe_palm_detection.py +++ b/depthai_nodes/node/parsers/mediapipe_palm_detection.py @@ -1,6 +1,5 @@ from typing import Any -import cv2 import depthai as dai import numpy as np @@ -8,7 +7,6 @@ from depthai_nodes.node.parsers.detection import DetectionParser from depthai_nodes.node.parsers.utils.medipipe import ( compute_mediapipe_palm_detections, - decode, generate_handtracker_anchors, ) diff --git a/depthai_nodes/node/parsers/utils/keypoints.py b/depthai_nodes/node/parsers/utils/keypoints.py index 731cc68d..5cbf0531 100644 --- a/depthai_nodes/node/parsers/utils/keypoints.py +++ b/depthai_nodes/node/parsers/utils/keypoints.py @@ -50,9 +50,7 @@ def compute_keypoints( num_coords = int(np.prod(parsed_keypoints.shape) / n_keypoints) if num_coords not in [2, 3]: - raise ValueError( - f"Expected 2 or 3 coordinates per keypoint, got {num_coords}." - ) + raise ValueError(f"Expected 2 or 3 coordinates per keypoint, got {num_coords}.") parsed_keypoints = parsed_keypoints.reshape(n_keypoints, num_coords) parsed_keypoints /= scale_factor diff --git a/depthai_nodes/node/parsers/utils/segmentation.py b/depthai_nodes/node/parsers/utils/segmentation.py index 7109e4f4..7603d140 100644 --- a/depthai_nodes/node/parsers/utils/segmentation.py +++ b/depthai_nodes/node/parsers/utils/segmentation.py @@ -35,9 +35,7 @@ def compute_segmentation_class_map( ) class_map = ( - reduce_fn(mask, axis=0) - .reshape(mask.shape[1], mask.shape[2]) - .astype(np.int16) + reduce_fn(mask, axis=0).reshape(mask.shape[1], mask.shape[2]).astype(np.int16) ) if adding_unassigned_class: diff --git a/depthai_nodes/node/parsers/utils/yolo.py b/depthai_nodes/node/parsers/utils/yolo.py index 99c8b3c9..83a7a33b 100644 --- a/depthai_nodes/node/parsers/utils/yolo.py +++ b/depthai_nodes/node/parsers/utils/yolo.py @@ -566,7 +566,9 @@ def compute_yolo_detections( if subtype == YOLOSubtype.V26: if input_shape is None: - raise ValueError("YOLO26 parsing requires model input shape in head_config.") + raise ValueError( + "YOLO26 parsing requires model input shape in head_config." + ) if len(outputs_values) != 1: raise ValueError("YOLO26 requires detection output layer.") @@ -666,9 +668,7 @@ def compute_yolo_detections( elif mode == seg_mode: if subtype == YOLOSubtype.V26: mask_coeff = v26_seg_mask_coeffs[i] - mask = process_single_mask( - v26_seg_protos, mask_coeff, mask_conf, bbox - ) + mask = process_single_mask(v26_seg_protos, mask_coeff, mask_conf, bbox) else: seg_coeff = other.astype(int) hi, ai, xi, yi = seg_coeff diff --git a/depthai_nodes/node/parsers/xfeat.py b/depthai_nodes/node/parsers/xfeat.py index 08b0e330..7342d61a 100644 --- a/depthai_nodes/node/parsers/xfeat.py +++ b/depthai_nodes/node/parsers/xfeat.py @@ -560,7 +560,9 @@ def compute( return { "status": "matched", - "match_result": compute_xfeat_matches(reference_result[0], target_result[0]), + "match_result": compute_xfeat_matches( + reference_result[0], target_result[0] + ), } def emit( diff --git a/depthai_nodes/node/parsers/yolo.py b/depthai_nodes/node/parsers/yolo.py index 8b176e41..5464a52a 100644 --- a/depthai_nodes/node/parsers/yolo.py +++ b/depthai_nodes/node/parsers/yolo.py @@ -1,25 +1,16 @@ from typing import Any -import cv2 import depthai as dai import numpy as np from depthai_nodes.message.creators import create_detection_message from depthai_nodes.node.parsers.base_parser import BaseParser -from depthai_nodes.node.parsers.utils import ( - normalize_bboxes, - xyxy_to_xywh, -) from depthai_nodes.node.parsers.utils.masks_utils import ( get_segmentation_outputs, - process_single_mask, ) from depthai_nodes.node.parsers.utils.yolo import ( YOLOSubtype, compute_yolo_detections, - decode_yolo26, - decode_yolo_output, - parse_kpts, ) @@ -441,7 +432,9 @@ def extract(self, output: dai.NNData) -> dict[str, Any]: if self.subtype == YOLOSubtype.V26: if any("output_masks" in name for name in layer_names): outputs_values = [ - output.getTensor("output_yolo26", dequantize=True).astype(np.float32) + output.getTensor("output_yolo26", dequantize=True).astype( + np.float32 + ) ] v26_mask_coeffs = output.getTensor( "output_masks", dequantize=True @@ -453,11 +446,13 @@ def extract(self, output: dai.NNData) -> dict[str, Any]: ).astype(np.float32)[0] elif any("kpt_output" in name for name in layer_names): outputs_values = [ - output.getTensor("output_yolo26", dequantize=True).astype(np.float32) + output.getTensor("output_yolo26", dequantize=True).astype( + np.float32 + ) ] - v26_pose_kpts = output.getTensor( - "kpt_output", dequantize=True - ).astype(np.float32) + v26_pose_kpts = output.getTensor("kpt_output", dequantize=True).astype( + np.float32 + ) else: outputs_values = [ output.getTensor(o, dequantize=True).astype(np.float32) @@ -476,7 +471,10 @@ def extract(self, output: dai.NNData) -> dict[str, Any]: for o in outputs_names ] - if any("kpt_output" in name for name in layer_names) and self.subtype != YOLOSubtype.P: + if ( + any("kpt_output" in name for name in layer_names) + and self.subtype != YOLOSubtype.P + ): kpts_output_names = sorted( [name for name in layer_names if "kpt_output" in name] ) @@ -484,7 +482,10 @@ def extract(self, output: dai.NNData) -> dict[str, Any]: output.getTensor(o, dequantize=True).astype(np.float32) for o in kpts_output_names ] - elif any("_masks" in name for name in layer_names) and self.subtype != YOLOSubtype.P: + elif ( + any("_masks" in name for name in layer_names) + and self.subtype != YOLOSubtype.P + ): ( masks_outputs_values, protos_output, diff --git a/depthai_nodes/node/parsers/yunet.py b/depthai_nodes/node/parsers/yunet.py index 914bf45f..7e4cc240 100644 --- a/depthai_nodes/node/parsers/yunet.py +++ b/depthai_nodes/node/parsers/yunet.py @@ -9,8 +9,6 @@ from depthai_nodes.node.parsers.utils.nms import nms_cv2 from depthai_nodes.node.parsers.utils.yunet import ( compute_yunet_detections, - decode_and_prune_detections, - format_detections, generate_anchors, ) From 3f928e4d380a209dc26d62673d5791ba138fd02a Mon Sep 17 00:00:00 2001 From: Dani Rogmans Date: Thu, 2 Jul 2026 13:32:44 +0200 Subject: [PATCH 05/11] Address pr-relevant coderabbit concerns --- depthai_nodes/node/parsers/embeddings.py | 31 +++++++++++++++----- depthai_nodes/node/parsers/hrnet.py | 5 ++-- depthai_nodes/node/parsers/utils/medipipe.py | 6 +++- depthai_nodes/node/parsers/utils/ppdet.py | 15 +++++++++- depthai_nodes/node/parsers/utils/yolo.py | 17 +++++++++-- depthai_nodes/node/parsers/utils/yunet.py | 11 +++---- depthai_nodes/node/parsers/xfeat.py | 7 +++++ 7 files changed, 72 insertions(+), 20 deletions(-) diff --git a/depthai_nodes/node/parsers/embeddings.py b/depthai_nodes/node/parsers/embeddings.py index 3dae4e4a..6ee38fb5 100644 --- a/depthai_nodes/node/parsers/embeddings.py +++ b/depthai_nodes/node/parsers/embeddings.py @@ -24,7 +24,7 @@ class EmbeddingsParser(BaseParser): def __init__(self) -> None: """Initialize the EmbeddingsParser node.""" super().__init__() - self.output_layer_name: str = None + self.output_layer_name: str | list[str] | None = None self._logger.debug( f"EmbeddingsParser initialized with output_layer_name={self.output_layer_name}" ) @@ -51,9 +51,12 @@ def build(self, head_config: dict[str, Any]) -> "EmbeddingsParser": """ self.output_layer_name = head_config["outputs"] - assert ( - len(self.output_layer_name) == 1 - ), "Embeddings head should have only one output layer" + output_names = self._normalize_output_layer_names( + self.output_layer_name + ) + assert len(output_names) == 1, ( + "Embeddings head should have only one output layer" + ) self._logger.debug( f"EmbeddingsParser built with output_layer_name={self.output_layer_name}" @@ -74,14 +77,26 @@ def run(self): self.emit(computed) def extract(self, output: dai.NNData) -> dai.NNData: - output_names = self.output_layer_name or output.getAllLayerNames() + output_names = self._normalize_output_layer_names( + self.output_layer_name or output.getAllLayerNames() + ) self._logger.debug(f"Processing input with layers: {output_names}") - assert ( - len(output_names) == 1 - ), "Embeddings head should have only one output layer" + assert len(output_names) == 1, ( + "Embeddings head should have only one output layer" + ) return output + @staticmethod + def _normalize_output_layer_names( + output_layer_name: str | list[str] | None, + ) -> list[str]: + if output_layer_name is None: + return [] + if isinstance(output_layer_name, str): + return [output_layer_name] + return list(output_layer_name) + @staticmethod def compute(output: dai.NNData) -> dai.NNData: return compute_embeddings_output(output) diff --git a/depthai_nodes/node/parsers/hrnet.py b/depthai_nodes/node/parsers/hrnet.py index 9f77fe45..9cd36c4e 100644 --- a/depthai_nodes/node/parsers/hrnet.py +++ b/depthai_nodes/node/parsers/hrnet.py @@ -116,14 +116,15 @@ def extract(self, output: dai.NNData) -> np.ndarray: dequantize=True, ) - def compute(self, heatmaps: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + @staticmethod + def compute(heatmaps: np.ndarray) -> tuple[np.ndarray, np.ndarray]: keypoints, scores = compute_hrnet_keypoints(heatmaps) - self.n_keypoints = len(keypoints) return keypoints, scores def emit( self, output: dai.NNData, keypoints: np.ndarray, scores: np.ndarray ) -> None: + self.n_keypoints = len(keypoints) keypoints_message = create_keypoints_message( keypoints=keypoints, scores=scores, diff --git a/depthai_nodes/node/parsers/utils/medipipe.py b/depthai_nodes/node/parsers/utils/medipipe.py index 38cfd268..2e3cd514 100644 --- a/depthai_nodes/node/parsers/utils/medipipe.py +++ b/depthai_nodes/node/parsers/utils/medipipe.py @@ -408,6 +408,7 @@ def compute_mediapipe_palm_detections( ) bbox_list = [] + nms_bbox_list = [] score_list = [] angle_list = [] for hand in decoded_bboxes: @@ -420,8 +421,11 @@ def compute_mediapipe_palm_detections( x_center, y_center = np.mean(extended_points, axis=0) width = np.linalg.norm(extended_points[0] - extended_points[3]) height = np.linalg.norm(extended_points[0] - extended_points[1]) + x_min = x_center - width / 2 + y_min = y_center - height / 2 bbox_list.append([x_center, y_center, width, height]) + nms_bbox_list.append([x_min, y_min, width, height]) angle_list.append(angle) score_list.append(hand.pd_score) @@ -435,7 +439,7 @@ def compute_mediapipe_palm_detections( ) indices = cv2.dnn.NMSBoxes( - bbox_list, + nms_bbox_list, score_list, conf_threshold, iou_threshold, diff --git a/depthai_nodes/node/parsers/utils/ppdet.py b/depthai_nodes/node/parsers/utils/ppdet.py index 31e9c29c..16fd2f79 100644 --- a/depthai_nodes/node/parsers/utils/ppdet.py +++ b/depthai_nodes/node/parsers/utils/ppdet.py @@ -183,7 +183,20 @@ def compute_pp_text_detections( max_det: int, ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: """Decode PaddleOCR text detection output into boxes, angles, and scores.""" - _, _, height, width = predictions.shape + if predictions.ndim != 4: + raise ValueError( + f"Predictions should be 4D array of shape (1, 1, H, W) or (1, H, W, 1), got {predictions.shape}." + ) + + if predictions.shape[0] == 1 and predictions.shape[1] == 1: + height, width = predictions.shape[2], predictions.shape[3] + elif predictions.shape[0] == 1 and predictions.shape[3] == 1: + height, width = predictions.shape[1], predictions.shape[2] + else: + raise ValueError( + f"Predictions should be either (1, 1, H, W) or (1, H, W, 1), got {predictions.shape}." + ) + return parse_paddle_detection_outputs( predictions, mask_threshold, diff --git a/depthai_nodes/node/parsers/utils/yolo.py b/depthai_nodes/node/parsers/utils/yolo.py index 83a7a33b..3003cb5c 100644 --- a/depthai_nodes/node/parsers/utils/yolo.py +++ b/depthai_nodes/node/parsers/utils/yolo.py @@ -641,7 +641,14 @@ def compute_yolo_detections( mapped_label_names = [] scores = [] additional_output = [] - final_mask = np.full(input_shape, 255, dtype=np.uint8) + final_mask = None + if mode == seg_mode: + 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." + ) + final_mask = np.full(input_shape, -1, dtype=np.int16) for i in range(results.shape[0]): bbox, conf, label, other = ( @@ -656,7 +663,11 @@ def compute_yolo_detections( bboxes.append(bbox) labels.append(int(label)) if label_names: - mapped_label_names.append(label_names[int(label)]) + label_idx = int(label) + if label_idx < len(label_names): + mapped_label_names.append(label_names[label_idx]) + else: + mapped_label_names.append(f"class_{label_idx}") scores.append(conf) if mode == kpts_mode: @@ -707,5 +718,5 @@ def compute_yolo_detections( "keypoints_scores": keypoints_scores, "keypoint_label_names": keypoint_label_names, "keypoint_edges": keypoint_edges, - "masks": final_mask if mode == seg_mode else None, + "masks": final_mask, } diff --git a/depthai_nodes/node/parsers/utils/yunet.py b/depthai_nodes/node/parsers/utils/yunet.py index 3b279a1b..11b30b6c 100644 --- a/depthai_nodes/node/parsers/utils/yunet.py +++ b/depthai_nodes/node/parsers/utils/yunet.py @@ -1,3 +1,4 @@ +from collections.abc import Callable from itertools import product import numpy as np @@ -287,8 +288,8 @@ def compute_yunet_detections( max_det: int, anchors: np.ndarray, label_names: list[str] | None = None, - nms_fn=None, - top_left_wh_to_xywh_fn=None, + nms_fn: Callable[..., np.ndarray], + top_left_wh_to_xywh_fn: Callable[[np.ndarray], np.ndarray], ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, list[str] | None]: """Decode YuNet outputs into final detection payloads.""" bboxes, keypoints, scores = decode_and_prune_detections( @@ -302,9 +303,9 @@ def compute_yunet_detections( if len(bboxes) == 0: return ( - np.array([]), - np.array([]), - np.array([]), + np.empty((0, 4), dtype=np.float32), + np.empty((0, 5, 2), dtype=np.float32), + np.empty((0, 1), dtype=np.float32), np.array([], dtype=int), [] if label_names is not None else None, ) diff --git a/depthai_nodes/node/parsers/xfeat.py b/depthai_nodes/node/parsers/xfeat.py index 7342d61a..535065d2 100644 --- a/depthai_nodes/node/parsers/xfeat.py +++ b/depthai_nodes/node/parsers/xfeat.py @@ -388,6 +388,7 @@ def emit(self, output: dai.NNData, result: dict[str, Any] | None) -> None: matched_points = dai.TrackedFeatures() matched_points.setTimestamp(output.getTimestamp()) matched_points.setSequenceNum(output.getSequenceNum()) + matched_points.setTimestampDevice(output.getTimestampDevice()) self._logger.debug("No keypoints found, sending TrackedFeatures message") self.out.send(matched_points) self._logger.debug("TrackedFeatures message sent") @@ -398,6 +399,7 @@ def emit(self, output: dai.NNData, result: dict[str, Any] | None) -> None: matched_points = create_tracked_features_message(mkpts0, mkpts1) matched_points.setTimestamp(output.getTimestamp()) matched_points.setSequenceNum(output.getSequenceNum()) + matched_points.setTimestampDevice(output.getTimestampDevice()) self._logger.debug("Keypoints found, sending TrackedFeatures message") self.out.send(matched_points) self._logger.debug("TrackedFeatures message sent") @@ -405,6 +407,7 @@ def emit(self, output: dai.NNData, result: dict[str, Any] | None) -> None: matched_points = dai.TrackedFeatures() matched_points.setTimestamp(output.getTimestamp()) matched_points.setSequenceNum(output.getSequenceNum()) + matched_points.setTimestampDevice(output.getTimestampDevice()) self._logger.debug("No previous results, sending TrackedFeatures message") self.out.send(matched_points) self._logger.debug("TrackedFeatures message sent") @@ -577,11 +580,15 @@ def emit( matched_points = dai.TrackedFeatures() if status == "reference_missing": matched_points.setTimestamp(reference_output.getTimestamp()) + matched_points.setTimestampDevice( + reference_output.getTimestampDevice() + ) self._logger.debug( "No reference keypoints found, sending TrackedFeatures message" ) else: matched_points.setTimestamp(target_output.getTimestamp()) + matched_points.setTimestampDevice(target_output.getTimestampDevice()) self._logger.debug( "No target keypoints found, sending TrackedFeatures message" ) From 91b0c6f6ca49edd5c0afbea925d048a1dda1a7ac Mon Sep 17 00:00:00 2001 From: Dani Rogmans Date: Thu, 2 Jul 2026 13:32:58 +0200 Subject: [PATCH 06/11] Address pr-relevant coderabbit concerns --- depthai_nodes/node/parsers/embeddings.py | 16 +++++++--------- depthai_nodes/node/parsers/xfeat.py | 4 +--- 2 files changed, 8 insertions(+), 12 deletions(-) diff --git a/depthai_nodes/node/parsers/embeddings.py b/depthai_nodes/node/parsers/embeddings.py index 6ee38fb5..2f9f0100 100644 --- a/depthai_nodes/node/parsers/embeddings.py +++ b/depthai_nodes/node/parsers/embeddings.py @@ -51,12 +51,10 @@ def build(self, head_config: dict[str, Any]) -> "EmbeddingsParser": """ self.output_layer_name = head_config["outputs"] - output_names = self._normalize_output_layer_names( - self.output_layer_name - ) - assert len(output_names) == 1, ( - "Embeddings head should have only one output layer" - ) + output_names = self._normalize_output_layer_names(self.output_layer_name) + assert ( + len(output_names) == 1 + ), "Embeddings head should have only one output layer" self._logger.debug( f"EmbeddingsParser built with output_layer_name={self.output_layer_name}" @@ -82,9 +80,9 @@ def extract(self, output: dai.NNData) -> dai.NNData: ) self._logger.debug(f"Processing input with layers: {output_names}") - assert len(output_names) == 1, ( - "Embeddings head should have only one output layer" - ) + assert ( + len(output_names) == 1 + ), "Embeddings head should have only one output layer" return output @staticmethod diff --git a/depthai_nodes/node/parsers/xfeat.py b/depthai_nodes/node/parsers/xfeat.py index 535065d2..bf1482fc 100644 --- a/depthai_nodes/node/parsers/xfeat.py +++ b/depthai_nodes/node/parsers/xfeat.py @@ -580,9 +580,7 @@ def emit( matched_points = dai.TrackedFeatures() if status == "reference_missing": matched_points.setTimestamp(reference_output.getTimestamp()) - matched_points.setTimestampDevice( - reference_output.getTimestampDevice() - ) + matched_points.setTimestampDevice(reference_output.getTimestampDevice()) self._logger.debug( "No reference keypoints found, sending TrackedFeatures message" ) From 6cd7bf3bb77db3aeaeb21edc8eeb4121b6c73138 Mon Sep 17 00:00:00 2001 From: Dani Rogmans Date: Mon, 6 Jul 2026 14:39:34 +0200 Subject: [PATCH 07/11] fixes for unnecessary defensive behavior, mention arguments explicitly --- depthai_nodes/node/parsers/classification.py | 4 +-- .../node/parsers/classification_sequence.py | 9 ------ depthai_nodes/node/parsers/embeddings.py | 13 ++++---- depthai_nodes/node/parsers/hrnet.py | 1 - depthai_nodes/node/parsers/scrfd.py | 4 --- .../node/parsers/utils/classification.py | 2 +- depthai_nodes/node/parsers/utils/scrfd.py | 3 ++ depthai_nodes/node/parsers/yunet.py | 30 +++++++++++++++++-- 8 files changed, 40 insertions(+), 26 deletions(-) diff --git a/depthai_nodes/node/parsers/classification.py b/depthai_nodes/node/parsers/classification.py index 6b08ac45..f6fde7a8 100644 --- a/depthai_nodes/node/parsers/classification.py +++ b/depthai_nodes/node/parsers/classification.py @@ -145,9 +145,7 @@ def extract(self, output: dai.NNData) -> np.ndarray: f"Expected 1 output layer, got {len(layers)} layers. Please provide the output_layer_name." ) - scores = np.asarray( - output.getTensor(self.output_layer_name, dequantize=True) - ).flatten() + scores = output.getTensor(self.output_layer_name, dequantize=True).flatten() if len(scores) != self.n_classes and self.n_classes != 0: raise ValueError( diff --git a/depthai_nodes/node/parsers/classification_sequence.py b/depthai_nodes/node/parsers/classification_sequence.py index 5290a931..88966499 100644 --- a/depthai_nodes/node/parsers/classification_sequence.py +++ b/depthai_nodes/node/parsers/classification_sequence.py @@ -151,15 +151,6 @@ 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) diff --git a/depthai_nodes/node/parsers/embeddings.py b/depthai_nodes/node/parsers/embeddings.py index 2f9f0100..b0fe97bf 100644 --- a/depthai_nodes/node/parsers/embeddings.py +++ b/depthai_nodes/node/parsers/embeddings.py @@ -24,7 +24,7 @@ class EmbeddingsParser(BaseParser): def __init__(self) -> None: """Initialize the EmbeddingsParser node.""" super().__init__() - self.output_layer_name: str | list[str] | None = None + self.output_layer_name: str | None = None self._logger.debug( f"EmbeddingsParser initialized with output_layer_name={self.output_layer_name}" ) @@ -49,12 +49,11 @@ def build(self, head_config: dict[str, Any]) -> "EmbeddingsParser": @return: The parser object with the head configuration set. @rtype: EmbeddingsParser """ - - self.output_layer_name = head_config["outputs"] - output_names = self._normalize_output_layer_names(self.output_layer_name) + output_names = self._normalize_output_layer_names(head_config["outputs"]) assert ( len(output_names) == 1 ), "Embeddings head should have only one output layer" + self.output_layer_name = output_names[0] self._logger.debug( f"EmbeddingsParser built with output_layer_name={self.output_layer_name}" @@ -75,8 +74,10 @@ def run(self): self.emit(computed) def extract(self, output: dai.NNData) -> dai.NNData: - output_names = self._normalize_output_layer_names( - self.output_layer_name or output.getAllLayerNames() + output_names = ( + [self.output_layer_name] + if self.output_layer_name is not None + else output.getAllLayerNames() ) self._logger.debug(f"Processing input with layers: {output_names}") diff --git a/depthai_nodes/node/parsers/hrnet.py b/depthai_nodes/node/parsers/hrnet.py index 9cd36c4e..db541707 100644 --- a/depthai_nodes/node/parsers/hrnet.py +++ b/depthai_nodes/node/parsers/hrnet.py @@ -124,7 +124,6 @@ def compute(heatmaps: np.ndarray) -> tuple[np.ndarray, np.ndarray]: def emit( self, output: dai.NNData, keypoints: np.ndarray, scores: np.ndarray ) -> None: - self.n_keypoints = len(keypoints) keypoints_message = create_keypoints_message( keypoints=keypoints, scores=scores, diff --git a/depthai_nodes/node/parsers/scrfd.py b/depthai_nodes/node/parsers/scrfd.py index 8df7a8fa..b1d38308 100644 --- a/depthai_nodes/node/parsers/scrfd.py +++ b/depthai_nodes/node/parsers/scrfd.py @@ -5,7 +5,6 @@ from depthai_nodes.message.creators import create_detection_message from depthai_nodes.node.parsers.detection import DetectionParser -from depthai_nodes.node.parsers.utils import xyxy_to_xywh from depthai_nodes.node.parsers.utils.scrfd import ( compute_anchor_centers, compute_scrfd_detections, @@ -262,9 +261,6 @@ def emit( labels: np.ndarray, label_names: list[str] | None, ) -> None: - bboxes = xyxy_to_xywh(bboxes) - bboxes = np.clip(bboxes, 0, 1) - message = create_detection_message( bboxes=bboxes, scores=scores, diff --git a/depthai_nodes/node/parsers/utils/classification.py b/depthai_nodes/node/parsers/utils/classification.py index bc802d5a..86dc204a 100644 --- a/depthai_nodes/node/parsers/utils/classification.py +++ b/depthai_nodes/node/parsers/utils/classification.py @@ -9,7 +9,7 @@ def compute_classification_scores( is_softmax: bool = True, ) -> np.ndarray: """Return classification scores, applying softmax when needed.""" - computed_scores = np.asarray(scores).flatten() + computed_scores = scores.flatten() if not is_softmax: computed_scores = softmax(computed_scores) return computed_scores diff --git a/depthai_nodes/node/parsers/utils/scrfd.py b/depthai_nodes/node/parsers/utils/scrfd.py index c4e5e556..2f03afaa 100644 --- a/depthai_nodes/node/parsers/utils/scrfd.py +++ b/depthai_nodes/node/parsers/utils/scrfd.py @@ -1,5 +1,6 @@ import numpy as np +from depthai_nodes.node.parsers.utils.bbox_format_converters import xyxy_to_xywh from depthai_nodes.node.parsers.utils.nms import nms @@ -193,6 +194,8 @@ def compute_scrfd_detections( nms_threshold=nms_threshold, anchors=anchors, ) + bboxes = xyxy_to_xywh(bboxes) + bboxes = np.clip(bboxes, 0, 1) labels = np.zeros(len(bboxes), dtype=int) mapped_label_names = ( diff --git a/depthai_nodes/node/parsers/yunet.py b/depthai_nodes/node/parsers/yunet.py index 7e4cc240..b1eb8ea7 100644 --- a/depthai_nodes/node/parsers/yunet.py +++ b/depthai_nodes/node/parsers/yunet.py @@ -1,3 +1,4 @@ +from collections.abc import Callable from typing import Any import depthai as dai @@ -290,8 +291,33 @@ def extract(self, output: dai.NNData) -> tuple[np.ndarray, np.ndarray, np.ndarra return loc, conf, iou @staticmethod - def compute(**kwargs): - return compute_yunet_detections(**kwargs) + def compute( + *, + input_size: tuple[int, int], + loc: np.ndarray, + conf: np.ndarray, + iou: np.ndarray, + conf_threshold: float, + iou_threshold: float, + max_det: int, + anchors: np.ndarray, + label_names: list[str] | None = None, + nms_fn: Callable[..., np.ndarray], + top_left_wh_to_xywh_fn: Callable[[np.ndarray], np.ndarray], + ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, list[str] | None]: + return compute_yunet_detections( + input_size=input_size, + loc=loc, + conf=conf, + iou=iou, + conf_threshold=conf_threshold, + iou_threshold=iou_threshold, + max_det=max_det, + anchors=anchors, + label_names=label_names, + nms_fn=nms_fn, + top_left_wh_to_xywh_fn=top_left_wh_to_xywh_fn, + ) def emit( self, From 0b6c38352f04ce74390ce1041c7162ef009392ca Mon Sep 17 00:00:00 2001 From: Dani Rogmans Date: Mon, 6 Jul 2026 14:57:55 +0200 Subject: [PATCH 08/11] YOLOComputeInputs dataclass, compute() no longer takes **kwargs --- depthai_nodes/node/parsers/yolo.py | 102 +++++++++++++++++++++-------- 1 file changed, 75 insertions(+), 27 deletions(-) diff --git a/depthai_nodes/node/parsers/yolo.py b/depthai_nodes/node/parsers/yolo.py index 5464a52a..cc5fa0ee 100644 --- a/depthai_nodes/node/parsers/yolo.py +++ b/depthai_nodes/node/parsers/yolo.py @@ -1,3 +1,4 @@ +from dataclasses import dataclass from typing import Any import depthai as dai @@ -14,6 +15,31 @@ ) +@dataclass(frozen=True) +class YOLOComputeInputs: + subtype: YOLOSubtype + layer_names: list[str] + outputs_values: list[np.ndarray] + conf_threshold: float + n_classes: int + iou_threshold: float + max_det: int + anchors: list[list[list[float]]] | np.ndarray | None + n_keypoints: int + label_names: list[str] | None + keypoint_label_names: list[str] | None + keypoint_edges: list[tuple[int, int]] | None + input_shape: tuple[int, int] | None + kpts_outputs: list[np.ndarray] | None + masks_outputs_values: list[np.ndarray] | None + protos_output: np.ndarray | None + protos_len: int | None + mask_conf: float + v26_mask_coeffs: np.ndarray | None + v26_protos: np.ndarray | None + v26_pose_kpts: np.ndarray | None + + class YOLOExtendedParser(BaseParser): """Parser class for parsing the output of the YOLO Instance Segmentation and Pose Estimation models. @@ -414,10 +440,10 @@ def run(self): except dai.MessageQueue.QueueException: break # Pipeline was stopped, no more data extracted = self.extract(output) - payload = self.compute(**extracted) + payload = self.compute(extracted) self.emit(output, payload) - def extract(self, output: dai.NNData) -> dict[str, Any]: + def extract(self, output: dai.NNData) -> YOLOComputeInputs: layer_names = self.output_layer_names or output.getAllLayerNames() self._logger.debug(f"Processing input with layers: {layer_names}") @@ -492,33 +518,55 @@ def extract(self, output: dai.NNData) -> dict[str, Any]: protos_len, ) = get_segmentation_outputs(output) - return { - "subtype": self.subtype, - "layer_names": list(layer_names), - "outputs_values": outputs_values, - "conf_threshold": self.conf_threshold, - "n_classes": self.n_classes, - "iou_threshold": self.iou_threshold, - "max_det": self.max_det, - "anchors": self.anchors, - "n_keypoints": self.n_keypoints, - "label_names": self.label_names, - "keypoint_label_names": self.keypoint_label_names, - "keypoint_edges": self.keypoint_edges, - "input_shape": self.input_shape, - "kpts_outputs": kpts_outputs, - "masks_outputs_values": masks_outputs_values, - "protos_output": protos_output, - "protos_len": protos_len, - "mask_conf": self.mask_conf, - "v26_mask_coeffs": v26_mask_coeffs, - "v26_protos": v26_protos, - "v26_pose_kpts": v26_pose_kpts, - } + return YOLOComputeInputs( + subtype=self.subtype, + layer_names=list(layer_names), + outputs_values=outputs_values, + conf_threshold=self.conf_threshold, + n_classes=self.n_classes, + iou_threshold=self.iou_threshold, + max_det=self.max_det, + anchors=self.anchors, + n_keypoints=self.n_keypoints, + label_names=self.label_names, + keypoint_label_names=self.keypoint_label_names, + keypoint_edges=self.keypoint_edges, + input_shape=self.input_shape, + kpts_outputs=kpts_outputs, + masks_outputs_values=masks_outputs_values, + protos_output=protos_output, + protos_len=protos_len, + mask_conf=self.mask_conf, + v26_mask_coeffs=v26_mask_coeffs, + v26_protos=v26_protos, + v26_pose_kpts=v26_pose_kpts, + ) @staticmethod - def compute(**kwargs) -> dict[str, Any]: - return compute_yolo_detections(**kwargs) + def compute(inputs: YOLOComputeInputs) -> dict[str, Any]: + return compute_yolo_detections( + subtype=inputs.subtype, + layer_names=inputs.layer_names, + outputs_values=inputs.outputs_values, + conf_threshold=inputs.conf_threshold, + n_classes=inputs.n_classes, + iou_threshold=inputs.iou_threshold, + max_det=inputs.max_det, + anchors=inputs.anchors, + n_keypoints=inputs.n_keypoints, + label_names=inputs.label_names, + keypoint_label_names=inputs.keypoint_label_names, + keypoint_edges=inputs.keypoint_edges, + input_shape=inputs.input_shape, + kpts_outputs=inputs.kpts_outputs, + masks_outputs_values=inputs.masks_outputs_values, + protos_output=inputs.protos_output, + protos_len=inputs.protos_len, + mask_conf=inputs.mask_conf, + v26_mask_coeffs=inputs.v26_mask_coeffs, + v26_protos=inputs.v26_protos, + v26_pose_kpts=inputs.v26_pose_kpts, + ) def emit(self, output: dai.NNData, payload: dict[str, Any]) -> None: mode = payload["mode"] From 692acdbd175e5c33b9893ff3b77be435cfd425ba Mon Sep 17 00:00:00 2001 From: Dani Rogmans Date: Wed, 8 Jul 2026 11:41:36 +0200 Subject: [PATCH 09/11] stride resolving behavior re-implemented after merge conflict for non-standard strides in nnarchive --- depthai_nodes/node/parsers/utils/yolo.py | 19 +++++++++++-------- depthai_nodes/node/parsers/yolo.py | 4 ++++ 2 files changed, 15 insertions(+), 8 deletions(-) diff --git a/depthai_nodes/node/parsers/utils/yolo.py b/depthai_nodes/node/parsers/utils/yolo.py index 9ab9411c..393a5db0 100644 --- a/depthai_nodes/node/parsers/utils/yolo.py +++ b/depthai_nodes/node/parsers/utils/yolo.py @@ -575,6 +575,7 @@ def compute_yolo_detections( subtype: YOLOSubtype, layer_names: list[str], outputs_values: list[np.ndarray], + strides: list[int] | None = None, conf_threshold: float, n_classes: int, iou_threshold: float, @@ -642,19 +643,21 @@ def compute_yolo_detections( elif mode == kpts_mode: v26_pose_kpts = extra else: - strides = ( - [8, 16, 32] - if subtype not in [YOLOSubtype.V3UT, YOLOSubtype.V3T, YOLOSubtype.V4T] - else [16, 32] - ) + resolved_strides = strides + if resolved_strides is None: + resolved_strides = ( + [8, 16, 32] + if subtype not in [YOLOSubtype.V3UT, YOLOSubtype.V3T, YOLOSubtype.V4T] + else [16, 32] + ) if input_shape is None: input_shape = tuple( - dim * strides[0] for dim in outputs_values[0].shape[2:4] + dim * resolved_strides[0] for dim in outputs_values[0].shape[2:4] ) if anchors is not None: - anchors = np.array(anchors).reshape(len(strides), -1) + anchors = np.array(anchors).reshape(len(resolved_strides), -1) num_classes_check = ( outputs_values[0].shape[1] - 5 @@ -675,7 +678,7 @@ def compute_yolo_detections( results = decode_yolo_output( outputs_values, - strides, + resolved_strides, anchors, kpts=kpts_outputs if mode == kpts_mode else None, conf_thres=conf_threshold, diff --git a/depthai_nodes/node/parsers/yolo.py b/depthai_nodes/node/parsers/yolo.py index dce9f0e8..99308818 100644 --- a/depthai_nodes/node/parsers/yolo.py +++ b/depthai_nodes/node/parsers/yolo.py @@ -21,6 +21,7 @@ class YOLOComputeInputs: subtype: YOLOSubtype layer_names: list[str] outputs_values: list[np.ndarray] + strides: list[int] | None conf_threshold: float n_classes: int iou_threshold: float @@ -452,6 +453,7 @@ def extract(self, output: dai.NNData) -> YOLOComputeInputs: layer_names = self.output_layer_names or output.getAllLayerNames() self._logger.debug(f"Processing input with layers: {layer_names}") + strides = None kpts_outputs = None masks_outputs_values = None protos_output = None @@ -543,6 +545,7 @@ def extract(self, output: dai.NNData) -> YOLOComputeInputs: subtype=self.subtype, layer_names=list(layer_names), outputs_values=outputs_values, + strides=strides, conf_threshold=self.conf_threshold, n_classes=self.n_classes, iou_threshold=self.iou_threshold, @@ -569,6 +572,7 @@ def compute(inputs: YOLOComputeInputs) -> dict[str, Any]: subtype=inputs.subtype, layer_names=inputs.layer_names, outputs_values=inputs.outputs_values, + strides=inputs.strides, conf_threshold=inputs.conf_threshold, n_classes=inputs.n_classes, iou_threshold=inputs.iou_threshold, From 977a34651a97d08199265d79f85c25244fb6d80b Mon Sep 17 00:00:00 2001 From: Dani Rogmans Date: Wed, 8 Jul 2026 12:18:01 +0200 Subject: [PATCH 10/11] smart merge of 304 into 302 --- .github/workflows/e2e_tests.yaml | 9 +- .github/workflows/stability_tests.yaml | 4 +- depthai_nodes/message/README.md | 10 -- depthai_nodes/message/__init__.py | 2 - .../message/creators/segmentation.py | 13 +- depthai_nodes/message/segmentation.py | 144 ------------------ depthai_nodes/message/utils/copy_message.py | 17 ++- depthai_nodes/node/README.md | 4 +- depthai_nodes/node/apply_colormap.py | 11 +- depthai_nodes/node/parser_generator.py | 16 +- depthai_nodes/node/parsers/fastsam.py | 4 +- depthai_nodes/node/parsers/segmentation.py | 53 ++++++- depthai_nodes/node/parsers/utils/fastsam.py | 2 +- .../node/parsers/utils/segmentation.py | 9 +- depthai_nodes/node/parsers/yolo.py | 18 +++ depthai_nodes/node/utils/message_remapping.py | 22 ++- depthai_nodes/node/utils/util_constants.py | 3 +- requirements.txt | 2 +- tests/stability_tests/check_messages.py | 9 +- tests/stability_tests/run_parser_test.py | 5 +- .../test_creators/test_segmentation.py | 23 +-- .../test_messages/test_copy_message.py | 29 ++++ .../test_messages/test_segmentation_msg.py | 49 ------ .../test_apply_colormap_node.py | 11 +- tests/utils/messages/creators/arrays.py | 2 +- tests/utils/nodes/mocks/pipeline.py | 3 + 26 files changed, 206 insertions(+), 268 deletions(-) delete mode 100644 depthai_nodes/message/segmentation.py delete mode 100644 tests/unittests/test_messages/test_segmentation_msg.py diff --git a/.github/workflows/e2e_tests.yaml b/.github/workflows/e2e_tests.yaml index c3a34fc6..98513002 100644 --- a/.github/workflows/e2e_tests.yaml +++ b/.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 diff --git a/.github/workflows/stability_tests.yaml b/.github/workflows/stability_tests.yaml index 5d16640d..6ea63ba1 100644 --- a/.github/workflows/stability_tests.yaml +++ b/.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 diff --git a/depthai_nodes/message/README.md b/depthai_nodes/message/README.md index e645a9b8..ccf1bdfa 100644 --- a/depthai_nodes/message/README.md +++ b/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. diff --git a/depthai_nodes/message/__init__.py b/depthai_nodes/message/__init__.py index 6a1d15a4..a14cf509 100644 --- a/depthai_nodes/message/__init__.py +++ b/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", diff --git a/depthai_nodes/message/creators/segmentation.py b/depthai_nodes/message/creators/segmentation.py index 164db8b8..9a9132d7 100644 --- a/depthai_nodes/message/creators/segmentation.py +++ b/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 diff --git a/depthai_nodes/message/segmentation.py b/depthai_nodes/message/segmentation.py deleted file mode 100644 index bc47f298..00000000 --- a/depthai_nodes/message/segmentation.py +++ /dev/null @@ -1,144 +0,0 @@ -import copy - -import cv2 -import depthai as dai -import numpy as np -from numpy.typing import NDArray - - -class SegmentationMask(dai.Buffer): - """SegmentationMask class for a single- or multi-object segmentation mask. - Unassigned pixels are represented with "-1" and foreground classes with non-negative - integers. - - Attributes - ---------- - mask: NDArray[np.int16] - Segmentation mask. - transformation : dai.ImgTransformation - Image transformation object. - """ - - def __init__(self): - """Initializes the SegmentationMask object.""" - super().__init__() - self._mask: NDArray[np.int16] = np.empty(0, dtype=np.int16) - self._transformation: dai.ImgTransformation | None = None - - def copy(self): - """Creates a new instance of the SegmentationMask class and copies the - attributes. - - @return: A new instance of the SegmentationMask class. - @rtype: SegmentationMask - """ - new_obj = SegmentationMask() - new_obj.mask = copy.deepcopy(self._mask) - new_obj.setSequenceNum(self.getSequenceNum()) - new_obj.setTimestamp(self.getTimestamp()) - new_obj.setTimestampDevice(self.getTimestampDevice()) - new_obj.setTransformation(self.getTransformation()) - return new_obj - - @property - def mask(self) -> NDArray[np.int16]: - """Returns the segmentation mask. - - @return: Segmentation mask. - @rtype: NDArray[np.int16] - """ - return self._mask - - @mask.setter - def mask(self, value: NDArray[np.int16]): - """Sets the segmentation mask. - - @param value: Segmentation mask. - @type value: NDArray[np.int16]) - @raise TypeError: If value is not a numpy array. - @raise ValueError: If value is not a 2D numpy array. - @raise ValueError: If each element is not of type int16. - @raise ValueError: If any element is smaller than -1. - """ - if not isinstance(value, np.ndarray): - raise TypeError("Mask must be a numpy array.") - if not (value.size == 0 or value.ndim == 2): - raise ValueError("Mask must be 2D or empty.") - if value.dtype != np.int16: - raise ValueError("Mask must be an array of int16.") - if np.any(value < -1): - raise ValueError("Mask must be an array of integers larger or equal to -1.") - self._mask = value - - @property - def transformation(self) -> dai.ImgTransformation | None: - """Returns the Image Transformation object. - - @return: The Image Transformation object. - @rtype: dai.ImgTransformation - """ - return self._transformation - - @transformation.setter - def transformation(self, value: dai.ImgTransformation | None): - """Sets the Image Transformation object. - - @param value: The Image Transformation object. - @type value: dai.ImgTransformation - @raise TypeError: If value is not a dai.ImgTransformation object. - """ - - if value is not None: - if not isinstance(value, dai.ImgTransformation): - raise TypeError( - f"Transformation must be a dai.ImgTransformation object, instead got {type(value)}." - ) - self._transformation = value - - def setTransformation(self, transformation: dai.ImgTransformation | None): - """Sets the Image Transformation object. - - @param transformation: The Image Transformation object. - @type transformation: dai.ImgTransformation - @raise TypeError: If value is not a dai.ImgTransformation object. - """ - self.transformation = transformation - - def getTransformation(self) -> dai.ImgTransformation | None: - """Returns the Image Transformation object. - - @return: The Image Transformation object. - @rtype: dai.ImgTransformation - """ - return self.transformation - - def getVisualizationMessage(self) -> dai.ImgFrame: - """Returns the default visualization message for segmentation masks.""" - img_frame = dai.ImgFrame() - img_frame.setTimestamp(self.getTimestamp()) - img_frame.setTimestampDevice(self.getTimestampDevice()) - img_frame.setSequenceNum(self.getSequenceNum()) - if self.transformation is not None: - img_frame.setTransformation(self.transformation) - mask = self._mask.copy() - - unique_values = np.unique(mask[mask >= 0]) - scaled_mask = np.zeros_like(mask, dtype=np.uint8) - - if unique_values.size == 0: - empty_bgr_mask = cv2.cvtColor(scaled_mask, cv2.COLOR_GRAY2BGR) - return img_frame.setCvFrame(empty_bgr_mask, dai.ImgFrame.Type.BGR888i) - - min_val, max_val = unique_values.min(), unique_values.max() - - if min_val == max_val: - scaled_mask = np.ones_like(mask, dtype=np.uint8) * 255 - else: - scaled_mask = ((mask - min_val) / (max_val - min_val) * 255).astype( - np.uint8 - ) - scaled_mask[mask == -1] = 0 - colored_mask = cv2.applyColorMap(scaled_mask, cv2.COLORMAP_RAINBOW) - colored_mask[mask == -1] = [0, 0, 0] - - return img_frame.setCvFrame(colored_mask, dai.ImgFrame.Type.BGR888i) diff --git a/depthai_nodes/message/utils/copy_message.py b/depthai_nodes/message/utils/copy_message.py index 0fd55d24..f78a6be9 100644 --- a/depthai_nodes/message/utils/copy_message.py +++ b/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) diff --git a/depthai_nodes/node/README.md b/depthai_nodes/node/README.md index f5e68695..50af6f01 100644 --- a/depthai_nodes/node/README.md +++ b/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 diff --git a/depthai_nodes/node/apply_colormap.py b/depthai_nodes/node/apply_colormap.py index ce0058fa..06b60f44 100644 --- a/depthai_nodes/node/apply_colormap.py +++ b/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: diff --git a/depthai_nodes/node/parser_generator.py b/depthai_nodes/node/parser_generator.py index 4ab75295..8a718adc 100644 --- a/depthai_nodes/node/parser_generator.py +++ b/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: diff --git a/depthai_nodes/node/parsers/fastsam.py b/depthai_nodes/node/parsers/fastsam.py index 409f3e98..ac6d43aa 100644 --- a/depthai_nodes/node/parsers/fastsam.py +++ b/depthai_nodes/node/parsers/fastsam.py @@ -41,9 +41,9 @@ class FastSAMParser(BaseParser): Output Message/s ---------------- - **Type**: SegmentationMask + **Type**: dai.SegmentationMask - **Description**: SegmentationMask message containing the resulting segmentation masks given the prompt. + **Description**: dai.SegmentationMask message containing the resulting segmentation masks given the prompt. Error Handling -------------- diff --git a/depthai_nodes/node/parsers/segmentation.py b/depthai_nodes/node/parsers/segmentation.py index 4b4080e9..527e7671 100644 --- a/depthai_nodes/node/parsers/segmentation.py +++ b/depthai_nodes/node/parsers/segmentation.py @@ -22,9 +22,9 @@ class SegmentationParser(BaseParser): Output Message/s ---------------- - **Type**: SegmentationMask + **Type**: dai.SegmentationMask - **Description**: Segmentation message containing the segmentation mask. Every pixel belongs to exactly one class. Unassigned pixels are represented with "-1" and class pixels with non-negative integers. + **Description**: dai.SegmentationMask containing the segmentation mask. Every pixel belongs to exactly one class. Unassigned pixels are represented with "255" and class pixels with non-negative integers. Error Handling -------------- @@ -34,7 +34,10 @@ class SegmentationParser(BaseParser): """ def __init__( - self, output_layer_name: str = "", classes_in_one_layer: bool = False + self, + output_layer_name: str = "", + classes_in_one_layer: bool = False, + background_class: bool = False, ) -> None: """Initializes the parser node. @@ -44,12 +47,19 @@ def __init__( class segmentation model. Default is False. If True, the parser will use np.max instead of np.argmax to get the class map. @type classes_in_one_layer: bool + @param background_class: Whether class index 0 should be treated as background. + @type background_class: bool """ super().__init__() self.output_layer_name = output_layer_name self.classes_in_one_layer = classes_in_one_layer + self.background_class = background_class + self._background_class_ignored_warning_sent = False self._logger.debug( - f"SegmentationParser initialized with output_layer_name='{output_layer_name}', classes_in_one_layer={classes_in_one_layer}" + "SegmentationParser initialized with " + f"output_layer_name='{output_layer_name}', " + f"classes_in_one_layer={classes_in_one_layer}, " + f"background_class={background_class}" ) def setOutputLayerName(self, output_layer_name: str) -> None: @@ -74,6 +84,28 @@ def setClassesInOneLayer(self, classes_in_one_layer: bool) -> None: self.classes_in_one_layer = classes_in_one_layer self._logger.debug(f"Classes in one layer set to {self.classes_in_one_layer}") + def setBackgroundClass(self, background_class: bool) -> None: + """Sets whether class index 0 should be treated as background. + + @param background_class: Whether class index 0 is background. + @type background_class: bool + """ + if not isinstance(background_class, bool): + raise ValueError("background_class must be a boolean.") + self.background_class = background_class + self._logger.debug(f"Background class set to {self.background_class}") + + def _warn_if_background_class_ignored(self) -> None: + if ( + self.background_class + and self.classes_in_one_layer + and not self._background_class_ignored_warning_sent + ): + self._logger.warning( + "background_class=True is ignored when classes_in_one_layer=True." + ) + self._background_class_ignored_warning_sent = True + def build( self, head_config: dict[str, Any], @@ -95,11 +127,19 @@ def build( self.classes_in_one_layer = head_config.get( "classes_in_one_layer", self.classes_in_one_layer ) + self.background_class = head_config.get( + "background_class", self.background_class + ) self._logger.debug( - f"SegmentationParser built with output_layer_name='{self.output_layer_name}', classes_in_one_layer={self.classes_in_one_layer}" + "SegmentationParser built with " + f"output_layer_name='{self.output_layer_name}', " + f"classes_in_one_layer={self.classes_in_one_layer}, " + f"background_class={self.background_class}" ) + self._warn_if_background_class_ignored() + return self def run(self): @@ -114,6 +154,7 @@ def run(self): class_map = self.compute( segmentation_mask, classes_in_one_layer=self.classes_in_one_layer, + background_class=self.background_class, ) self.emit(output, class_map) @@ -134,10 +175,12 @@ def compute( segmentation_mask: np.ndarray, *, classes_in_one_layer: bool = False, + background_class: bool = False, ) -> np.ndarray: return compute_segmentation_class_map( segmentation_mask, classes_in_one_layer=classes_in_one_layer, + background_class=background_class, ) def emit(self, output: dai.NNData, class_map: np.ndarray) -> None: diff --git a/depthai_nodes/node/parsers/utils/fastsam.py b/depthai_nodes/node/parsers/utils/fastsam.py index 20162920..88d68d5e 100644 --- a/depthai_nodes/node/parsers/utils/fastsam.py +++ b/depthai_nodes/node/parsers/utils/fastsam.py @@ -386,7 +386,7 @@ def merge_masks(masks: np.ndarray) -> np.ndarray: else: raise ValueError("Masks must be a 3D array.") - merged_masks = np.full((height, width), -1, dtype=np.int16) + merged_masks = np.full((height, width), 255, dtype=np.uint8) for i in range(n): merged_masks[masks[i] > 0] = i diff --git a/depthai_nodes/node/parsers/utils/segmentation.py b/depthai_nodes/node/parsers/utils/segmentation.py index 7603d140..b60b3a8d 100644 --- a/depthai_nodes/node/parsers/utils/segmentation.py +++ b/depthai_nodes/node/parsers/utils/segmentation.py @@ -5,6 +5,7 @@ def compute_segmentation_class_map( segmentation_mask: np.ndarray, *, classes_in_one_layer: bool = False, + background_class: bool = False, ) -> np.ndarray: """Convert segmentation logits into a class map.""" mask = np.asarray(segmentation_mask) @@ -41,4 +42,10 @@ def compute_segmentation_class_map( if adding_unassigned_class: class_map = class_map - 1 - return class_map + if background_class and not classes_in_one_layer: + class_map = np.where(class_map == 0, 255, class_map) + + if np.any((class_map < 0) | (class_map > 255)): + raise ValueError("Segmentation class map values must be in the range [0, 255].") + + return class_map.astype(np.uint8) diff --git a/depthai_nodes/node/parsers/yolo.py b/depthai_nodes/node/parsers/yolo.py index 99308818..53bff573 100644 --- a/depthai_nodes/node/parsers/yolo.py +++ b/depthai_nodes/node/parsers/yolo.py @@ -254,6 +254,24 @@ def setAnchors(self, anchors: list[list[list[float]]]) -> None: self.anchors = anchors self._logger.debug(f"Anchors set to {self.anchors}") + def setStrides(self, strides: list[int] | tuple[int, ...]) -> None: + """Sets the strides for YOLO output heads. + + @param strides: Strides for YOLO output heads. + @type strides: list[int] | tuple[int, ...] + """ + if not isinstance(strides, (list, tuple)): + raise ValueError("Strides must be a list or tuple.") + if not strides: + raise ValueError("Strides must not be empty.") + if not all(isinstance(stride, int) for stride in strides): + raise ValueError("Strides must be a list or tuple of integers.") + if not all(stride > 0 for stride in strides): + raise ValueError("Strides must be positive.") + + self.strides = strides + self._logger.debug(f"Strides set to {self.strides}") + def setSubtype(self, subtype: str) -> None: """Sets the subtype of the YOLO model. diff --git a/depthai_nodes/node/utils/message_remapping.py b/depthai_nodes/node/utils/message_remapping.py index 1662d46a..a694264f 100644 --- a/depthai_nodes/node/utils/message_remapping.py +++ b/depthai_nodes/node/utils/message_remapping.py @@ -8,7 +8,6 @@ from depthai_nodes.message.lines import Line, Lines from depthai_nodes.message.map import Map2D from depthai_nodes.message.prediction import Prediction, Predictions -from depthai_nodes.message.segmentation import SegmentationMask from depthai_nodes.node.utils.util_constants import UNASSIGNED_MASK_LABEL, GMessage @@ -21,7 +20,7 @@ def remap_message( return remap_img_detections(from_transformation, to_transformation, message) elif isinstance(message, Keypoints): return remap_keypoints(from_transformation, to_transformation, message) - elif isinstance(message, SegmentationMask): + elif isinstance(message, dai.SegmentationMask): return remap_segmentation_mask(from_transformation, to_transformation, message) elif isinstance(message, Clusters): return remap_clusters(from_transformation, to_transformation, message) @@ -68,13 +67,18 @@ def remap_segmentation_mask_array( dst_matrix = np.array(to_transformation.getMatrix()) src_matrix = np.array(from_transformation.getMatrixInv()) trans_matrix = dst_matrix @ src_matrix + border_value = ( + 255 + if np.issubdtype(segmentation_mask.dtype, np.uint8) + else UNASSIGNED_MASK_LABEL + ) new_mask = cv2.warpPerspective( segmentation_mask, trans_matrix, to_transformation.getSize(), flags=cv2.INTER_NEAREST, borderMode=cv2.BORDER_CONSTANT, - borderValue=UNASSIGNED_MASK_LABEL, # type: ignore + borderValue=border_value, # type: ignore ) return new_mask @@ -82,11 +86,13 @@ def remap_segmentation_mask_array( def remap_segmentation_mask( from_transformation: dai.ImgTransformation, to_transformation: dai.ImgTransformation, - segmentation_mask: SegmentationMask, -) -> SegmentationMask: - new_mask = SegmentationMask() - new_mask.mask = remap_segmentation_mask_array( - from_transformation, to_transformation, segmentation_mask.mask + segmentation_mask: dai.SegmentationMask, +) -> dai.SegmentationMask: + new_mask = dai.SegmentationMask() + new_mask.setCvMask( + remap_segmentation_mask_array( + from_transformation, to_transformation, segmentation_mask.getCvMask() + ) ) return new_mask diff --git a/depthai_nodes/node/utils/util_constants.py b/depthai_nodes/node/utils/util_constants.py index ed14a26e..ed1025c7 100644 --- a/depthai_nodes/node/utils/util_constants.py +++ b/depthai_nodes/node/utils/util_constants.py @@ -8,13 +8,12 @@ from depthai_nodes.message.lines import Lines from depthai_nodes.message.map import Map2D from depthai_nodes.message.prediction import Predictions -from depthai_nodes.message.segmentation import SegmentationMask GMessage = TypeVar( "GMessage", bound=dai.ImgDetections | Keypoints - | SegmentationMask + | dai.SegmentationMask | Clusters | Map2D | Lines diff --git a/requirements.txt b/requirements.txt index 29117abd..55d7c9bc 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,3 @@ -depthai>=3.6.1 +depthai>=3.7.1 opencv-python-headless~=4.10.0 numpy>=1.22 diff --git a/tests/stability_tests/check_messages.py b/tests/stability_tests/check_messages.py index b78c0a84..02c21f2a 100644 --- a/tests/stability_tests/check_messages.py +++ b/tests/stability_tests/check_messages.py @@ -11,7 +11,6 @@ Lines, Map2D, Predictions, - SegmentationMask, ) from .utils import extract_main_slug @@ -107,7 +106,7 @@ def check_embeddings_msg( def check_segmentation_msg( - message: SegmentationMask, + message: dai.SegmentationMask, expected_output: dict[str, Any], threshold: float = 0.9, verbose: bool = False, @@ -121,10 +120,10 @@ def check_segmentation_msg( } """ assert isinstance( - message, SegmentationMask - ), f"The message is not a SegmentationMask. Got {type(message)}." + message, dai.SegmentationMask + ), f"The message is not a dai.SegmentationMask. Got {type(message)}." - mask = message.mask + mask = message.getCvMask() expected_mask = expected_output["mask"] if verbose: diff --git a/tests/stability_tests/run_parser_test.py b/tests/stability_tests/run_parser_test.py index 5725739e..8d5e3f37 100644 --- a/tests/stability_tests/run_parser_test.py +++ b/tests/stability_tests/run_parser_test.py @@ -90,7 +90,10 @@ def test_parser(parser_generator, model: str, parser_name: str, duration: int): exit(9) # Create and build parser - parser = parser_generator.build(nnArchive=nn_archive)[0] + parser = parser_generator.build( + nnArchive=nn_archive, + hostOnly=parser_name == "SegmentationParser", + )[0] parser.input._queue.duration = duration # Load and send test data diff --git a/tests/unittests/test_creators/test_segmentation.py b/tests/unittests/test_creators/test_segmentation.py index 21b42788..0d35c5b9 100644 --- a/tests/unittests/test_creators/test_segmentation.py +++ b/tests/unittests/test_creators/test_segmentation.py @@ -1,18 +1,19 @@ +import depthai as dai import numpy as np import pytest -from depthai_nodes import SegmentationMask from depthai_nodes.message.creators import create_segmentation_message -MASK = np.random.randint(0, 256, (480, 640), dtype=np.int16) +MASK = np.random.randint(0, 256, (480, 640), dtype=np.uint8) def test_valid_input(): message = create_segmentation_message(MASK) + mask = message.getCvMask() - assert isinstance(message, SegmentationMask) - assert message.mask.shape == (480, 640) - assert message.mask.dtype == np.int16 + assert isinstance(message, dai.SegmentationMask) + assert mask.shape == (480, 640) + assert mask.dtype == np.uint8 def test_invalid_type(): @@ -21,20 +22,20 @@ def test_invalid_type(): def test_invalid_shape(): - mask = np.random.randint(0, 256, (480, 640, 3), dtype=np.int16) + mask = np.random.randint(0, 256, (480, 640, 3), dtype=np.uint8) with pytest.raises(ValueError): create_segmentation_message(mask) def test_invalid_dtype(): with pytest.raises(ValueError): - create_segmentation_message(MASK.astype(np.uint8)) + create_segmentation_message(MASK.astype(np.int16)) def test_empty_mask(): - mask = np.empty((0, 0), dtype=np.int16) + mask = np.empty((0, 0), dtype=np.uint8) message = create_segmentation_message(mask) + mask = message.getCvMask() - assert isinstance(message, SegmentationMask) - assert message.mask.shape == (0, 0) - assert message.mask.dtype == np.int16 + assert isinstance(message, dai.SegmentationMask) + assert mask is None diff --git a/tests/unittests/test_messages/test_copy_message.py b/tests/unittests/test_messages/test_copy_message.py index 93f6da0b..c5d2b9b3 100644 --- a/tests/unittests/test_messages/test_copy_message.py +++ b/tests/unittests/test_messages/test_copy_message.py @@ -1,6 +1,7 @@ import inspect from collections.abc import Callable +import depthai as dai import numpy as np import pytest @@ -67,3 +68,31 @@ def test_message_copying(message_creator: tuple[str, Callable]): except TypeError: # copying not implemented for all messages pass + + +def test_copy_native_segmentation_mask(): + mask = np.array( + [ + [0, 1, 255], + [2, 3, 255], + ], + dtype=np.uint8, + ) + + msg = dai.SegmentationMask() + msg.setCvMask(mask) + msg.setSequenceNum(123) + + if hasattr(msg, "setLabels"): + msg.setLabels(["class_0", "class_1", "class_2", "class_3"]) + + msg_copy = copy_message(msg) + + assert isinstance(msg_copy, dai.SegmentationMask) + assert msg_copy is not msg + assert np.array_equal(msg_copy.getCvMask(), msg.getCvMask()) + assert msg_copy.getCvMask().dtype == np.uint8 + assert msg_copy.getSequenceNum() == msg.getSequenceNum() + + if hasattr(msg, "getLabels"): + assert msg_copy.getLabels() == msg.getLabels() diff --git a/tests/unittests/test_messages/test_segmentation_msg.py b/tests/unittests/test_messages/test_segmentation_msg.py deleted file mode 100644 index 1495f6f9..00000000 --- a/tests/unittests/test_messages/test_segmentation_msg.py +++ /dev/null @@ -1,49 +0,0 @@ -import depthai as dai -import numpy as np -import pytest - -from depthai_nodes import SegmentationMask - - -@pytest.fixture -def segmentation_mask(): - return SegmentationMask() - - -def test_segmentation_mask_initialization(segmentation_mask: SegmentationMask): - assert np.array_equal(segmentation_mask.mask, np.array([])) - assert segmentation_mask.transformation is None - - -def test_segmentation_mask_set_mask(segmentation_mask: SegmentationMask): - mask_array = np.random.randint(-1, 256, (480, 640), dtype=np.int16) - segmentation_mask.mask = mask_array - assert np.array_equal(segmentation_mask.mask, mask_array) - - with pytest.raises(TypeError): - segmentation_mask.mask = "not a numpy array" - - with pytest.raises(ValueError): - segmentation_mask.mask = np.random.randint( - -1, 256, (480, 640, 3), dtype=np.int16 - ) - - with pytest.raises(ValueError): - segmentation_mask.mask = np.random.randint(-1, 256, (480, 640), dtype=np.uint8) - - with pytest.raises(ValueError): - segmentation_mask.mask = np.random.randint(-2, 256, (480, 640), dtype=np.int16) - - -def test_segmentation_mask_set_transformation(segmentation_mask: SegmentationMask): - transformation = dai.ImgTransformation() - segmentation_mask.transformation = transformation - assert segmentation_mask.transformation == transformation - - with pytest.raises(TypeError): - segmentation_mask.transformation = "not a dai.ImgTransformation" - - -def test_segmentation_mask_set_transformation_none(segmentation_mask: SegmentationMask): - segmentation_mask.transformation = None - assert segmentation_mask.transformation is None diff --git a/tests/unittests/test_nodes/test_host_nodes/test_apply_colormap_node.py b/tests/unittests/test_nodes/test_host_nodes/test_apply_colormap_node.py index 359e1158..1111e61a 100644 --- a/tests/unittests/test_nodes/test_host_nodes/test_apply_colormap_node.py +++ b/tests/unittests/test_nodes/test_host_nodes/test_apply_colormap_node.py @@ -22,12 +22,21 @@ # Your parameter lists colormap_values = [cv2.COLORMAP_HOT, cv2.COLORMAP_PLASMA, cv2.COLORMAP_INFERNO] + + +def _create_segmentation_mask(arr: np.ndarray) -> dai.SegmentationMask: + msg = dai.SegmentationMask() + msg.setCvMask(arr.astype(np.uint8)) + return msg + + arr_creators = [ lambda: create_img_frame( image=ARR[..., np.newaxis], img_frame_type=dai.ImgFrame.Type.RAW8 ), lambda: create_map(ARR.astype(np.float32)), lambda: _create_img_detections_with_mask(ARR), + lambda: _create_segmentation_mask(ARR), ] @@ -137,7 +146,7 @@ def test_processing( modified_arr = ( np.where(ARR == 255, 0, ARR + 1) - if isinstance(arr, dai.ImgDetections) + if isinstance(arr, (dai.ImgDetections, dai.SegmentationMask)) else ARR ) assert np.array_equal( diff --git a/tests/utils/messages/creators/arrays.py b/tests/utils/messages/creators/arrays.py index 00003b06..771eb553 100644 --- a/tests/utils/messages/creators/arrays.py +++ b/tests/utils/messages/creators/arrays.py @@ -20,4 +20,4 @@ def create_map(map: np.ndarray = ARRAYS["2d"]): def create_segmentation(mask: np.ndarray = ARRAYS["2d"]): - return creators.create_segmentation_message(mask=mask.astype(np.int16)) + return creators.create_segmentation_message(mask=mask.astype(np.uint8)) diff --git a/tests/utils/nodes/mocks/pipeline.py b/tests/utils/nodes/mocks/pipeline.py index 49fe8729..13bf6a0d 100644 --- a/tests/utils/nodes/mocks/pipeline.py +++ b/tests/utils/nodes/mocks/pipeline.py @@ -187,6 +187,9 @@ def build(self): def setNNArchive(self, nn_archive): self._nn_archive = nn_archive + def setNNArchiveHead(self, head) -> None: + self._nn_archive_head = head + def setRunOnHost(self, run_on_host: bool): self._run_on_host = run_on_host From 26b22c6b27cddae160be1c13d4b8ccb777ef87fd Mon Sep 17 00:00:00 2001 From: Dani Rogmans Date: Wed, 8 Jul 2026 12:36:48 +0200 Subject: [PATCH 11/11] fix incorrect shift from new semantic segmentation changes --- depthai_nodes/node/parsers/utils/segmentation.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/depthai_nodes/node/parsers/utils/segmentation.py b/depthai_nodes/node/parsers/utils/segmentation.py index b60b3a8d..ddabf672 100644 --- a/depthai_nodes/node/parsers/utils/segmentation.py +++ b/depthai_nodes/node/parsers/utils/segmentation.py @@ -41,8 +41,9 @@ def compute_segmentation_class_map( if adding_unassigned_class: class_map = class_map - 1 + class_map = np.where(class_map == -1, 255, class_map) - if background_class and not classes_in_one_layer: + if background_class and not classes_in_one_layer and not adding_unassigned_class: class_map = np.where(class_map == 0, 255, class_map) if np.any((class_map < 0) | (class_map > 255)):