Add dai.ImgTransformations to messages and improvements to Text detection parser

depthai_nodes/ml/messages/classification.py

@@ -14,13 +14,16 @@ class Classifications(dai.Buffer):
         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
 
     @property
     def classes(self) -> List:
@@ -91,3 +94,26 @@ def top_score(self) -> float:
         @rtype: float
         """
         return self._scores[0]
+
+    @property
+    def transformation(self) -> dai.ImgTransformation:
+        """Returns the Image Transformation object.
+
+        @return: The Image Transformation object.
+        @rtype: dai.ImgTransformation
+        """
+        return self._transformation
+
+    @transformation.setter
+    def transformation(self, value: dai.ImgTransformation):
+        """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 not isinstance(value, dai.ImgTransformation):
+            raise TypeError(
+                f"Transformation must be a dai.ImgTransformation object, instead got {type(value)}."
+            )
+        self._transformation = value

depthai_nodes/ml/messages/clusters.py

@@ -73,12 +73,15 @@ class Clusters(dai.Buffer):
     ----------
     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
 
     @property
     def clusters(self) -> List[Cluster]:
@@ -103,3 +106,26 @@ def clusters(self, value: List[Cluster]):
         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:
+        """Returns the Image Transformation object.
+
+        @return: The Image Transformation object.
+        @rtype: dai.ImgTransformation
+        """
+        return self._transformation
+
+    @transformation.setter
+    def transformation(self, value: dai.ImgTransformation):
+        """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 not isinstance(value, dai.ImgTransformation):
+            raise TypeError(
+                f"Transformation must be a dai.ImgTransformation object, instead got {type(value)}."
+            )
+        self._transformation = value

depthai_nodes/ml/messages/img_detections.py

@@ -161,13 +161,16 @@ class ImgDetectionsExtended(dai.Buffer):
         Image detections with keypoints.
     masks: np.ndarray
         The segmentation masks of the image. All masks are stored in a single numpy array.
+    transformation : dai.ImgTransformation
+        Image transformation object.
     """
 
     def __init__(self) -> None:
         """Initializes the ImgDetectionsExtended object."""
         super().__init__()
         self._detections: List[ImgDetectionExtended] = []
         self._masks: SegmentationMask = SegmentationMask()
+        self._transformation: dai.ImgTransformation = None
 
     @property
     def detections(self) -> List[ImgDetectionExtended]:
@@ -226,3 +229,26 @@ def masks(self, value: NDArray[np.int16]):
         masks_msg = SegmentationMask()
         masks_msg.mask = value
         self._masks = masks_msg
+
+    @property
+    def transformation(self) -> dai.ImgTransformation:
+        """Returns the Image Transformation object.
+
+        @return: The Image Transformation object.
+        @rtype: dai.ImgTransformation
+        """
+        return self._transformation
+
+    @transformation.setter
+    def transformation(self, value: dai.ImgTransformation):
+        """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 not isinstance(value, dai.ImgTransformation):
+            raise TypeError(
+                f"Transformation must be a dai.ImgTransformation object, instead got {type(value)}."
+            )
+        self._transformation = value

depthai_nodes/ml/messages/keypoints.py

@@ -127,12 +127,15 @@ class Keypoints(dai.Buffer):
     ----------
     keypoints: List[Keypoint]
         List of Keypoint objects, each representing a keypoint.
+    transformation : dai.ImgTransformation
+        Image transformation object.
     """
 
     def __init__(self):
         """Initializes the Keypoints object."""
         super().__init__()
         self._keypoints: List[Keypoint] = []
+        self._transformation: dai.ImgTransformation = None
 
     @property
     def keypoints(self) -> List[Keypoint]:
@@ -157,3 +160,26 @@ def keypoints(self, value: List[Keypoint]):
         if not all(isinstance(item, Keypoint) for item in value):
             raise ValueError("keypoints must be a list of Keypoint objects.")
         self._keypoints = value
+
+    @property
+    def transformation(self) -> dai.ImgTransformation:
+        """Returns the Image Transformation object.
+
+        @return: The Image Transformation object.
+        @rtype: dai.ImgTransformation
+        """
+        return self._transformation
+
+    @transformation.setter
+    def transformation(self, value: dai.ImgTransformation):
+        """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 not isinstance(value, dai.ImgTransformation):
+            raise TypeError(
+                f"Transformation must be a dai.ImgTransformation object, instead got {type(value)}."
+            )
+        self._transformation = value

depthai_nodes/ml/messages/lines.py

@@ -100,12 +100,15 @@ class Lines(dai.Buffer):
     ----------
     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
 
     @property
     def lines(self) -> List[Line]:
@@ -130,3 +133,26 @@ def lines(self, value: List[Line]):
         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:
+        """Returns the Image Transformation object.
+
+        @return: The Image Transformation object.
+        @rtype: dai.ImgTransformation
+        """
+        return self._transformation
+
+    @transformation.setter
+    def transformation(self, value: dai.ImgTransformation):
+        """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 not isinstance(value, dai.ImgTransformation):
+            raise TypeError(
+                f"Transformation must be a dai.ImgTransformation object, instead got {type(value)}."
+            )
+        self._transformation = value

depthai_nodes/ml/messages/map.py

@@ -14,6 +14,8 @@ class Map2D(dai.Buffer):
         2D Map width.
     height : int
         2D Map height.
+    transformation : dai.ImgTransformation
+        Image transformation object.
     """
 
     def __init__(self):
@@ -22,6 +24,7 @@ def __init__(self):
         self._map: NDArray[np.float32] = np.array([])
         self._width: int = None
         self._height: int = None
+        self._transformation: dai.ImgTransformation = None
 
     @property
     def map(self) -> NDArray[np.float32]:
@@ -71,3 +74,26 @@ def height(self) -> int:
         @rtype: int
         """
         return self._height
+
+    @property
+    def transformation(self) -> dai.ImgTransformation:
+        """Returns the Image Transformation object.
+
+        @return: The Image Transformation object.
+        @rtype: dai.ImgTransformation
+        """
+        return self._transformation
+
+    @transformation.setter
+    def transformation(self, value: dai.ImgTransformation):
+        """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 not isinstance(value, dai.ImgTransformation):
+            raise TypeError(
+                f"Transformation must be a dai.ImgTransformation object, instead got {type(value)}."
+            )
+        self._transformation = value

depthai_nodes/ml/messages/prediction.py

@@ -48,12 +48,15 @@ class Predictions(dai.Buffer):
     ----------
     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
 
     @property
     def predictions(self) -> List[Prediction]:
@@ -89,3 +92,26 @@ def prediction(self) -> float:
         @rtype: float
         """
         return self._predictions[0].prediction
+
+    @property
+    def transformation(self) -> dai.ImgTransformation:
+        """Returns the Image Transformation object.
+
+        @return: The Image Transformation object.
+        @rtype: dai.ImgTransformation
+        """
+        return self._transformation
+
+    @transformation.setter
+    def transformation(self, value: dai.ImgTransformation):
+        """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 not isinstance(value, dai.ImgTransformation):
+            raise TypeError(
+                f"Transformation must be a dai.ImgTransformation object, instead got {type(value)}."
+            )
+        self._transformation = value

depthai_nodes/ml/messages/segmentation.py

@@ -12,12 +12,15 @@ class SegmentationMask(dai.Buffer):
     ----------
     mask: NDArray[np.int16]
         Segmentation mask.
+    transformation : dai.ImgTransformation
+        Image transformation object.
     """
 
     def __init__(self):
         """Initializes the SegmentationMask object."""
         super().__init__()
         self._mask: NDArray[np.int16] = np.array([])
+        self._transformation: dai.ImgTransformation = None
 
     @property
     def mask(self) -> NDArray[np.int16]:
@@ -48,3 +51,26 @@ def mask(self, value: NDArray[np.int16]):
         if np.any((value < -1)):
             raise ValueError("Mask must be an array of integers larger or equal to -1.")
         self._mask = value
+
+    @property
+    def transformation(self) -> dai.ImgTransformation:
+        """Returns the Image Transformation object.
+
+        @return: The Image Transformation object.
+        @rtype: dai.ImgTransformation
+        """
+        return self._transformation
+
+    @transformation.setter
+    def transformation(self, value: dai.ImgTransformation):
+        """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 not isinstance(value, dai.ImgTransformation):
+            raise TypeError(
+                f"Transformation must be a dai.ImgTransformation object, instead got {type(value)}."
+            )
+        self._transformation = value

depthai_nodes/ml/parsers/classification.py

@@ -141,6 +141,7 @@ def run(self):
                 scores = softmax(scores)
 
             msg = create_classification_message(self.classes, scores)
+            msg.transformation = output.getTransformation()
             msg.setTimestamp(output.getTimestamp())
 
             self.out.send(msg)

depthai_nodes/ml/parsers/classification_sequence.py

@@ -178,6 +178,7 @@ def run(self):
                 ignored_indexes=self.ignored_indexes,
                 concatenate_classes=self.concatenate_classes,
             )
+            msg.transformation = output.getTransformation()
             msg.setTimestamp(output.getTimestamp())
 
             self.out.send(msg)

depthai_nodes/ml/parsers/fastsam.py

@@ -353,6 +353,6 @@ def run(self):
             results_masks = merge_masks(results_masks)
 
             segmentation_message = create_segmentation_message(results_masks)
+            segmentation_message.transformation = output.getTransformation()
             segmentation_message.setTimestamp(output.getTimestamp())
-
             self.out.send(segmentation_message)

depthai_nodes/ml/parsers/hrnet.py

@@ -113,5 +113,6 @@ def run(self):
                 confidence_threshold=self.score_threshold,
             )
             keypoints_message.setTimestamp(output.getTimestamp())
+            keypoints_message.transformation = output.getTransformation()
 
             self.out.send(keypoints_message)

depthai_nodes/ml/parsers/image_output.py

@@ -113,5 +113,6 @@ def run(self):
                 is_bgr=self.output_is_bgr,
             )
             image_message.setTimestamp(output.getTimestamp())
+            image_message.transformation = output.getTransformation()
 
             self.out.send(image_message)

depthai_nodes/ml/parsers/keypoints.py

@@ -172,5 +172,6 @@ def run(self):
 
             msg = create_keypoints_message(keypoints)
             msg.setTimestamp(output.getTimestamp())
+            msg.transformation = output.getTransformation()
 
             self.out.send(msg)