diff --git a/depthai_nodes/ml/helpers/constants.py b/depthai_nodes/ml/helpers/constants.py new file mode 100644 index 00000000..39ccbcac --- /dev/null +++ b/depthai_nodes/ml/helpers/constants.py @@ -0,0 +1,6 @@ +import depthai as dai + +OUTLINE_COLOR = dai.Color(1.0, 0.5, 0.5, 1.0) +TEXT_COLOR = dai.Color(0.5, 0.5, 1.0, 1.0) +BACKGROUND_COLOR = dai.Color(1.0, 1.0, 0.5, 1.0) +KEYPOINT_COLOR = dai.Color(1.0, 0.35, 0.367, 1.0) diff --git a/depthai_nodes/ml/messages/classification.py b/depthai_nodes/ml/messages/classification.py index 911889a1..67483c0d 100644 --- a/depthai_nodes/ml/messages/classification.py +++ b/depthai_nodes/ml/messages/classification.py @@ -4,6 +4,11 @@ import numpy as np from numpy.typing import NDArray +from depthai_nodes.ml.helpers.constants import ( + BACKGROUND_COLOR, + TEXT_COLOR, +) + class Classifications(dai.Buffer): """Classification class for storing the classes and their respective scores. @@ -119,3 +124,25 @@ def transformation(self, value: dai.ImgTransformation): f"Transformation must be a dai.ImgTransformation object, instead got {type(value)}." ) self._transformation = value + + def getVisualizationMessage(self) -> dai.ImgAnnotations: + """Returns default visualization message for classification. + + The message adds the top five classes and their scores to the right side of the + image. + """ + img_annotations = dai.ImgAnnotations() + annotation = dai.ImgAnnotation() + + for i in range(len(self._scores)): + text = dai.TextAnnotation() + text.position = dai.Point2f(1.05, 0.1 + i * 0.1) + text.text = f"{self._classes[i]} {self._scores[i] * 100:.0f}%" + text.fontSize = 15 + text.textColor = TEXT_COLOR + text.backgroundColor = BACKGROUND_COLOR + annotation.texts.append(text) + + img_annotations.annotations.append(annotation) + img_annotations.setTimestamp(self.getTimestamp()) + return img_annotations diff --git a/depthai_nodes/ml/messages/clusters.py b/depthai_nodes/ml/messages/clusters.py index 145d5abb..3950cf69 100644 --- a/depthai_nodes/ml/messages/clusters.py +++ b/depthai_nodes/ml/messages/clusters.py @@ -1,6 +1,8 @@ from typing import List +import cv2 import depthai as dai +import numpy as np class Cluster(dai.Buffer): @@ -18,7 +20,7 @@ def __init__(self): """Initializes the Cluster object.""" super().__init__() self._label: int = None - self.points: List[dai.Point2f] = [] + self._points: List[dai.Point2f] = [] @property def label(self) -> int: @@ -131,3 +133,30 @@ def transformation(self, value: dai.ImgTransformation): f"Transformation must be a dai.ImgTransformation object, instead got {type(value)}." ) self._transformation = value + + def getVisualizationMessage(self) -> dai.ImgAnnotations: + """Creates a default visualization message for clusters and colors each one + separately.""" + img_annotations = dai.ImgAnnotations() + annotation = dai.ImgAnnotation() + + num_clusters = len(self.clusters) + color_mask = np.array(range(0, 255, 255 // num_clusters), dtype=np.uint8) + color_mask = cv2.applyColorMap(color_mask, cv2.COLORMAP_RAINBOW) + color_mask = color_mask / 255 + color_mask = color_mask.reshape(-1, 3) + + for i, cluster in enumerate(self.clusters): + pointsAnnotation = dai.PointsAnnotation() + pointsAnnotation.type = dai.PointsAnnotationType.POINTS + pointsAnnotation.points = dai.VectorPoint2f(cluster.points) + r, g, b = color_mask[i] + color = dai.Color(r, g, b) + pointsAnnotation.outlineColor = color + pointsAnnotation.fillColor = color + pointsAnnotation.thickness = 2.0 + annotation.points.append(pointsAnnotation) + + img_annotations.annotations.append(annotation) + img_annotations.setTimestamp(self.getTimestamp()) + return img_annotations diff --git a/depthai_nodes/ml/messages/img_detections.py b/depthai_nodes/ml/messages/img_detections.py index b0443698..0fd418af 100644 --- a/depthai_nodes/ml/messages/img_detections.py +++ b/depthai_nodes/ml/messages/img_detections.py @@ -4,6 +4,13 @@ import numpy as np from numpy.typing import NDArray +from depthai_nodes.ml.helpers.constants import ( + BACKGROUND_COLOR, + KEYPOINT_COLOR, + OUTLINE_COLOR, + TEXT_COLOR, +) + from .keypoints import Keypoint from .segmentation import SegmentationMask @@ -199,7 +206,7 @@ def detections(self, value: List[ImgDetectionExtended]): self._detections = value @property - def masks(self) -> NDArray[np.int8]: + def masks(self) -> NDArray[np.int16]: """Returns the segmentation masks stored in a single numpy array. @return: Segmentation masks. @@ -253,3 +260,49 @@ def transformation(self, value: dai.ImgTransformation): f"Transformation must be a dai.ImgTransformation object, instead got {type(value)}." ) self._transformation = value + + def getVisualizationMessage(self) -> dai.ImgAnnotations: + img_annotations = dai.ImgAnnotations() + annotation = dai.ImgAnnotation() + transformation = self.transformation + w, h = 1, 1 + if transformation is not None: # remove once RVC2 supports ImgTransformation + w, h = transformation.getSize() + + for detection in self.detections: + detection: ImgDetectionExtended = detection + rotated_rect = detection.rotated_rect + rotated_rect = rotated_rect.denormalize(w, h) + points = rotated_rect.getPoints() + points = [dai.Point2f(point.x / w, point.y / h) for point in points] + pointsAnnotation = dai.PointsAnnotation() + pointsAnnotation.type = dai.PointsAnnotationType.LINE_LOOP + pointsAnnotation.points = dai.VectorPoint2f(points) + pointsAnnotation.outlineColor = OUTLINE_COLOR + pointsAnnotation.thickness = 2.0 + annotation.points.append(pointsAnnotation) + + text = dai.TextAnnotation() + text.position = points[0] + text.text = f"{detection.label_name} {int(detection.confidence * 100)}%" + text.fontSize = 15 + text.textColor = TEXT_COLOR + text.backgroundColor = BACKGROUND_COLOR + annotation.texts.append(text) + + if len(detection.keypoints) > 0: + keypoints = [ + dai.Point2f(keypoint.x, keypoint.y) + for keypoint in detection.keypoints + ] + keypointAnnotation = dai.PointsAnnotation() + keypointAnnotation.type = dai.PointsAnnotationType.POINTS + keypointAnnotation.points = dai.VectorPoint2f(keypoints) + keypointAnnotation.outlineColor = KEYPOINT_COLOR + keypointAnnotation.fillColor = KEYPOINT_COLOR + keypointAnnotation.thickness = 2 + annotation.points.append(keypointAnnotation) + + img_annotations.annotations.append(annotation) + img_annotations.setTimestamp(self.getTimestamp()) + return img_annotations diff --git a/depthai_nodes/ml/messages/keypoints.py b/depthai_nodes/ml/messages/keypoints.py index 27c32d22..7eca21b9 100644 --- a/depthai_nodes/ml/messages/keypoints.py +++ b/depthai_nodes/ml/messages/keypoints.py @@ -2,6 +2,8 @@ import depthai as dai +from depthai_nodes.ml.helpers.constants import KEYPOINT_COLOR + class Keypoint(dai.Buffer): """Keypoint class for storing a keypoint. @@ -185,3 +187,20 @@ def transformation(self, value: dai.ImgTransformation): f"Transformation must be a dai.ImgTransformation object, instead got {type(value)}." ) self._transformation = value + + def getVisualizationMessage(self) -> dai.ImgAnnotations: + """Creates a default visualization message for the keypoints.""" + img_annotations = dai.ImgAnnotations() + annotation = dai.ImgAnnotation() + keypoints = [dai.Point2f(keypoint.x, keypoint.y) for keypoint in self.keypoints] + pointsAnnotation = dai.PointsAnnotation() + pointsAnnotation.type = dai.PointsAnnotationType.POINTS + pointsAnnotation.points = dai.VectorPoint2f(keypoints) + pointsAnnotation.outlineColor = KEYPOINT_COLOR + pointsAnnotation.fillColor = KEYPOINT_COLOR + pointsAnnotation.thickness = 2 + annotation.points.append(pointsAnnotation) + + img_annotations.annotations.append(annotation) + img_annotations.setTimestamp(self.getTimestamp()) + return img_annotations diff --git a/depthai_nodes/ml/messages/lines.py b/depthai_nodes/ml/messages/lines.py index f9d7015a..6d8d9418 100644 --- a/depthai_nodes/ml/messages/lines.py +++ b/depthai_nodes/ml/messages/lines.py @@ -2,6 +2,8 @@ import depthai as dai +from depthai_nodes.ml.helpers.constants import OUTLINE_COLOR + class Line(dai.Buffer): """Line class for storing a line. @@ -158,3 +160,25 @@ def transformation(self, value: dai.ImgTransformation): f"Transformation must be a dai.ImgTransformation object, instead got {type(value)}." ) self._transformation = value + + def getVisualizationMessage(self) -> dai.ImgAnnotations: + """Returns default visualization message for lines. + + The message adds lines to the image. + """ + img_annotation = dai.ImgAnnotations() + annotation = dai.ImgAnnotation() + + for line in self.lines: + pointsAnnotation = dai.PointsAnnotation() + pointsAnnotation.type = dai.PointsAnnotationType.LINE_STRIP + pointsAnnotation.points = dai.VectorPoint2f( + [line.start_point, line.end_point] + ) + pointsAnnotation.outlineColor = OUTLINE_COLOR + pointsAnnotation.thickness = 2.0 + annotation.points.append(pointsAnnotation) + + img_annotation.annotations.append(annotation) + img_annotation.setTimestamp(self.getTimestamp()) + return img_annotation diff --git a/depthai_nodes/ml/messages/map.py b/depthai_nodes/ml/messages/map.py index 0158961e..191bf209 100644 --- a/depthai_nodes/ml/messages/map.py +++ b/depthai_nodes/ml/messages/map.py @@ -1,3 +1,4 @@ +import cv2 import depthai as dai import numpy as np from numpy.typing import NDArray @@ -92,10 +93,22 @@ def transformation(self, value: dai.ImgTransformation): @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 getVisualizationMessage(self) -> dai.ImgFrame: + """Returns default visualization message for 2D maps in the form of a + colormapped image.""" + img_frame = dai.ImgFrame() + mask = self._map.copy() + if np.any(mask < 1): + mask = mask * 255 + mask = mask.astype(np.uint8) + + colored_mask = cv2.applyColorMap(mask, cv2.COLORMAP_PLASMA) + return img_frame.setCvFrame(colored_mask, dai.ImgFrame.Type.BGR888i) diff --git a/depthai_nodes/ml/messages/prediction.py b/depthai_nodes/ml/messages/prediction.py index 29e0cdb6..ad054938 100644 --- a/depthai_nodes/ml/messages/prediction.py +++ b/depthai_nodes/ml/messages/prediction.py @@ -2,6 +2,11 @@ import depthai as dai +from depthai_nodes.ml.helpers.constants import ( + BACKGROUND_COLOR, + TEXT_COLOR, +) + class Prediction(dai.Buffer): """Prediction class for storing a prediction. @@ -117,3 +122,24 @@ def transformation(self, value: dai.ImgTransformation): f"Transformation must be a dai.ImgTransformation object, instead got {type(value)}." ) self._transformation = value + + def getVisualizationMessage(self) -> dai.ImgAnnotations: + """Returns the visualization message for the predictions. + + The message adds text representing the predictions to the right of the image. + """ + img_annotations = dai.ImgAnnotations() + annotation = dai.ImgAnnotation() + + for i, prediction in enumerate(self.predictions): + text = dai.TextAnnotation() + text.position = dai.Point2f(1.05, 0.1 + i * 0.1) + text.text = f"{prediction.prediction:.2f}" + text.fontSize = 15 + text.textColor = TEXT_COLOR + text.backgroundColor = BACKGROUND_COLOR + annotation.texts.append(text) + + img_annotations.annotations.append(annotation) + img_annotations.setTimestamp(self.getTimestamp()) + return img_annotations diff --git a/depthai_nodes/ml/messages/segmentation.py b/depthai_nodes/ml/messages/segmentation.py index 27952281..8def8da7 100644 --- a/depthai_nodes/ml/messages/segmentation.py +++ b/depthai_nodes/ml/messages/segmentation.py @@ -1,3 +1,4 @@ +import cv2 import depthai as dai import numpy as np from numpy.typing import NDArray @@ -76,3 +77,28 @@ def transformation(self, value: dai.ImgTransformation): f"Transformation must be a dai.ImgTransformation object, instead got {type(value)}." ) self._transformation = value + + def getVisualizationMessage(self) -> dai.ImgFrame: + """Returns the default visualization message for segmentation masks.""" + img_frame = dai.ImgFrame() + 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: + return img_frame.setCvFrame(scaled_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/ml/parsers/mediapipe_palm_detection.py b/depthai_nodes/ml/parsers/mediapipe_palm_detection.py index 5b22905b..deddf2bf 100644 --- a/depthai_nodes/ml/parsers/mediapipe_palm_detection.py +++ b/depthai_nodes/ml/parsers/mediapipe_palm_detection.py @@ -195,7 +195,7 @@ def run(self): angles = np.round(angles, 0) detections_msg = create_detection_message( - bboxes=bboxes, scores=scores, angles=angles + bboxes=bboxes, scores=scores, angles=angles, keypoints=points ) detections_msg.setTimestamp(output.getTimestamp()) detections_msg.transformation = output.getTransformation()