From e5375fdeefa786cf99247bdd6cc51fc5d8eaa0b3 Mon Sep 17 00:00:00 2001 From: klemen1999 Date: Tue, 18 Aug 2026 13:51:13 +0200 Subject: [PATCH] Replace dai-nodes messages with DAI native messages --- depthai_nodes/message/README.md | 150 ++------- depthai_nodes/message/__init__.py | 15 - depthai_nodes/message/classification.py | 192 ------------ depthai_nodes/message/clusters.py | 210 ------------- .../message/creators/classification.py | 15 +- depthai_nodes/message/creators/clusters.py | 18 +- depthai_nodes/message/creators/keypoints.py | 10 +- depthai_nodes/message/creators/line.py | 16 +- depthai_nodes/message/creators/map.py | 15 +- depthai_nodes/message/creators/regression.py | 10 +- depthai_nodes/message/keypoints.py | 104 ------- depthai_nodes/message/lines.py | 234 --------------- depthai_nodes/message/map.py | 149 --------- depthai_nodes/message/prediction.py | 195 ------------ depthai_nodes/message/utils/copy_message.py | 58 ++++ depthai_nodes/node/README.md | 12 +- depthai_nodes/node/apply_colormap.py | 10 +- depthai_nodes/node/parsers/classification.py | 2 +- .../node/parsers/classification_sequence.py | 2 +- depthai_nodes/node/parsers/hrnet.py | 2 +- depthai_nodes/node/parsers/keypoints.py | 2 +- depthai_nodes/node/parsers/lane_detection.py | 2 +- depthai_nodes/node/parsers/map_output.py | 4 +- depthai_nodes/node/parsers/mlsd.py | 4 +- depthai_nodes/node/parsers/regression.py | 2 +- .../node/parsers/superanimal_landmarker.py | 2 +- depthai_nodes/node/utils/message_remapping.py | 284 ++---------------- depthai_nodes/node/utils/util_constants.py | 20 +- tests/stability_tests/check_messages.py | 79 ++--- .../test_classification_sequence.py | 4 +- .../test_creators/test_classifications.py | 6 +- .../unittests/test_creators/test_clusters.py | 10 +- .../unittests/test_creators/test_keypoints.py | 9 +- .../test_creators/test_line_detections.py | 16 +- tests/unittests/test_creators/test_map.py | 44 +-- .../test_creators/test_regression.py | 14 +- .../test_messages/test_classifications_msg.py | 69 ----- .../test_messages/test_clusters_msg.py | 72 ----- .../test_messages/test_keypoints_msg.py | 48 --- .../unittests/test_messages/test_lines_msg.py | 93 ------ tests/unittests/test_messages/test_map_msg.py | 48 --- .../test_messages/test_predictions_msg.py | 68 ----- 42 files changed, 261 insertions(+), 2058 deletions(-) delete mode 100644 depthai_nodes/message/classification.py delete mode 100644 depthai_nodes/message/clusters.py delete mode 100644 depthai_nodes/message/keypoints.py delete mode 100644 depthai_nodes/message/lines.py delete mode 100644 depthai_nodes/message/map.py delete mode 100644 depthai_nodes/message/prediction.py delete mode 100644 tests/unittests/test_messages/test_classifications_msg.py delete mode 100644 tests/unittests/test_messages/test_clusters_msg.py delete mode 100644 tests/unittests/test_messages/test_keypoints_msg.py delete mode 100644 tests/unittests/test_messages/test_lines_msg.py delete mode 100644 tests/unittests/test_messages/test_map_msg.py delete mode 100644 tests/unittests/test_messages/test_predictions_msg.py diff --git a/depthai_nodes/message/README.md b/depthai_nodes/message/README.md index ccf1bdfa..0a524325 100644 --- a/depthai_nodes/message/README.md +++ b/depthai_nodes/message/README.md @@ -1,146 +1,30 @@ # Message Types -Here are the custom message types that we introduce in this package. They are used as output types of the parsers. +Parser creators return native DepthAI messages. The parser-specific message types +are available under `dai.beta`: -**Table of Contents** +- `dai.beta.Classifications` +- `dai.beta.Clusters` and `dai.beta.Cluster` +- `dai.beta.Keypoints` +- `dai.beta.Lines` and `dai.beta.Line` +- `dai.beta.Map2D` +- `dai.beta.Predictions` and `dai.beta.Prediction` -- [Message Types](#message-types) - - [Classifications](#classifications) - - [Attributes](#attributes) - - [Cluster](#cluster) - - [Attributes](#attributes-1) - - [Clusters](#clusters) - - [Attributes](#attributes-2) - - [Collection](#collection) - - [Attributes](#attributes-3) - - [GatheredData](#gathereddata) - - [Attributes](#attributes-4) - - [Keypoints](#keypoints) - - [Attributes](#attributes-5) - - [Line](#line) - - [Attributes](#attributes-6) - - [Lines](#lines) - - [Attributes](#attributes-7) - - [Map2D](#map2d) - - [Attributes](#attributes-8) - - [Prediction](#prediction) - - [Attributes](#attributes-9) - - [Predictions](#predictions) - - [Attributes](#attributes-10) - - [SnapData](#snapdata) - - [Attributes](#attributes-12) - -## Classifications - -Classification class for storing the classes and their respective scores. - -### Attributes - -- **classes** (list\[str\]): A list of classes. -- **scores** (NDArray\[np.float32\]): Corresponding probability scores. - -## Cluster - -Cluster class for storing a cluster. - -### Attributes - -- **label** (int): Label of the cluster. -- **points** (List\[dai.Point2f\]): List of points in the cluster. - -## Clusters - -Clusters class for storing clusters. - -### Attributes - -- **clusters** (List\[[Cluster](#cluster)\]): List of clusters. +depthai-nodes retains only the package-specific messages below because DepthAI has +no native equivalent for them. ## Collection -Collection class for storing a list of messages or other items of the same type. - -### Attributes - -- **items** (List\[T\]): List of collected items. -- **item_cls** (Optional\[Type\[T\]\]): Runtime item type inferred from the first item once the collection is non-empty. - -Items can be added with `append(...)` or `extend(...)`. The collection enforces that all items have the same inferred type. +`Collection` stores a list of messages or other items of the same runtime type. +Items can be added with `append(...)` or `extend(...)`. ## GatheredData -GatheredData class for storing a reference message and the messages gathered for that reference. - -### Attributes - -- **reference_data** (TReference): Reference message used to determine how many items to gather. -- **items** (List\[TGathered\]): List of gathered messages. -- **item_cls** (Optional\[Type\[TGathered\]\]): Runtime gathered-item type inferred from the first item once the collection is non-empty. - -GatheredData inherits [Collection](#collection) behavior, so all gathered items must have the same inferred type. Setting `reference_data` copies the reference message sequence number, timestamp, and device timestamp to the GatheredData message. - -## Keypoints - -Keypoints class for storing keypoints and optional skeleton edges. - -### Attributes - -- **keypoints_list** (dai.KeypointsList): Native DepthAI keypoints list containing keypoints and edges. -- **transformation** (Optional\[dai.ImgTransformation\]): Optional image transformation associated with the keypoints. - -The keypoints can be accessed with `getKeypoints()` and set with `setKeypoints(...)`. Each keypoint is a `dai.Keypoint` with image coordinates, confidence, and optional label name. Edges can be accessed with `getEdges()` and set with `setEdges(...)`. - -## Line - -Line class for storing a line. - -### Attributes - -- **start_point** (dai.Point2f): Start point of the line with x and y coordinates. -- **end_point** (dai.Point2f): End point of the line with x and y coordinates. -- **confidence** (float): Confidence of the line. - -## Lines - -Lines class for storing lines. - -### Attributes - -- **lines** (List\[[Line](#line)\]): List of detected lines. - -## Map2D - -Map2D class for storing a 2D map of floats. - -### Attributes - -- **map** (NDArray\[np.float32\]): 2D map. -- **width** (int): 2D Map width. -- **height** (int): 2D Map height. - -## Prediction - -Prediction class for storing a prediction. - -### Attributes - -- **prediction** (float): The predicted value. - -## Predictions - -Predictions class for storing predictions. - -### Attributes - -- **predictions** (List\[[Prediction](#prediction)\]): List of predictions. +`GatheredData` stores a reference message and the messages gathered for that +reference. It inherits `Collection` and copies timestamp and sequence metadata from +the reference message. ## SnapData -SnapData class for representing a single snap event to be uploaded to DepthAI Hub. - -### Attributes - -- **snap_name** (str): Logical name of the snap. -- **file_group** (dai.FileGroup): Object containing the snap image and associated data (e.g., images, detections). -- **tags** (List\[str\]): Optional list of tags for categorizing the snap. -- **extras** (Dict\[str, str\]): Additional metadata as key-value pairs. +`SnapData` represents a snap event uploaded to DepthAI Hub. It contains the snap +name, a `dai.FileGroup`, optional tags, and string metadata. diff --git a/depthai_nodes/message/__init__.py b/depthai_nodes/message/__init__.py index a14cf509..1781ed4c 100644 --- a/depthai_nodes/message/__init__.py +++ b/depthai_nodes/message/__init__.py @@ -1,24 +1,9 @@ -from .classification import Classifications -from .clusters import Cluster, Clusters from .collection import Collection from .gathered_data import GatheredData -from .keypoints import Keypoints -from .lines import Line, Lines -from .map import Map2D -from .prediction import Prediction, Predictions from .snap_data import SnapData __all__ = [ - "Line", - "Lines", - "Classifications", - "Map2D", - "Clusters", - "Cluster", - "Prediction", - "Predictions", "GatheredData", - "Keypoints", "SnapData", "Collection", ] diff --git a/depthai_nodes/message/classification.py b/depthai_nodes/message/classification.py deleted file mode 100644 index 103f9e4f..00000000 --- a/depthai_nodes/message/classification.py +++ /dev/null @@ -1,192 +0,0 @@ -import copy - -import depthai as dai -import numpy as np -from numpy.typing import NDArray - -from depthai_nodes import FONT_BACKGROUND_COLOR, FONT_COLOR -from depthai_nodes.utils import AnnotationHelper, AnnotationSizes - - -class Classifications(dai.Buffer): - """Classification class for storing the classes and their respective scores. - - Attributes - ---------- - classes : list[str] - A list of classes. - scores : NDArray[np.float32] - Corresponding probability scores. - transformation : dai.ImgTransformation - Image transformation object. - """ - - def __init__(self): - """Initializes the Classifications object.""" - dai.Buffer.__init__(self) - self._classes: list[str] = [] - self._scores: NDArray[np.float32] = np.array([]) - self._transformation: dai.ImgTransformation | None = None - - def copy(self): - """Creates a new instance of the Classifications class and copies the - attributes. - - @return: A new instance of the Classifications class. - @rtype: Classifications - """ - new_obj = Classifications() - new_obj.classes = copy.deepcopy(self.classes) - new_obj.scores = copy.deepcopy(self.scores) - new_obj.setSequenceNum(self.getSequenceNum()) - new_obj.setTimestamp(self.getTimestamp()) - new_obj.setTimestampDevice(self.getTimestampDevice()) - new_obj.setTransformation(self.transformation) - return new_obj - - @property - def classes(self) -> list: - """Returns the list of classes. - - @return: List of classes. - @rtype: list[str] - """ - return self._classes - - @classes.setter - def classes(self, value: list[str]): - """Sets the classes. - - @param value: A list of class names. - @type value: list[str] - @raise TypeError: If value is not a list. - @raise ValueError: If each element is not of type string. - """ - if not isinstance(value, list): - raise TypeError(f"Classes must be a list, instead got {type(value)}.") - if not all(isinstance(class_name, str) for class_name in value): - raise ValueError("Classes must be a list of strings.") - self._classes = value - - @property - def scores(self) -> NDArray: - """Returns the list of scores. - - @return: List of scores. - @rtype: NDArray[np.float32] - """ - return self._scores - - @scores.setter - def scores(self, value: NDArray[np.float32]): - """Sets the scores. - - @param value: A list of scores. - @type value: NDArray[np.float32] - @raise TypeError: If value is not a numpy array. - @raise ValueError: If value is not a 1D numpy array. - @raise ValueError: If each element is not of type float. - """ - if not isinstance(value, np.ndarray): - raise TypeError(f"Scores must be a np.ndarray, instead got {type(value)}.") - if value.ndim != 1: - raise ValueError("Scores must be a 1D a np.ndarray.") - if value.size > 0 and value.dtype != np.float32: - raise ValueError("Scores must be a np.ndarray of floats.") - self._scores = value - - @property - def top_class(self) -> str: - """Returns the most probable class. Only works if classes are sorted by scores. - - @return: The top class. - @rtype: str - """ - return self._classes[0] - - @property - def top_score(self) -> float: - """Returns the probability of the most probable class. Only works if scores are - sorted by descending order. - - @return: The top score. - @rtype: float - """ - return self._scores[0] - - @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.ImgAnnotations: - """Returns default visualization message for classification. - - The message adds the top five classes and their scores to the right side of the - image. - """ - if self.transformation is None: - raise ValueError("Transformation must be set to get visualization message.") - - w, h = self.transformation.getSize() - annotation_sizes = AnnotationSizes(w, h) - x_offset = 2 / w - y_offset = 2 / h - - annotation_helper = AnnotationHelper() - for i in range(min(5, len(self._classes))): - y_position = ( - y_offset - + (annotation_sizes.relative_font_size) - + i * (annotation_sizes.relative_font_size) - ) - annotation_helper.draw_text( - text=f"{self._classes[i]} {self._scores[i] * 100:.0f}%", - position=( - x_offset, - y_position, - ), - color=FONT_COLOR, - background_color=FONT_BACKGROUND_COLOR, - size=annotation_sizes.font_size, - ) - return annotation_helper.build( - timestamp=self.getTimestamp(), sequence_num=self.getSequenceNum() - ) diff --git a/depthai_nodes/message/clusters.py b/depthai_nodes/message/clusters.py deleted file mode 100644 index d3dce6b2..00000000 --- a/depthai_nodes/message/clusters.py +++ /dev/null @@ -1,210 +0,0 @@ -import copy - -import cv2 -import depthai as dai -import numpy as np - -from .utils import ( - copy_message, -) - - -class Cluster(dai.Buffer): - """Cluster class for storing a cluster. - - Attributes - ---------- - label : int - Label of the cluster. - points : list[dai.Point2f] - List of points in the cluster. - """ - - def __init__(self): - """Initializes the Cluster object.""" - super().__init__() - self._label: int = None - self._points: list[dai.Point2f] = [] - - def copy(self): - """Creates a new instance of the Cluster class and copies the attributes. - - @return: A new instance of the Cluster class. - @rtype: Cluster - """ - new_obj = Cluster() - new_obj.label = copy.deepcopy(self.label) - new_obj.points = [copy_message(p) for p in self.points] - return new_obj - - @property - def label(self) -> int: - """Returns the label of the cluster. - - @return: Label of the cluster. - @rtype: int - """ - return self._label - - @label.setter - def label(self, value: int): - """Sets the label of the cluster. - - @param value: Label of the cluster. - @type value: int - @raise TypeError: If value is not an int. - """ - if not isinstance(value, int): - raise TypeError(f"Label must be of type int, instead got {type(value)}.") - self._label = value - - @property - def points(self) -> list[dai.Point2f]: - """Returns the points in the cluster. - - @return: List of points in the cluster. - @rtype: list[dai.Point2f - """ - return self._points - - @points.setter - def points(self, value: list[dai.Point2f]): - """Sets the points in the cluster. - - @param value: List of points in the cluster. - @type value: list[dai.Point2f] - @raise TypeError: If value is not a list. - @raise TypeError: If each element is not of type dai.Point2f. - """ - if not isinstance(value, list): - raise TypeError(f"Points must be a list, instead got {type(value)}.") - if not all(isinstance(point, dai.Point2f) for point in value): - raise ValueError("Points must be a list of dai.Point2f objects") - self._points = value - - -class Clusters(dai.Buffer): - """Clusters class for storing clusters. - - Attributes - ---------- - clusters : list[Cluster] - List of clusters. - transformation : dai.ImgTransformation - Image transformation object. - """ - - def __init__(self): - """Initializes the Clusters object.""" - super().__init__() - self._clusters: list[Cluster] = [] - self._transformation: dai.ImgTransformation | None = None - - def copy(self): - """Creates a new instance of the Clusters class and copies the attributes. - - @return: A new instance of the Clusters class. - @rtype: Clusters - """ - new_obj = Clusters() - new_obj.clusters = [cluster.copy() for cluster in self.clusters] - new_obj.setSequenceNum(self.getSequenceNum()) - new_obj.setTimestamp(self.getTimestamp()) - new_obj.setTimestampDevice(self.getTimestampDevice()) - new_obj.setTransformation(self.transformation) - return new_obj - - @property - def clusters(self) -> list[Cluster]: - """Returns the clusters. - - @return: List of clusters. - @rtype: list[Cluster] - """ - return self._clusters - - @clusters.setter - def clusters(self, value: list[Cluster]): - """Sets the clusters. - - @param value: List of clusters. - @type value: list[Cluster] - @raise TypeError: If value is not a list. - @raise ValueError: If each element is not of type Cluster. - """ - if not isinstance(value, list): - raise TypeError("Clusters must be a list.") - if not all(isinstance(cluster, Cluster) for cluster in value): - raise ValueError("Clusters must be a list of Cluster objects.") - self._clusters = 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.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()) - img_annotations.setSequenceNum(self.getSequenceNum()) - img_annotations.setTimestampDevice(self.getTimestampDevice()) - return img_annotations diff --git a/depthai_nodes/message/creators/classification.py b/depthai_nodes/message/creators/classification.py index 2a426696..c184ce69 100644 --- a/depthai_nodes/message/creators/classification.py +++ b/depthai_nodes/message/creators/classification.py @@ -1,11 +1,10 @@ +import depthai as dai import numpy as np -from depthai_nodes import Classifications - def create_classification_message( classes: list[str], scores: np.ndarray | list -) -> Classifications: +) -> dai.beta.Classifications: """Create a message for classification. The message contains the class names and their respective scores, sorted in descending order of scores. @@ -14,7 +13,7 @@ def create_classification_message( @type scores: np.ndarray @return: A message with attributes `classes` and `scores`. `classes` is a list of classes, sorted in descending order of scores. `scores` is a list of the corresponding scores. - @rtype: Classifications + @rtype: dai.beta.Classifications @raises ValueError: If the provided classes are None. @raises ValueError: If the provided classes are not a list. @@ -73,7 +72,7 @@ def create_classification_message( f"Number of labels and scores mismatch. Provided {len(scores)} scores and {len(classes)} class names." ) - classification_msg = Classifications() + classification_msg = dai.beta.Classifications() sorted_args = np.argsort(-scores, kind="stable") scores = scores[sorted_args] @@ -89,7 +88,7 @@ def create_classification_sequence_message( ignored_indexes: list[int] | None = None, remove_duplicates: bool = False, concatenate_classes: bool = False, -) -> Classifications: +) -> dai.beta.Classifications: """Creates a message for a multi-class sequence. The message contains the class names and their respective scores, ordered according to the sequence. The 'scores' array is a sequence of probabilities for each class at each position in the @@ -104,7 +103,7 @@ def create_classification_sequence_message( @param concatenate_classes: If True, concatenates consecutive classes based on the space character. Defaults to False. @type concatenate_classes: bool @return: A Classification message with attributes `classes` and `scores`, where `classes` is a list of class names and `scores` is a list of corresponding scores. - @rtype: Classifications + @rtype: dai.beta.Classifications @raises ValueError: If 'classes' is not a list of strings. @raises ValueError: If 'scores' is not a 2D array of list of shape (sequence_length, n_classes). @raises ValueError: If the number of classes does not match the number of columns in 'scores'. @@ -190,7 +189,7 @@ def create_classification_sequence_message( mean_score = np.mean(score_list) score_list = np.array([mean_score]) - classification_msg = Classifications() + classification_msg = dai.beta.Classifications() classification_msg.classes = class_list classification_msg.scores = score_list diff --git a/depthai_nodes/message/creators/clusters.py b/depthai_nodes/message/creators/clusters.py index d46d7cb3..aef1adc3 100644 --- a/depthai_nodes/message/creators/clusters.py +++ b/depthai_nodes/message/creators/clusters.py @@ -1,16 +1,16 @@ import depthai as dai -from depthai_nodes import Cluster, Clusters - -def create_cluster_message(clusters: list[list[list[float | int]]]) -> Clusters: +def create_cluster_message( + clusters: list[list[list[float | int]]], +) -> dai.beta.Clusters: """Create a DepthAI message for clusters. @param clusters: List of clusters. Each cluster is a list of points with x and y coordinates. @type clusters: list[list[list[float | int]]] @return: Clusters message containing the detected clusters. - @rtype: Clusters + @rtype: dai.beta.Clusters @raise TypeError: If the clusters are not a list. @raise TypeError: If each cluster is not a list. @raise TypeError: If each point is not a list. @@ -35,14 +35,14 @@ def create_cluster_message(clusters: list[list[list[float | int]]]) -> Clusters: f"All items in points must be of type int or float, got {type(value)}" ) - message = Clusters() + message = dai.beta.Clusters() temp = [] for i, cluster in enumerate(clusters): - temp_cluster = Cluster() + temp_cluster = dai.beta.Cluster() temp_cluster.label = i - temp_cluster.points = [ - dai.Point2f(float(point[0]), float(point[1])) for point in cluster - ] + temp_cluster.points = dai.VectorPoint2f( + [dai.Point2f(float(point[0]), float(point[1])) for point in cluster] + ) temp.append(temp_cluster) diff --git a/depthai_nodes/message/creators/keypoints.py b/depthai_nodes/message/creators/keypoints.py index d0e520d5..4cf865b3 100644 --- a/depthai_nodes/message/creators/keypoints.py +++ b/depthai_nodes/message/creators/keypoints.py @@ -1,8 +1,6 @@ import depthai as dai import numpy as np -from depthai_nodes.message.keypoints import Keypoints - def create_keypoints_message( keypoints: np.ndarray | list[list[float]], @@ -10,8 +8,8 @@ def create_keypoints_message( confidence_threshold: float | None = None, label_names: list[str] | None = None, edges: list[tuple[int, int]] | None = None, -) -> Keypoints: - """Create a native DepthAI keypoints message.""" +) -> dai.beta.Keypoints: + """Create a native ``dai.beta.Keypoints`` message.""" if not isinstance(keypoints, (np.ndarray, list)): raise ValueError( @@ -140,6 +138,6 @@ def create_keypoints_message( keypoints_list.setEdges(filtered_edges) - keypoints_msg = Keypoints() - keypoints_msg.keypoints_list = keypoints_list + keypoints_msg = dai.beta.Keypoints() + keypoints_msg.keypointsList = keypoints_list return keypoints_msg diff --git a/depthai_nodes/message/creators/line.py b/depthai_nodes/message/creators/line.py index 4e27b7a6..7e404757 100644 --- a/depthai_nodes/message/creators/line.py +++ b/depthai_nodes/message/creators/line.py @@ -1,10 +1,10 @@ import depthai as dai import numpy as np -from depthai_nodes import Line, Lines - -def create_line_detection_message(lines: np.ndarray, scores: np.ndarray): +def create_line_detection_message( + lines: np.ndarray, scores: np.ndarray +) -> dai.beta.Lines: """Create a DepthAI message for a line detection. @param lines: Detected lines of shape (N,4) meaning [...,[x_start, y_start, x_end, y_end],...]. @@ -13,7 +13,7 @@ def create_line_detection_message(lines: np.ndarray, scores: np.ndarray): @type scores: np.ndarray @return: Message containing the lines and confidence scores of detected lines. - @rtype: Lines + @rtype: dai.beta.Lines @raise ValueError: If the lines are not a numpy array. @raise ValueError: If the lines are not of shape (N,4). @@ -57,12 +57,12 @@ def create_line_detection_message(lines: np.ndarray, scores: np.ndarray): line_detections = [] for i, line in enumerate(lines): - line_detection = Line() - line_detection.start_point = dai.Point2f(line[0], line[1]) - line_detection.end_point = dai.Point2f(line[2], line[3]) + line_detection = dai.beta.Line() + line_detection.startPoint = dai.Point2f(line[0], line[1]) + line_detection.endPoint = dai.Point2f(line[2], line[3]) line_detection.confidence = float(scores[i]) line_detections.append(line_detection) - lines_msg = Lines() + lines_msg = dai.beta.Lines() lines_msg.lines = line_detections return lines_msg diff --git a/depthai_nodes/message/creators/map.py b/depthai_nodes/message/creators/map.py index 1e5051ba..2cca43ce 100644 --- a/depthai_nodes/message/creators/map.py +++ b/depthai_nodes/message/creators/map.py @@ -1,9 +1,10 @@ +import depthai as dai import numpy as np -from depthai_nodes import Map2D - -def create_map_message(map: np.ndarray, min_max_scaling: bool = False) -> Map2D: +def create_map_message( + map: np.ndarray, min_max_scaling: bool = False +) -> dai.beta.Map2D: """Create a DepthAI message for a map of floats. @param map: A NumPy array representing the map with shape HW or NHW/HWN. Here N @@ -12,8 +13,8 @@ def create_map_message(map: np.ndarray, min_max_scaling: bool = False) -> Map2D: @param min_max_scaling: If True, the map is scaled to the range [0, 1]. Defaults to False. @type min_max_scaling: bool - @return: An Map2D object containing the density information. - @rtype: Map2D + @return: A native Map2D object containing the density information. + @rtype: dai.beta.Map2D @raise ValueError: If the density map is not a NumPy array. @raise ValueError: If the density map is not 2D or 3D. @raise ValueError: If the 3D density map shape is not NHW or HWN. @@ -44,7 +45,7 @@ def create_map_message(map: np.ndarray, min_max_scaling: bool = False) -> Map2D: if map.dtype != np.float32: map = map.astype(np.float32) - map_2d = Map2D() - map_2d.map = map + map_2d = dai.beta.Map2D() + map_2d.setMap(map) return map_2d diff --git a/depthai_nodes/message/creators/regression.py b/depthai_nodes/message/creators/regression.py index 6ae73a20..4f30e5ed 100644 --- a/depthai_nodes/message/creators/regression.py +++ b/depthai_nodes/message/creators/regression.py @@ -1,13 +1,13 @@ -from depthai_nodes import Prediction, Predictions +import depthai as dai -def create_regression_message(predictions: list[float]) -> Predictions: +def create_regression_message(predictions: list[float]) -> dai.beta.Predictions: """Create a DepthAI message for prediction models. @param predictions: Predicted value(s). @type predictions: list[float] @return: Predictions message containing the predicted value(s). - @rtype: Predictions + @rtype: dai.beta.Predictions @raise ValueError: If predictions is not a list. @raise ValueError: If each prediction is not a float. """ @@ -23,11 +23,11 @@ def create_regression_message(predictions: list[float]) -> Predictions: prediction_objects_list = [] for prediction in predictions: - prediction_object = Prediction() + prediction_object = dai.beta.Prediction() prediction_object.prediction = prediction prediction_objects_list.append(prediction_object) - predictions_message = Predictions() + predictions_message = dai.beta.Predictions() predictions_message.predictions = prediction_objects_list return predictions_message diff --git a/depthai_nodes/message/keypoints.py b/depthai_nodes/message/keypoints.py deleted file mode 100644 index c27d928e..00000000 --- a/depthai_nodes/message/keypoints.py +++ /dev/null @@ -1,104 +0,0 @@ -import copy - -import depthai as dai - -from depthai_nodes import KEYPOINT_COLOR, PRIMARY_COLOR -from depthai_nodes.utils import AnnotationHelper - - -class Keypoints(dai.Buffer): - """DepthAI Nodes keypoints message wrapping a native ``dai.KeypointsList``.""" - - def __init__(self): - super().__init__() - self._keypoints_list = dai.KeypointsList() - self._transformation: dai.ImgTransformation | None = None - - def copy(self): - new_obj = Keypoints() - native_copy = dai.KeypointsList() - native_copy.setKeypoints([copy.deepcopy(kp) for kp in self.getKeypoints()]) - native_copy.setEdges(copy.deepcopy(self.getEdges())) - new_obj.keypoints_list = native_copy - 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 keypoints_list(self) -> dai.KeypointsList: - return self._keypoints_list - - @keypoints_list.setter - def keypoints_list(self, value: dai.KeypointsList): - if not isinstance(value, dai.KeypointsList): - raise TypeError( - f"keypoints_list must be a dai.KeypointsList, got {type(value)}." - ) - self._keypoints_list = value - - @property - def transformation(self) -> dai.ImgTransformation | None: - return self._transformation - - @transformation.setter - def transformation(self, value: dai.ImgTransformation | None): - if value is not None and not isinstance(value, dai.ImgTransformation): - raise TypeError( - f"Transformation must be a dai.ImgTransformation object, got {type(value)}." - ) - self._transformation = value - - def setTransformation(self, transformation: dai.ImgTransformation | None): - self.transformation = transformation - - def getTransformation(self) -> dai.ImgTransformation | None: - return self.transformation - - def getKeypoints(self) -> list[dai.Keypoint]: - return self._keypoints_list.getKeypoints() - - def setKeypoints(self, value: list[dai.Keypoint]): - self._keypoints_list.setKeypoints(value) - - def getEdges(self) -> list[tuple[int, int]]: - return self._keypoints_list.getEdges() - - def setEdges(self, value: list[tuple[int, int]]): - self._keypoints_list.setEdges(value) - - def getPoints2f(self) -> dai.VectorPoint2f: - return dai.VectorPoint2f( - [ - dai.Point2f(kp.imageCoordinates.x, kp.imageCoordinates.y) - for kp in self.getKeypoints() - ] - ) - - def getPoints3f(self) -> list[dai.Point3f]: - return [ - dai.Point3f( - kp.imageCoordinates.x, kp.imageCoordinates.y, kp.imageCoordinates.z - ) - for kp in self.getKeypoints() - ] - - def getVisualizationMessage(self) -> dai.ImgAnnotations: - annotation_helper = AnnotationHelper() - annotation_helper.draw_points( - points=self.getPoints2f(), color=KEYPOINT_COLOR, thickness=1 - ) - for edge in self.getEdges(): - pt1_ix, pt2_ix = edge - pt1 = self.getKeypoints()[pt1_ix] - pt2 = self.getKeypoints()[pt2_ix] - annotation_helper.draw_line( - pt1=(pt1.imageCoordinates.x, pt1.imageCoordinates.y), - pt2=(pt2.imageCoordinates.x, pt2.imageCoordinates.y), - color=PRIMARY_COLOR, - thickness=1, - ) - return annotation_helper.build( - timestamp=self.getTimestamp(), sequence_num=self.getSequenceNum() - ) diff --git a/depthai_nodes/message/lines.py b/depthai_nodes/message/lines.py deleted file mode 100644 index 830c4b41..00000000 --- a/depthai_nodes/message/lines.py +++ /dev/null @@ -1,234 +0,0 @@ -import copy - -import depthai as dai - -from depthai_nodes import PRIMARY_COLOR -from depthai_nodes.logging import get_logger -from depthai_nodes.utils import AnnotationHelper - -from .utils import ( - copy_message, -) - - -class Line(dai.Buffer): - """Line class for storing a line. - - Attributes - ---------- - start_point : dai.Point2f - Start point of the line with x and y coordinate. - end_point : dai.Point2f - End point of the line with x and y coordinate. - confidence : float - Confidence of the line. - """ - - def __init__(self): - """Initializes the Line object.""" - super().__init__() - self._start_point: dai.Point2f = None - self._end_point: dai.Point2f = None - self._confidence: float = None - self._logger = get_logger(__name__) - - def copy(self): - """Creates a new instance of the Line class and copies the attributes. - - @return: A new instance of the Line class. - @rtype: Line - """ - new_obj = Line() - new_obj.start_point = copy_message(self.start_point) - new_obj.end_point = copy_message(self.end_point) - new_obj.confidence = copy.deepcopy(self.confidence) - return new_obj - - @property - def start_point(self) -> dai.Point2f: - """Returns the start point of the line. - - @return: Start point of the line. - @rtype: dai.Point2f - """ - return self._start_point - - @start_point.setter - def start_point(self, value: dai.Point2f): - """Sets the start point of the line. - - @param value: Start point of the line. - @type value: dai.Point2f - @raise TypeError: If value is not of type dai.Point2f. - """ - if not isinstance(value, dai.Point2f): - raise TypeError( - f"Start Point must be of type Point2f, instead got {type(value)}." - ) - self._start_point = value - - @property - def end_point(self) -> dai.Point2f: - """Returns the end point of the line. - - @return: End point of the line. - @rtype: dai.Point2f - """ - return self._end_point - - @end_point.setter - def end_point(self, value: dai.Point2f): - """Sets the end point of the line. - - @param value: End point of the line. - @type value: dai.Point2f - @raise TypeError: If value is not of type dai.Point2f. - """ - if not isinstance(value, dai.Point2f): - raise TypeError( - f"End Point must be of type Point2f, instead got {type(value)}." - ) - self._end_point = value - - @property - def confidence(self) -> float: - """Returns the confidence of the line. - - @return: Confidence of the line. - @rtype: float - """ - return self._confidence - - @confidence.setter - def confidence(self, value: float): - """Sets the confidence of the line. - - @param value: Confidence of the line. - @type value: float - @raise TypeError: If value is not a float. - @raise ValueError: If value is not between 0 and 1. - """ - if not isinstance(value, float): - raise TypeError("Confidence must be a float.") - if value < -0.1 or value > 1.1: - raise ValueError("Confidence must be between 0 and 1.") - if not (0 <= value <= 1): - value = float(max(0.0, min(1.0, value))) - self._logger.info("Confidence value was clipped to [0, 1].") - - self._confidence = value - - -class Lines(dai.Buffer): - """Lines class for storing lines. - - Attributes - ---------- - lines : list[Line] - List of detected lines. - transformation : dai.ImgTransformation - Image transformation object. - """ - - def __init__(self): - """Initializes the Lines object.""" - super().__init__() - self._lines: list[Line] = [] - self._transformation: dai.ImgTransformation | None = None - - def copy(self): - """Creates a new instance of the Lines class and copies the attributes. - - @return: A new instance of the Lines class. - @rtype: Lines - """ - new_obj = Lines() - new_obj.lines = [line.copy() for line in self.lines] - new_obj.setSequenceNum(self.getSequenceNum()) - new_obj.setTimestamp(self.getTimestamp()) - new_obj.setTimestampDevice(self.getTimestampDevice()) - new_obj.setTransformation(self.transformation) - return new_obj - - @property - def lines(self) -> list[Line]: - """Returns the lines. - - @return: List of lines. - @rtype: list[Line] - """ - return self._lines - - @lines.setter - def lines(self, value: list[Line]): - """Sets the lines. - - @param value: List of lines. - @type value: list[Line] - @raise TypeError: If value is not a list. - @raise TypeError: If each element is not of type Line. - """ - if not isinstance(value, list): - raise TypeError(f"lines must be a list, instead got {type(value)}.") - if not all(isinstance(item, Line) for item in value): - raise ValueError("Lines must be a list of Line objects.") - self._lines = 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.ImgAnnotations: - """Returns default visualization message for lines. - - The message adds lines to the image. - """ - annotation_helper = AnnotationHelper() - for line in self.lines: - annotation_helper.draw_line( - pt1=line.start_point, - pt2=line.end_point, - color=PRIMARY_COLOR, - thickness=2.0, - ) - return annotation_helper.build( - timestamp=self.getTimestamp(), sequence_num=self.getSequenceNum() - ) diff --git a/depthai_nodes/message/map.py b/depthai_nodes/message/map.py deleted file mode 100644 index 0dda99c5..00000000 --- a/depthai_nodes/message/map.py +++ /dev/null @@ -1,149 +0,0 @@ -import copy - -import cv2 -import depthai as dai -import numpy as np -from numpy.typing import NDArray - - -class Map2D(dai.Buffer): - """Map2D class for storing a 2D map of floats. - - Attributes - ---------- - map : NDArray[np.float32] - 2D map. - width : int - 2D Map width. - height : int - 2D Map height. - transformation : dai.ImgTransformation - Image transformation object. - """ - - def __init__(self): - """Initializes the Map2D object.""" - super().__init__() - self._map: NDArray[np.float32] = np.array([]) - self._width: int = None - self._height: int = None - self._transformation: dai.ImgTransformation | None = None - - def copy(self): - """Creates a new instance of the Map2D class and copies the attributes. - - @return: A new instance of the Map2D class. - @rtype: Map2D - """ - new_obj = Map2D() - new_obj.map = copy.deepcopy(self._map) - new_obj.transformation = self._transformation - return new_obj - - @property - def map(self) -> NDArray[np.float32]: - """Returns the 2D map. - - @return: 2D map. - @rtype: NDArray[np.float32] - """ - return self._map - - @map.setter - def map(self, value: np.ndarray): - """Sets the 2D map. - - @param value: 2D map. - @type value: NDArray[np.float32] - @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 float. - """ - if not isinstance(value, np.ndarray): - raise TypeError( - f"2D map must be of type np.ndarray, instead got {type(value)}." - ) - if value.ndim != 2: - raise ValueError("2D map must be a 2D array") - if value.dtype != np.float32: - raise ValueError("2D map must be an array of floats") - self._map = value - self._width = value.shape[1] - self._height = value.shape[0] - - @property - def width(self) -> int: - """Returns the 2D map width. - - @return: 2D map width. - @rtype: int - """ - return self._width - - @property - def height(self) -> int: - """Returns the 2D map height. - - @return: 2D map height. - @rtype: int - """ - return self._height - - @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): - """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 default visualization message for 2D maps in the form of a - colormapped image.""" - 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._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/message/prediction.py b/depthai_nodes/message/prediction.py deleted file mode 100644 index 601d2bac..00000000 --- a/depthai_nodes/message/prediction.py +++ /dev/null @@ -1,195 +0,0 @@ -import copy - -import depthai as dai - -from depthai_nodes import FONT_BACKGROUND_COLOR, FONT_COLOR -from depthai_nodes.utils import AnnotationHelper, AnnotationSizes - - -class Prediction(dai.Buffer): - """Prediction class for storing a prediction. - - Attributes - ---------- - prediction : float - The predicted value. - """ - - def __init__(self): - """Initializes the Prediction object.""" - super().__init__() - self._prediction: float = None - - def copy(self): - """Creates a new instance of the Prediction class and copies the attributes. - - @return: A new instance of the Prediction class. - @rtype: Prediction - """ - new_obj = Prediction() - new_obj.prediction = copy.deepcopy(self.prediction) - return new_obj - - @property - def prediction(self) -> float: - """Returns the prediction. - - @return: The predicted value. - @rtype: float - """ - return self._prediction - - @prediction.setter - def prediction(self, value: float): - """Sets the prediction. - - @param value: The predicted value. - @type value: float - @raise TypeError: If value is not of type float. - """ - if not isinstance(value, float): - raise TypeError( - f"Prediction must be of type float, instead got {type(value)}." - ) - self._prediction = value - - -class Predictions(dai.Buffer): - """Predictions class for storing predictions. - - Attributes - ---------- - predictions : list[Prediction] - List of predictions. - transformation : dai.ImgTransformation - Image transformation object. - """ - - def __init__(self): - """Initializes the Predictions object.""" - super().__init__() - self._predictions: list[Prediction] = [] - self._transformation: dai.ImgTransformation | None = None - - def copy(self): - """Creates a new instance of the Predictions class and copies the attributes. - - @return: A new instance of the Predictions class. - @rtype: Predictions - """ - new_obj = Predictions() - new_obj.predictions = [prediction.copy() for prediction in self.predictions] - new_obj.setSequenceNum(self.getSequenceNum()) - new_obj.setTimestamp(self.getTimestamp()) - new_obj.setTimestampDevice(self.getTimestampDevice()) - new_obj.setTransformation(self.transformation) - return new_obj - - @property - def predictions(self) -> list[Prediction]: - """Returns the predictions. - - @return: List of predictions. - @rtype: list[Prediction] - """ - return self._predictions - - @predictions.setter - def predictions(self, value: list[Prediction]): - """Sets the predictions. - - @param value: List of predicted values. - @type value: list[Prediction] - @raise TypeError: If value is not a list. - @raise ValueError: If each element is not of type Prediction. - """ - if not isinstance(value, list): - raise TypeError( - f"Predictions must be of type list, instead got {type(value)}." - ) - if not all(isinstance(item, Prediction) for item in value): - raise ValueError("Predictions must be a list of Prediction objects.") - self._predictions = value - - @property - def prediction(self) -> float: - """Returns the first prediction. Useful for single predictions. - - @return: The predicted value. - @rtype: float - """ - return self._predictions[0].prediction - - @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.ImgAnnotations: - """Returns the visualization message for the predictions. - - The message adds text representing the predictions to the right of the image. - """ - if self.transformation is None: - raise ValueError("Transformation must be set to get visualization message.") - w, h = self.transformation.getSize() - annotation_helper = AnnotationHelper() - annotation_sizes = AnnotationSizes(w, h) - - x_offset = 3 / w - y_offset = 3 / h - - for i, prediction in enumerate(self.predictions): - y_position = ( - y_offset - + (annotation_sizes.relative_font_size) - + i * (annotation_sizes.relative_font_size) - ) - annotation_helper.draw_text( - text=f"{prediction.prediction:.2f}", - position=(x_offset, y_position), - color=FONT_COLOR, - background_color=FONT_BACKGROUND_COLOR, - size=annotation_sizes.font_size, - ) - return annotation_helper.build( - timestamp=self.getTimestamp(), - sequence_num=self.getSequenceNum(), - ) diff --git a/depthai_nodes/message/utils/copy_message.py b/depthai_nodes/message/utils/copy_message.py index f78a6be9..52dc24d3 100644 --- a/depthai_nodes/message/utils/copy_message.py +++ b/depthai_nodes/message/utils/copy_message.py @@ -126,6 +126,7 @@ def _copy_keypoint(keypoint: dai.Keypoint) -> dai.Keypoint: keypoint_copy = _copy_metadata(keypoint) keypoint_copy.imageCoordinates = _copy_point3f(keypoint.imageCoordinates) keypoint_copy.confidence = keypoint.confidence + keypoint_copy.label = keypoint.label keypoint_copy.labelName = keypoint.labelName return keypoint_copy @@ -136,6 +137,51 @@ def _copy_rotated_rect(rotated_rect: dai.RotatedRect) -> dai.RotatedRect: rotated_rect_copy.angle = rotated_rect.angle return rotated_rect_copy + def _copy_beta_message(msg: dai.Buffer) -> dai.Buffer: + msg_copy = _copy_metadata(msg) + + if isinstance(msg, dai.beta.Classifications): + msg_copy.classes = list(msg.classes) + msg_copy.scores = msg.scores.copy() + elif isinstance(msg, dai.beta.Clusters): + clusters = [] + for cluster in msg.clusters: + cluster_copy = dai.beta.Cluster() + cluster_copy.label = cluster.label + cluster_copy.points = dai.VectorPoint2f( + [_copy_point2f(point) for point in cluster.points] + ) + clusters.append(cluster_copy) + msg_copy.clusters = clusters + elif isinstance(msg, dai.beta.Keypoints): + msg_copy.setKeypoints( + _copy_keypoints(msg.getKeypoints()), copy.deepcopy(msg.getEdges()) + ) + elif isinstance(msg, dai.beta.Lines): + lines = [] + for line in msg.lines: + line_copy = dai.beta.Line() + line_copy.startPoint = _copy_point2f(line.startPoint) + line_copy.endPoint = _copy_point2f(line.endPoint) + line_copy.confidence = line.confidence + lines.append(line_copy) + msg_copy.lines = lines + elif isinstance(msg, dai.beta.Map2D): + map_array = msg.getMap() + if map_array.size > 0: + msg_copy.setMap(map_array.copy()) + elif isinstance(msg, dai.beta.Predictions): + predictions = [] + for prediction in msg.predictions: + prediction_copy = dai.beta.Prediction() + prediction_copy.prediction = prediction.prediction + predictions.append(prediction_copy) + msg_copy.predictions = predictions + else: + raise TypeError(f"Unsupported beta message type {type(msg)}") + + return msg_copy + if isinstance(msg, dai.SegmentationMask): return _copy_segmentation_mask(msg) elif isinstance(msg, dai.ImgFrame): @@ -148,6 +194,18 @@ def _copy_rotated_rect(rotated_rect: dai.RotatedRect) -> dai.RotatedRect: return _copy_keypoints_list(msg) elif isinstance(msg, dai.Point2f): return _copy_point2f(msg) + elif isinstance( + msg, + ( + dai.beta.Classifications, + dai.beta.Clusters, + dai.beta.Keypoints, + dai.beta.Lines, + dai.beta.Map2D, + dai.beta.Predictions, + ), + ): + return _copy_beta_message(msg) else: # TODO: define logic for copying other message types raise TypeError(f"Copying of message type {type(msg)} is not supported.") diff --git a/depthai_nodes/node/README.md b/depthai_nodes/node/README.md index dfb2d811..cfc2c8f1 100644 --- a/depthai_nodes/node/README.md +++ b/depthai_nodes/node/README.md @@ -34,8 +34,8 @@ The parser classes listed below are the host-side implementations provided by de ### Classification -- `ClassificationParser`: General classification parser for postprocessing the output of a classification model and outputting the [`depthai_nodes.message.Classifications`](../message/README.md#classifications) message. -- `ClassificationSequenceParser`: Parser for models that predict the classes multiple times and return a list of predicted classes, where each item corresponds to the relative step in the sequence. In addition to time series classification, this parser can also be used for text recognition models where words can be interpreted as a sequence of characters (classes). It will output the [`depthai_nodes.message.Classifications`](../message/README.md#classifications) message. +- `ClassificationParser`: General classification parser for postprocessing the output of a classification model and outputting a native `dai.beta.Classifications` message. +- `ClassificationSequenceParser`: Parser for models that predict the classes multiple times and return a list of predicted classes, where each item corresponds to the relative step in the sequence. In addition to time series classification, this parser can also be used for text recognition models where words can be interpreted as a sequence of characters (classes). It outputs a native `dai.beta.Classifications` message. ### Segmentation @@ -49,11 +49,11 @@ The parser classes listed below are the host-side implementations provided by de ### Other -- `LaneDetectionParser`: Special parser for parsing the output of the Ultra-Fast-Lane-Detection model. It will output the [`depthai_nodes.message.Clusters`](../message/README.md#clusters) message. -- `MLSDParser`: Special parser for parsing the output of the MLSD model. It will output the [`depthai_nodes.message.Lines`](../message/README.md#lines) message. +- `LaneDetectionParser`: Special parser for parsing the output of the Ultra-Fast-Lane-Detection model. It outputs a native `dai.beta.Clusters` message. +- `MLSDParser`: Special parser for parsing the output of the MLSD model. It outputs a native `dai.beta.Lines` message. - `EmbeddingsParser`: Simple parser that will only forward the output of the neural network. It will output the `dai.NNData` message. -- `RegressionParser`: Special parser for parsing the output of the regression models. It will output the [`depthai_nodes.message.Predictions`](../message/README.md#predictions) message. -- `MapOutputParser`: Special parser for parsing the output of the model that produces a map (like depth estimation). It will output the [`depthai_nodes.message.Map2D`](../message/README.md#map2d) message. +- `RegressionParser`: Special parser for parsing the output of the regression models. It outputs a native `dai.beta.Predictions` message. +- `MapOutputParser`: Special parser for parsing the output of the model that produces a map (like depth estimation). It outputs a native `dai.beta.Map2D` message. - `ImageOutputParser`: Special parser for parsing the output of the model that produces an image (like super-resolution models). It will output the `dai.ImgFrame` message. ## Utility & Helper Nodes diff --git a/depthai_nodes/node/apply_colormap.py b/depthai_nodes/node/apply_colormap.py index 06b60f44..8e09636e 100644 --- a/depthai_nodes/node/apply_colormap.py +++ b/depthai_nodes/node/apply_colormap.py @@ -2,7 +2,6 @@ import depthai as dai import numpy as np -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 +24,7 @@ class ApplyColormap(BaseHostNode): Inputs ------ - frame : dai.ImgFrame | Map2D | dai.ImgDetections | dai.SegmentationMask + frame : dai.ImgFrame | dai.beta.Map2D | dai.ImgDetections | dai.SegmentationMask Input message containing a 2D array to be colorized. Outputs @@ -133,8 +132,8 @@ def _get_input_map(msg: dai.Buffer) -> np.ndarray: mask = msg_copy.getCvMask() return np.where(mask == 255, 0, mask + 1) - if isinstance(msg_copy, Map2D): - return msg_copy.map + if isinstance(msg_copy, dai.beta.Map2D): + return msg_copy.getMap() if isinstance(msg_copy, dai.ImgDetections): mask = msg_copy.getCvSegmentationMask() @@ -144,7 +143,8 @@ 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, dai.SegmentationMask, Map2D and dai.ImgDetections." + "dai.ImgFrame, dai.SegmentationMask, dai.beta.Map2D and " + "dai.ImgDetections." ) def _colorize(self, input_map: np.ndarray) -> np.ndarray: diff --git a/depthai_nodes/node/parsers/classification.py b/depthai_nodes/node/parsers/classification.py index f6fde7a8..0cb94a5d 100644 --- a/depthai_nodes/node/parsers/classification.py +++ b/depthai_nodes/node/parsers/classification.py @@ -25,7 +25,7 @@ class ClassificationParser(BaseParser): Output Message/s ---------------- - **Type** : Classifications(dai.Buffer) + **Type** : dai.beta.Classifications **Description**: An object with attributes `classes` and `scores`. `classes` is a list of classes, sorted in descending order of scores. `scores` is a list of corresponding scores. """ diff --git a/depthai_nodes/node/parsers/classification_sequence.py b/depthai_nodes/node/parsers/classification_sequence.py index 88966499..b0a43c25 100644 --- a/depthai_nodes/node/parsers/classification_sequence.py +++ b/depthai_nodes/node/parsers/classification_sequence.py @@ -35,7 +35,7 @@ class ClassificationSequenceParser(ClassificationParser): Output Message/s ---------------- - **Type**: Classifications(dai.Buffer) + **Type**: dai.beta.Classifications **Description**: An object with attributes `classes` and `scores`. `classes` is a list containing the predicted classes. `scores` is a list of corresponding probability scores. diff --git a/depthai_nodes/node/parsers/hrnet.py b/depthai_nodes/node/parsers/hrnet.py index db541707..5f9ecfc5 100644 --- a/depthai_nodes/node/parsers/hrnet.py +++ b/depthai_nodes/node/parsers/hrnet.py @@ -26,7 +26,7 @@ class HRNetParser(KeypointParser): Output Message/s ---------------- - **Type**: Keypoints + **Type**: dai.beta.Keypoints **Description**: Output containing detected body keypoints. """ diff --git a/depthai_nodes/node/parsers/keypoints.py b/depthai_nodes/node/parsers/keypoints.py index 4a3a97b9..d3990d50 100644 --- a/depthai_nodes/node/parsers/keypoints.py +++ b/depthai_nodes/node/parsers/keypoints.py @@ -30,7 +30,7 @@ class KeypointParser(BaseParser): Output Message/s ---------------- - **Type**: Keypoints + **Type**: dai.beta.Keypoints **Description**: Output containing 2D or 3D keypoints. diff --git a/depthai_nodes/node/parsers/lane_detection.py b/depthai_nodes/node/parsers/lane_detection.py index 9084f849..dcebe94a 100644 --- a/depthai_nodes/node/parsers/lane_detection.py +++ b/depthai_nodes/node/parsers/lane_detection.py @@ -30,7 +30,7 @@ class LaneDetectionParser(BaseParser): Output Message/s ---------------- - **Type**: Clusters + **Type**: dai.beta.Clusters **Description**: Detected lanes represented as clusters of points. Error Handling diff --git a/depthai_nodes/node/parsers/map_output.py b/depthai_nodes/node/parsers/map_output.py index d86f0ed8..e0019cce 100644 --- a/depthai_nodes/node/parsers/map_output.py +++ b/depthai_nodes/node/parsers/map_output.py @@ -20,9 +20,9 @@ class MapOutputParser(BaseParser): Output Message/s ---------------- - **Type**: Map2D + **Type**: dai.beta.Map2D - **Description**: Density message containing the density map. The density map is represented with Map2D object. + **Description**: Density message containing the density map as a native dai.beta.Map2D object. """ def __init__( diff --git a/depthai_nodes/node/parsers/mlsd.py b/depthai_nodes/node/parsers/mlsd.py index 5851a501..0fbd6bd6 100644 --- a/depthai_nodes/node/parsers/mlsd.py +++ b/depthai_nodes/node/parsers/mlsd.py @@ -28,9 +28,9 @@ class MLSDParser(BaseParser): Output Message/s ---------------- - **Type**: LineDetections + **Type**: dai.beta.Lines - **Description**: LineDetections message containing detected lines and confidence scores. + **Description**: Native message containing detected lines and confidence scores. """ def __init__( diff --git a/depthai_nodes/node/parsers/regression.py b/depthai_nodes/node/parsers/regression.py index 9403fe99..89c544e2 100644 --- a/depthai_nodes/node/parsers/regression.py +++ b/depthai_nodes/node/parsers/regression.py @@ -21,7 +21,7 @@ class RegressionParser(BaseParser): Output Message/s ---------------- - **Type**: Predictions + **Type**: dai.beta.Predictions **Description**: Message containing the prediction(s). """ diff --git a/depthai_nodes/node/parsers/superanimal_landmarker.py b/depthai_nodes/node/parsers/superanimal_landmarker.py index ae52fad8..b0b2eb02 100644 --- a/depthai_nodes/node/parsers/superanimal_landmarker.py +++ b/depthai_nodes/node/parsers/superanimal_landmarker.py @@ -30,7 +30,7 @@ class SuperAnimalParser(KeypointParser): Output Message/s ---------------- - **Type**: Keypoints + **Type**: dai.beta.Keypoints **Description**: Output containing detected keypoints that exceed the confidence threshold. """ diff --git a/depthai_nodes/node/utils/message_remapping.py b/depthai_nodes/node/utils/message_remapping.py index a694264f..834468fa 100644 --- a/depthai_nodes/node/utils/message_remapping.py +++ b/depthai_nodes/node/utils/message_remapping.py @@ -1,14 +1,6 @@ -import cv2 import depthai as dai -import numpy as np -from depthai_nodes.message.classification import Classifications -from depthai_nodes.message.clusters import Cluster, Clusters -from depthai_nodes.message.keypoints import Keypoints -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.node.utils.util_constants import UNASSIGNED_MASK_LABEL, GMessage +from depthai_nodes.node.utils.util_constants import GMessage def remap_message( @@ -16,258 +8,30 @@ def remap_message( from_transformation: dai.ImgTransformation, to_transformation: dai.ImgTransformation, ) -> GMessage: - if isinstance(message, dai.ImgDetections): - return remap_img_detections(from_transformation, to_transformation, message) - elif isinstance(message, Keypoints): - return remap_keypoints(from_transformation, to_transformation, message) - 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) - elif isinstance(message, Map2D): - return remap_map2d(from_transformation, to_transformation, message) - elif isinstance(message, Lines): - return remap_lines(from_transformation, to_transformation, message) - elif isinstance(message, Predictions): - return remap_predictions(from_transformation, to_transformation, message) - elif isinstance(message, Classifications): - return remap_classifications(from_transformation, to_transformation, message) - else: + """Remap a transformable DepthAI message to a target image transformation. + + ``from_transformation`` remains part of the API for callers which resolve the + source transformation explicitly. Native messages carry that transformation + themselves, so remapping is delegated to their ``transformTo`` implementation. + """ + + if not isinstance( + message, + ( + dai.ImgDetections, + dai.SegmentationMask, + dai.beta.Keypoints, + dai.beta.Clusters, + dai.beta.Map2D, + dai.beta.Lines, + dai.beta.Predictions, + dai.beta.Classifications, + ), + ): raise TypeError( f"Cannot remap message. Unsupported message type: {type(message)}" ) - -def remap_img_detections( - from_transformation: dai.ImgTransformation, - to_transformation: dai.ImgTransformation, - detections: dai.ImgDetections, -) -> dai.ImgDetections: - new_detections = dai.ImgDetections() - new_detections.detections = [ - remap_img_detection(from_transformation, to_transformation, det) - for det in detections.detections - ] - if ( - detections.getCvSegmentationMask() is not None - and detections.getCvSegmentationMask().size > 0 - ): - new_mask = remap_segmentation_mask_array( - from_transformation, to_transformation, detections.getCvSegmentationMask() - ) - new_detections.setCvSegmentationMask(new_mask) - return new_detections - - -def remap_segmentation_mask_array( - from_transformation: dai.ImgTransformation, - to_transformation: dai.ImgTransformation, - segmentation_mask: np.ndarray, -) -> np.ndarray: - 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=border_value, # type: ignore - ) - return new_mask - - -def remap_segmentation_mask( - from_transformation: dai.ImgTransformation, - to_transformation: dai.ImgTransformation, - 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 - - -def remap_img_detection( - from_transformation: dai.ImgTransformation, - to_transformation: dai.ImgTransformation, - img_detection: dai.ImgDetection, -) -> dai.ImgDetection: - new_det = dai.ImgDetection() - new_rect = from_transformation.remapRectTo( - to_transformation, img_detection.getBoundingBox() - ) - new_det.setBoundingBox(new_rect) - new_det.label = img_detection.label - new_det.labelName = img_detection.labelName - new_det.confidence = img_detection.confidence - - new_kpts_list = [] - for kpt in img_detection.getKeypoints(): - new_kpt = remap_keypoint(from_transformation, to_transformation, kpt) - new_kpts_list.append(new_kpt) - new_kpts = dai.KeypointsList() - new_kpts.setEdges(img_detection.getEdges()) - new_kpts.setKeypoints(new_kpts_list) - new_det.setKeypoints(new_kpts) - return new_det - - -def remap_keypoint( - from_transformation: dai.ImgTransformation, - to_transformation: dai.ImgTransformation, - keypoint: dai.Keypoint, -) -> dai.Keypoint: - new_kpt = dai.Keypoint() - transformed = from_transformation.remapPointTo( - to_transformation, - dai.Point2f(keypoint.imageCoordinates.x, keypoint.imageCoordinates.y), - ) - new_kpt.imageCoordinates.x = transformed.x - new_kpt.imageCoordinates.y = transformed.y - new_kpt.imageCoordinates.z = keypoint.imageCoordinates.z - new_kpt.confidence = keypoint.confidence - new_kpt.labelName = keypoint.labelName - new_kpt.label = keypoint.label - return new_kpt - - -def remap_native_keypoints( - from_transformation: dai.ImgTransformation, - to_transformation: dai.ImgTransformation, - keypoints: dai.KeypointsList, -) -> dai.KeypointsList: - new_kpts_list = [] - for kpt in keypoints.getKeypoints(): - new_kpt = remap_keypoint(from_transformation, to_transformation, kpt) - new_kpts_list.append(new_kpt) - new_kpts = dai.KeypointsList() - new_kpts.setKeypoints(new_kpts_list) - new_kpts.setEdges(keypoints.getEdges()) - return new_kpts - - -def remap_keypoints( - from_transformation: dai.ImgTransformation, - to_transformation: dai.ImgTransformation, - keypoints: Keypoints, -) -> Keypoints: - new_kpts = Keypoints() - new_kpts.keypoints_list = remap_native_keypoints( - from_transformation, to_transformation, keypoints.keypoints_list - ) - new_kpts.setSequenceNum(keypoints.getSequenceNum()) - new_kpts.setTimestamp(keypoints.getTimestamp()) - new_kpts.setTimestampDevice(keypoints.getTimestampDevice()) - new_kpts.setTransformation(keypoints.getTransformation()) - return new_kpts - - -def remap_clusters( - from_transformation: dai.ImgTransformation, - to_transformation: dai.ImgTransformation, - clusters: Clusters, -) -> Clusters: - new_clusters = Clusters() - new_clusters.clusters = [ - remap_cluster(from_transformation, to_transformation, cluster) - for cluster in clusters.clusters - ] - return new_clusters - - -def remap_cluster( - from_transformation: dai.ImgTransformation, - to_transformation: dai.ImgTransformation, - cluster: Cluster, -) -> Cluster: - new_cluster = Cluster() - new_cluster.label = cluster.label - new_cluster.points = [ - from_transformation.remapPointTo(to_transformation, pt) for pt in cluster.points - ] - return new_cluster - - -def remap_map2d( - from_transformation: dai.ImgTransformation, - to_transformation: dai.ImgTransformation, - map2d: Map2D, -) -> Map2D: - new_map2d = Map2D() - new_map_arr = remap_segmentation_mask_array( - from_transformation, to_transformation, map2d.map - ) - new_map2d.map = new_map_arr - return new_map2d - - -def remap_lines( - from_transformation: dai.ImgTransformation, - to_transformation: dai.ImgTransformation, - lines: Lines, -) -> Lines: - new_lines = Lines() - new_lines.lines = [ - remap_line(from_transformation, to_transformation, line) for line in lines.lines - ] - return new_lines - - -def remap_line( - from_transformation: dai.ImgTransformation, - to_transformation: dai.ImgTransformation, - line: Line, -) -> Line: - new_line = Line() - new_line.confidence = line.confidence - new_line.start_point = from_transformation.remapPointTo( - to_transformation, line.start_point - ) - new_line.end_point = from_transformation.remapPointTo( - to_transformation, line.end_point - ) - return new_line - - -def remap_predictions( - from_transformation: dai.ImgTransformation, - to_transformation: dai.ImgTransformation, - predictions: Predictions, -) -> Predictions: - new_predictions = Predictions() - new_predictions.predictions = [ - remap_prediction(from_transformation, to_transformation, prediction) - for prediction in predictions.predictions - ] - return new_predictions - - -def remap_prediction( - from_transformation: dai.ImgTransformation, - to_transformation: dai.ImgTransformation, - prediction: Prediction, -) -> Prediction: - new_prediction = Prediction() - new_prediction.prediction = prediction.prediction - return new_prediction - - -def remap_classifications( - from_transformation: dai.ImgTransformation, - to_transformation: dai.ImgTransformation, - classifications: Classifications, -) -> Classifications: - new_classifications = Classifications() - new_classifications.classes = classifications.classes - new_classifications.scores = classifications.scores - return new_classifications + if message.getTransformation() is None: + message.setTransformation(from_transformation) + return message.transformTo(to_transformation) diff --git a/depthai_nodes/node/utils/util_constants.py b/depthai_nodes/node/utils/util_constants.py index ed1025c7..819b4b32 100644 --- a/depthai_nodes/node/utils/util_constants.py +++ b/depthai_nodes/node/utils/util_constants.py @@ -2,22 +2,14 @@ import depthai as dai -from depthai_nodes.message.classification import Classifications -from depthai_nodes.message.clusters import Clusters -from depthai_nodes.message.keypoints import Keypoints -from depthai_nodes.message.lines import Lines -from depthai_nodes.message.map import Map2D -from depthai_nodes.message.prediction import Predictions - GMessage = TypeVar( "GMessage", bound=dai.ImgDetections - | Keypoints + | dai.beta.Keypoints | dai.SegmentationMask - | Clusters - | Map2D - | Lines - | Predictions - | Classifications, + | dai.beta.Clusters + | dai.beta.Map2D + | dai.beta.Lines + | dai.beta.Predictions + | dai.beta.Classifications, ) -UNASSIGNED_MASK_LABEL = -1 diff --git a/tests/stability_tests/check_messages.py b/tests/stability_tests/check_messages.py index 02c21f2a..dba74e2f 100644 --- a/tests/stability_tests/check_messages.py +++ b/tests/stability_tests/check_messages.py @@ -4,15 +4,6 @@ import depthai as dai import numpy as np -from depthai_nodes import ( - Classifications, - Clusters, - Keypoints, - Lines, - Map2D, - Predictions, -) - from .utils import extract_main_slug @@ -23,7 +14,9 @@ def load_expected_output(model: str, parser: str) -> dict[str, Any]: def check_classification_msg( - message: Classifications, expected_output: dict[str, Any], verbose: bool = False + message: dai.beta.Classifications, + expected_output: dict[str, Any], + verbose: bool = False, ): """ Expected output format: @@ -35,22 +28,26 @@ def check_classification_msg( } """ assert isinstance( - message, Classifications - ), f"The message is not a Classifications. Got {type(message)}." + message, dai.beta.Classifications + ), f"The message is not a dai.beta.Classifications. Got {type(message)}." if verbose: print( - f"Expected top class: {expected_output['class']}, predicted top class: {message.top_class}" + f"Expected top class: {expected_output['class']}, predicted top class: {message.getTopClass()}" ) print( - f"Expected top score: {expected_output['score']}, predicted top score: {message.top_score}" + f"Expected top score: {expected_output['score']}, predicted top score: {message.getTopScore()}" ) - assert message.top_class == expected_output["class"] - np.testing.assert_allclose(message.top_score, expected_output["score"], rtol=1e-2) + assert message.getTopClass() == expected_output["class"] + np.testing.assert_allclose( + message.getTopScore(), expected_output["score"], rtol=1e-2 + ) def check_classification_sequence_msg( - message: Classifications, expected_output: dict[str, Any], verbose: bool = False + message: dai.beta.Classifications, + expected_output: dict[str, Any], + verbose: bool = False, ): """ Expected output format: @@ -60,8 +57,8 @@ def check_classification_sequence_msg( "class": ['HELLO'] """ assert isinstance( - message, Classifications - ), f"The message is not a Classifications. Got {type(message)}." + message, dai.beta.Classifications + ), f"The message is not a dai.beta.Classifications. Got {type(message)}." if verbose: print( @@ -148,7 +145,7 @@ def check_segmentation_msg( def check_keypoints_msg( - message: Keypoints, + message: dai.beta.Keypoints, expected_output: dict[str, Any], verbose: bool = False, ): @@ -160,8 +157,8 @@ def check_keypoints_msg( "keypoints": [[0.1, 0.2], ...] """ assert isinstance( - message, Keypoints - ), f"The message is not a Keypoints. Got {type(message)}." + message, dai.beta.Keypoints + ), f"The message is not a dai.beta.Keypoints. Got {type(message)}." keypoints = [ [kp.imageCoordinates.x, kp.imageCoordinates.y] for kp in message.getKeypoints() @@ -209,7 +206,9 @@ def check_image_msg( def check_cluster_msg( - message: Clusters, expected_output: dict[str, Any], verbose: bool = False + message: dai.beta.Clusters, + expected_output: dict[str, Any], + verbose: bool = False, ): """ Expected output format: @@ -219,8 +218,8 @@ def check_cluster_msg( "clusters": [[[0.1, 0.2], ...]] """ assert isinstance( - message, Clusters - ), f"The message is not a Clusters. Got {type(message)}." + message, dai.beta.Clusters + ), f"The message is not a dai.beta.Clusters. Got {type(message)}." clusters = message.clusters expected_clusters = expected_output["clusters"] @@ -243,7 +242,9 @@ def check_cluster_msg( def check_map_msg( - message: Map2D, expected_output: dict[str, Any], verbose: bool = False + message: dai.beta.Map2D, + expected_output: dict[str, Any], + verbose: bool = False, ): """ Expected output format: @@ -254,10 +255,10 @@ def check_map_msg( } """ assert isinstance( - message, Map2D - ), f"The message is not a Map2D. Got {type(message)}." + message, dai.beta.Map2D + ), f"The message is not a dai.beta.Map2D. Got {type(message)}." - map_tensor = message.map + map_tensor = message.getMap() expected_map = expected_output["map"] if verbose: print( @@ -385,7 +386,9 @@ def check_detection_msg( def check_line_msg( - message: Lines, expected_output: dict[str, Any], verbose: bool = False + message: dai.beta.Lines, + expected_output: dict[str, Any], + verbose: bool = False, ): """ Expected output format: @@ -403,16 +406,16 @@ def check_line_msg( } """ assert isinstance( - message, Lines - ), f"The message is not a Lines. Got {type(message)}." + message, dai.beta.Lines + ), f"The message is not a dai.beta.Lines. Got {type(message)}." expected_lines: list[dict[str, Any]] = expected_output["lines"] predicted_lines = [] for line in message.lines: line_dict = { "confidence": line.confidence, - "start_point": [line.start_point.x, line.start_point.y], - "end_point": [line.end_point.x, line.end_point.y], + "start_point": [line.startPoint.x, line.startPoint.y], + "end_point": [line.endPoint.x, line.endPoint.y], } predicted_lines.append(line_dict) @@ -438,7 +441,9 @@ def check_line_msg( def check_regression_msg( - message: Predictions, expected_output: dict[str, Any], verbose: bool = False + message: dai.beta.Predictions, + expected_output: dict[str, Any], + verbose: bool = False, ): """ Expected output format: @@ -449,8 +454,8 @@ def check_regression_msg( } """ assert isinstance( - message, Predictions - ), f"The message is not a Predictions. Got {type(message)}." + message, dai.beta.Predictions + ), f"The message is not a dai.beta.Predictions. Got {type(message)}." predictions = message.predictions predictions = np.array([pred.prediction for pred in predictions]) diff --git a/tests/unittests/test_creators/test_classification_sequence.py b/tests/unittests/test_creators/test_classification_sequence.py index b4220070..86080dda 100644 --- a/tests/unittests/test_creators/test_classification_sequence.py +++ b/tests/unittests/test_creators/test_classification_sequence.py @@ -1,7 +1,7 @@ +import depthai as dai import numpy as np import pytest -from depthai_nodes import Classifications from depthai_nodes.message.creators import ( create_classification_sequence_message, ) @@ -21,7 +21,7 @@ def test_valid_input(): message = create_classification_sequence_message(CLASSES, SCORES) - assert isinstance(message, Classifications) + assert isinstance(message, dai.beta.Classifications) assert message.classes == CLASSES assert np.array_equal(message.scores, np.array(CORRECT_SCORES, dtype=np.float32)) diff --git a/tests/unittests/test_creators/test_classifications.py b/tests/unittests/test_creators/test_classifications.py index c99bb6c2..349bb7bc 100644 --- a/tests/unittests/test_creators/test_classifications.py +++ b/tests/unittests/test_creators/test_classifications.py @@ -1,7 +1,7 @@ +import depthai as dai import numpy as np import pytest -from depthai_nodes import Classifications from depthai_nodes.message.creators import ( create_classification_message, ) @@ -13,7 +13,7 @@ def test_valid_input(): message = create_classification_message(CLASSES, SCORES) - assert isinstance(message, Classifications) + assert isinstance(message, dai.beta.Classifications) assert message.classes == ["cat", "dog", "bird"] assert np.array_equal(message.scores, np.array(SCORES, dtype=np.float32)) @@ -53,7 +53,7 @@ def test_very_small_scores(): message = create_classification_message(CLASSES, scores) - assert isinstance(message, Classifications) + assert isinstance(message, dai.beta.Classifications) assert message.classes == ["bird", "cat", "dog"] assert np.all( message.scores == np.array([1 - 2e-10, 1e-10, 1e-10], dtype=np.float32) diff --git a/tests/unittests/test_creators/test_clusters.py b/tests/unittests/test_creators/test_clusters.py index 5759feab..afa2bb87 100644 --- a/tests/unittests/test_creators/test_clusters.py +++ b/tests/unittests/test_creators/test_clusters.py @@ -1,6 +1,6 @@ +import depthai as dai import pytest -from depthai_nodes import Cluster, Clusters from depthai_nodes.message.creators import ( create_cluster_message, ) @@ -10,9 +10,9 @@ def test_valid_input(): clusters = [[[1.0, 2.0], [3.0, 4.0]], [[5.0, 6.0], [7.0, 8.0]]] message = create_cluster_message(clusters) - assert isinstance(message, Clusters) + assert isinstance(message, dai.beta.Clusters) assert len(message.clusters) == 2 - assert all(isinstance(cluster, Cluster) for cluster in message.clusters) + assert all(isinstance(cluster, dai.beta.Cluster) for cluster in message.clusters) assert message.clusters[0].label == 0 assert message.clusters[1].label == 1 assert len(message.clusters[0].points) == 2 @@ -58,7 +58,7 @@ def test_empty_clusters(): clusters = [] message = create_cluster_message(clusters) - assert isinstance(message, Clusters) + assert isinstance(message, dai.beta.Clusters) assert len(message.clusters) == 0 @@ -66,6 +66,6 @@ def test_empty_cluster(): clusters = [[]] message = create_cluster_message(clusters) - assert isinstance(message, Clusters) + assert isinstance(message, dai.beta.Clusters) assert len(message.clusters) == 1 assert len(message.clusters[0].points) == 0 diff --git a/tests/unittests/test_creators/test_keypoints.py b/tests/unittests/test_creators/test_keypoints.py index 6dd9e572..425bfa70 100644 --- a/tests/unittests/test_creators/test_keypoints.py +++ b/tests/unittests/test_creators/test_keypoints.py @@ -1,7 +1,6 @@ import depthai as dai import pytest -from depthai_nodes import Keypoints from depthai_nodes.message.creators import create_keypoints_message KPTS = [[0.1, 0.2], [0.3, 0.4]] @@ -11,8 +10,8 @@ def test_valid_2d_keypoints(): message = create_keypoints_message(KPTS, SCORES) - assert isinstance(message, Keypoints) - assert isinstance(message.keypoints_list, dai.KeypointsList) + assert isinstance(message, dai.beta.Keypoints) + assert isinstance(message.keypointsList, dai.KeypointsList) assert len(message.getKeypoints()) == 2 for i, kp in enumerate(message.getKeypoints()): assert kp.imageCoordinates.x == pytest.approx(KPTS[i][0]) @@ -26,7 +25,7 @@ def test_valid_3d_keypoints(): scores = [0.9, 0.8] message = create_keypoints_message(keypoints, scores) - assert isinstance(message, Keypoints) + assert isinstance(message, dai.beta.Keypoints) assert len(message.getKeypoints()) == 2 for i, kp in enumerate(message.getKeypoints()): assert kp.imageCoordinates.x == pytest.approx(keypoints[i][0]) @@ -38,7 +37,7 @@ def test_valid_3d_keypoints(): def test_valid_keypoints_no_scores(): message = create_keypoints_message(KPTS) - assert isinstance(message, Keypoints) + assert isinstance(message, dai.beta.Keypoints) assert len(message.getKeypoints()) == 2 for kp in message.getKeypoints(): assert kp.confidence == -1 diff --git a/tests/unittests/test_creators/test_line_detections.py b/tests/unittests/test_creators/test_line_detections.py index 36376b1c..1390cf61 100644 --- a/tests/unittests/test_creators/test_line_detections.py +++ b/tests/unittests/test_creators/test_line_detections.py @@ -1,7 +1,7 @@ +import depthai as dai import numpy as np import pytest -from depthai_nodes import Line, Lines from depthai_nodes.message.creators import create_line_detection_message LINE = np.array([[0.1, 0.2, 0.3, 0.4]]) @@ -13,15 +13,15 @@ def test_valid_input(): scores = np.array([0.9, 0.8]) message = create_line_detection_message(lines, scores) - assert isinstance(message, Lines) + assert isinstance(message, dai.beta.Lines) assert len(message.lines) == 2 for i, line in enumerate(message.lines): - assert isinstance(line, Line) - assert np.allclose(line.start_point.x, lines[i][0], atol=1e-3) - assert np.allclose(line.start_point.y, lines[i][1], atol=1e-3) - assert np.allclose(line.end_point.x, lines[i][2], atol=1e-3) - assert np.allclose(line.end_point.y, lines[i][3], atol=1e-3) + assert isinstance(line, dai.beta.Line) + assert np.allclose(line.startPoint.x, lines[i][0], atol=1e-3) + assert np.allclose(line.startPoint.y, lines[i][1], atol=1e-3) + assert np.allclose(line.endPoint.x, lines[i][2], atol=1e-3) + assert np.allclose(line.endPoint.y, lines[i][3], atol=1e-3) assert np.allclose(line.confidence, scores[i], atol=1e-3) @@ -30,7 +30,7 @@ def test_empty_lines(): scores = np.array([]) message = create_line_detection_message(lines, scores) - assert isinstance(message, Lines) + assert isinstance(message, dai.beta.Lines) assert len(message.lines) == 0 diff --git a/tests/unittests/test_creators/test_map.py b/tests/unittests/test_creators/test_map.py index ce027cf0..ec1c37b3 100644 --- a/tests/unittests/test_creators/test_map.py +++ b/tests/unittests/test_creators/test_map.py @@ -1,7 +1,7 @@ +import depthai as dai import numpy as np import pytest -from depthai_nodes import Map2D from depthai_nodes.message.creators import create_map_message MAP_ARRAY = np.random.rand(1, 480, 640).astype(np.float32) @@ -10,39 +10,39 @@ def test_valid_2d_input(): message = create_map_message(MAP_ARRAY[0]) - assert isinstance(message, Map2D) - assert message.map.shape == (480, 640) - assert message.map.dtype == np.float32 - assert np.allclose(message.map, MAP_ARRAY[0]) + assert isinstance(message, dai.beta.Map2D) + assert message.getMap().shape == (480, 640) + assert message.getMap().dtype == np.float32 + assert np.allclose(message.getMap(), MAP_ARRAY[0]) def test_valid_3d_input_nhw(): message = create_map_message(MAP_ARRAY) - assert isinstance(message, Map2D) - assert message.map.shape == (480, 640) - assert message.map.dtype == np.float32 - assert np.allclose(message.map, MAP_ARRAY[0]) + assert isinstance(message, dai.beta.Map2D) + assert message.getMap().shape == (480, 640) + assert message.getMap().dtype == np.float32 + assert np.allclose(message.getMap(), MAP_ARRAY[0]) def test_valid_3d_input_hwn(): message = create_map_message(MAP_ARRAY.transpose(1, 2, 0)) - assert isinstance(message, Map2D) - assert message.map.shape == (480, 640) - assert message.map.dtype == np.float32 - assert np.allclose(message.map, MAP_ARRAY[0]) + assert isinstance(message, dai.beta.Map2D) + assert message.getMap().shape == (480, 640) + assert message.getMap().dtype == np.float32 + assert np.allclose(message.getMap(), MAP_ARRAY[0]) def test_min_max_scaling(): map_array = MAP_ARRAY[0] * 100 message = create_map_message(map_array, min_max_scaling=True) - assert isinstance(message, Map2D) - assert message.map.shape == (480, 640) - assert message.map.dtype == np.float32 - assert np.all(message.map >= 0) and np.all(message.map <= 1) - assert np.allclose(message.map, MAP_ARRAY[0], atol=1e-3) + assert isinstance(message, dai.beta.Map2D) + assert message.getMap().shape == (480, 640) + assert message.getMap().dtype == np.float32 + assert np.all(message.getMap() >= 0) and np.all(message.getMap() <= 1) + assert np.allclose(message.getMap(), MAP_ARRAY[0], atol=1e-3) def test_invalid_type(): @@ -64,7 +64,7 @@ def test_valid_input_non_float(): map_array = np.random.randint(0, 256, (480, 640), dtype=np.uint8) message = create_map_message(map_array) - assert isinstance(message, Map2D) - assert message.map.shape == (480, 640) - assert message.map.dtype == np.float32 - assert np.allclose(message.map, map_array, atol=1e-3) + assert isinstance(message, dai.beta.Map2D) + assert message.getMap().shape == (480, 640) + assert message.getMap().dtype == np.float32 + assert np.allclose(message.getMap(), map_array, atol=1e-3) diff --git a/tests/unittests/test_creators/test_regression.py b/tests/unittests/test_creators/test_regression.py index febcfc10..808242d0 100644 --- a/tests/unittests/test_creators/test_regression.py +++ b/tests/unittests/test_creators/test_regression.py @@ -1,6 +1,6 @@ +import depthai as dai import pytest -from depthai_nodes import Prediction, Predictions from depthai_nodes.message.creators import create_regression_message @@ -8,19 +8,19 @@ def test_valid_input(): predictions = [0.1, 0.2, 0.3] message = create_regression_message(predictions) - assert isinstance(message, Predictions) + assert isinstance(message, dai.beta.Predictions) assert len(message.predictions) == 3 - assert all(isinstance(pred, Prediction) for pred in message.predictions) - assert message.predictions[0].prediction == 0.1 - assert message.predictions[1].prediction == 0.2 - assert message.predictions[2].prediction == 0.3 + assert all(isinstance(pred, dai.beta.Prediction) for pred in message.predictions) + assert message.predictions[0].prediction == pytest.approx(0.1) + assert message.predictions[1].prediction == pytest.approx(0.2) + assert message.predictions[2].prediction == pytest.approx(0.3) def test_empty_list(): predictions = [] message = create_regression_message(predictions) - assert isinstance(message, Predictions) + assert isinstance(message, dai.beta.Predictions) assert len(message.predictions) == 0 diff --git a/tests/unittests/test_messages/test_classifications_msg.py b/tests/unittests/test_messages/test_classifications_msg.py deleted file mode 100644 index 781009c4..00000000 --- a/tests/unittests/test_messages/test_classifications_msg.py +++ /dev/null @@ -1,69 +0,0 @@ -import depthai as dai -import numpy as np -import pytest - -from depthai_nodes import Classifications - - -@pytest.fixture -def classifications(): - return Classifications() - - -def test_initialization(classifications: Classifications): - assert classifications.classes == [] - assert np.array_equal(classifications.scores, np.array([])) - assert classifications.transformation is None - - -def test_set_classes(classifications: Classifications): - classes = ["cat", "dog", "bird"] - classifications.classes = classes - assert classifications.classes == classes - - with pytest.raises(TypeError): - classifications.classes = "not a list" - - with pytest.raises(ValueError): - classifications.classes = ["cat", 123, "bird"] - - -def test_set_scores(classifications: Classifications): - scores = np.array([0.9, 0.05, 0.05], dtype=np.float32) - classifications.scores = scores - assert np.array_equal(classifications.scores, scores) - - with pytest.raises(TypeError): - classifications.scores = "not an array" - - with pytest.raises(ValueError): - classifications.scores = np.array([[0.9, 0.05, 0.05]], dtype=np.float32) - - with pytest.raises(ValueError): - classifications.scores = np.array([0.9, 0.05, "not a float"], dtype=object) - - -def test_top_class(classifications: Classifications): - classes = ["cat", "dog", "bird"] - classifications.classes = classes - assert classifications.top_class == "cat" - - -def test_top_score(classifications: Classifications): - scores = np.array([0.9, 0.05, 0.05], dtype=np.float32) - classifications.scores = scores - assert np.allclose(classifications.top_score, 0.9, atol=1e-3) - - -def test_set_transformation(classifications: Classifications): - transformation = dai.ImgTransformation() - classifications.transformation = transformation - assert classifications.transformation == transformation - - with pytest.raises(TypeError): - classifications.transformation = "not a dai.ImgTransformation" - - -def test_set_transformation_none(classifications: Classifications): - classifications.transformation = None - assert classifications.transformation is None diff --git a/tests/unittests/test_messages/test_clusters_msg.py b/tests/unittests/test_messages/test_clusters_msg.py deleted file mode 100644 index 910f9e08..00000000 --- a/tests/unittests/test_messages/test_clusters_msg.py +++ /dev/null @@ -1,72 +0,0 @@ -import depthai as dai -import pytest - -from depthai_nodes import Cluster, Clusters - - -@pytest.fixture -def cluster(): - return Cluster() - - -@pytest.fixture -def clusters(): - return Clusters() - - -def test_cluster_initialization(cluster: Cluster): - assert cluster.label is None - assert cluster.points == [] - - -def test_cluster_set_label(cluster: Cluster): - cluster.label = 1 - assert cluster.label == 1 - - with pytest.raises(TypeError): - cluster.label = "not an int" - - -def test_cluster_set_points(cluster: Cluster): - points = [dai.Point2f(0.1, 0.2), dai.Point2f(0.3, 0.4)] - cluster.points = points - assert cluster.points == points - - with pytest.raises(TypeError): - cluster.points = "not a list" - - with pytest.raises(ValueError): - cluster.points = [dai.Point2f(0.1, 0.2), "not a Point2f"] - - -def test_clusters_initialization(clusters: Clusters): - assert clusters.clusters == [] - assert clusters.transformation is None - - -def test_clusters_set_clusters(clusters: Clusters): - cluster1 = Cluster() - cluster2 = Cluster() - clusters_list = [cluster1, cluster2] - clusters.clusters = clusters_list - assert clusters.clusters == clusters_list - - with pytest.raises(TypeError): - clusters.clusters = "not a list" - - with pytest.raises(ValueError): - clusters.clusters = [cluster1, "not a Cluster"] - - -def test_clusters_set_transformation(clusters: Clusters): - transformation = dai.ImgTransformation() - clusters.transformation = transformation - assert clusters.transformation == transformation - - with pytest.raises(TypeError): - clusters.transformation = "not a dai.ImgTransformation" - - -def test_clusters_set_transformation_none(clusters: Clusters): - clusters.transformation = None - assert clusters.transformation is None diff --git a/tests/unittests/test_messages/test_keypoints_msg.py b/tests/unittests/test_messages/test_keypoints_msg.py deleted file mode 100644 index cff01f56..00000000 --- a/tests/unittests/test_messages/test_keypoints_msg.py +++ /dev/null @@ -1,48 +0,0 @@ -import depthai as dai -import pytest - -from depthai_nodes import Keypoints - - -@pytest.fixture -def keypoints(): - return Keypoints() - - -def test_keypoints_initialization(keypoints: Keypoints): - assert isinstance(keypoints.keypoints_list, dai.KeypointsList) - assert keypoints.getKeypoints() == [] - assert keypoints.getEdges() == [] - assert keypoints.transformation is None - - -def test_keypoints_set_keypoints_list(keypoints: Keypoints): - native = dai.KeypointsList() - point1 = dai.Keypoint() - point1.imageCoordinates = dai.Point3f(0.1, 0.2, 0.0) - point2 = dai.Keypoint() - point2.imageCoordinates = dai.Point3f(0.3, 0.4, 0.0) - native.setKeypoints([point1, point2]) - native.setEdges([(0, 1)]) - keypoints.keypoints_list = native - - assert keypoints.keypoints_list is native - assert keypoints.getKeypoints()[0].imageCoordinates.x == pytest.approx(0.1) - assert keypoints.getEdges() == [[0, 1]] - - with pytest.raises(TypeError): - keypoints.keypoints_list = "not a dai.KeypointsList" - - -def test_keypoints_set_transformation(keypoints: Keypoints): - transformation = dai.ImgTransformation() - keypoints.transformation = transformation - assert keypoints.transformation == transformation - - with pytest.raises(TypeError): - keypoints.transformation = "not a dai.ImgTransformation" - - -def test_keypoints_set_transformation_none(keypoints: Keypoints): - keypoints.transformation = None - assert keypoints.transformation is None diff --git a/tests/unittests/test_messages/test_lines_msg.py b/tests/unittests/test_messages/test_lines_msg.py deleted file mode 100644 index 3a9d0626..00000000 --- a/tests/unittests/test_messages/test_lines_msg.py +++ /dev/null @@ -1,93 +0,0 @@ -import depthai as dai -import numpy as np -import pytest - -from depthai_nodes import Line, Lines - - -@pytest.fixture -def line(): - return Line() - - -@pytest.fixture -def lines(): - return Lines() - - -def test_line_initialization(line: Line): - assert line.start_point is None - assert line.end_point is None - assert line.confidence is None - - -def test_line_set_start_point(line: Line): - start_point = dai.Point2f(0.1, 0.2) - line.start_point = start_point - assert np.allclose(line.start_point.x, 0.1, atol=1e-3) - assert np.allclose(line.start_point.y, 0.2, atol=1e-3) - - with pytest.raises(TypeError): - line.start_point = "not a Point2f" - - -def test_line_set_end_point(line: Line): - end_point = dai.Point2f(0.3, 0.4) - line.end_point = end_point - assert np.allclose(line.end_point.x, 0.3, atol=1e-3) - assert np.allclose(line.end_point.y, 0.4, atol=1e-3) - - with pytest.raises(TypeError): - line.end_point = "not a Point2f" - - -def test_line_set_confidence(line: Line): - line.confidence = 0.9 - assert line.confidence == 0.9 - - line.confidence = 1.05 - assert line.confidence == 1.0 - assert isinstance(line.confidence, float) - - line.confidence = -0.05 - assert line.confidence == 0.0 - assert isinstance(line.confidence, float) - - with pytest.raises(TypeError): - line.confidence = "not a float" - - with pytest.raises(ValueError): - line.confidence = 1.5 - - -def test_lines_initialization(lines: Lines): - assert lines.lines == [] - assert lines.transformation is None - - -def test_lines_set_lines(lines: Lines): - line1 = Line() - line2 = Line() - lines_list = [line1, line2] - lines.lines = lines_list - assert lines.lines == lines_list - - with pytest.raises(TypeError): - lines.lines = "not a list" - - with pytest.raises(ValueError): - lines.lines = [line1, "not a Line"] - - -def test_lines_set_transformation(lines: Lines): - transformation = dai.ImgTransformation() - lines.transformation = transformation - assert lines.transformation == transformation - - with pytest.raises(TypeError): - lines.transformation = "not a dai.ImgTransformation" - - -def test_lines_set_transformation_none(lines: Lines): - lines.transformation = None - assert lines.transformation is None diff --git a/tests/unittests/test_messages/test_map_msg.py b/tests/unittests/test_messages/test_map_msg.py deleted file mode 100644 index 250021e6..00000000 --- a/tests/unittests/test_messages/test_map_msg.py +++ /dev/null @@ -1,48 +0,0 @@ -import depthai as dai -import numpy as np -import pytest - -from depthai_nodes import Map2D - - -@pytest.fixture -def map2d(): - return Map2D() - - -def test_map2d_initialization(map2d: Map2D): - assert np.array_equal(map2d.map, np.array([])) - assert map2d.width is None - assert map2d.height is None - assert map2d.transformation is None - - -def test_map2d_set_map(map2d: Map2D): - map_array = np.random.rand(480, 640).astype(np.float32) - map2d.map = map_array - assert np.array_equal(map2d.map, map_array) - assert map2d.width == 640 - assert map2d.height == 480 - - with pytest.raises(TypeError): - map2d.map = "not a numpy array" - - with pytest.raises(ValueError): - map2d.map = np.random.rand(480, 640, 3).astype(np.float32) - - with pytest.raises(ValueError): - map2d.map = np.random.rand(480, 640).astype(np.float64) - - -def test_map2d_set_transformation(map2d: Map2D): - transformation = dai.ImgTransformation() - map2d.transformation = transformation - assert map2d.transformation == transformation - - with pytest.raises(TypeError): - map2d.transformation = "not a dai.ImgTransformation" - - -def test_map2d_set_transformation_none(map2d: Map2D): - map2d.transformation = None - assert map2d.transformation is None diff --git a/tests/unittests/test_messages/test_predictions_msg.py b/tests/unittests/test_messages/test_predictions_msg.py deleted file mode 100644 index 633f6841..00000000 --- a/tests/unittests/test_messages/test_predictions_msg.py +++ /dev/null @@ -1,68 +0,0 @@ -import depthai as dai -import pytest - -from depthai_nodes import Prediction, Predictions - - -@pytest.fixture -def prediction(): - return Prediction() - - -@pytest.fixture -def predictions(): - return Predictions() - - -def test_prediction_initialization(prediction: Prediction): - assert prediction.prediction is None - - -def test_prediction_set_prediction(prediction: Prediction): - prediction.prediction = 0.9 - assert prediction.prediction == 0.9 - - with pytest.raises(TypeError): - prediction.prediction = "not a float" - - -def test_predictions_initialization(predictions: Predictions): - assert predictions.predictions == [] - assert predictions.transformation is None - - -def test_predictions_set_predictions(predictions: Predictions): - pred1 = Prediction() - pred1.prediction = 0.1 - pred2 = Prediction() - pred2.prediction = 0.2 - predictions_list = [pred1, pred2] - predictions.predictions = predictions_list - assert predictions.predictions == predictions_list - - with pytest.raises(TypeError): - predictions.predictions = "not a list" - - with pytest.raises(ValueError): - predictions.predictions = [pred1, "not a Prediction"] - - -def test_predictions_get_prediction(predictions: Predictions): - pred1 = Prediction() - pred1.prediction = 0.1 - predictions.predictions = [pred1] - assert predictions.prediction == 0.1 - - -def test_predictions_set_transformation(predictions: Predictions): - transformation = dai.ImgTransformation() - predictions.transformation = transformation - assert predictions.transformation == transformation - - with pytest.raises(TypeError): - predictions.transformation = "not a dai.ImgTransformation" - - -def test_predictions_set_transformation_none(predictions: Predictions): - predictions.transformation = None - assert predictions.transformation is None