depthai-nodes 0.0.1__py3-none-any.whl

depthai_nodes/__init__.py

@@ -0,0 +1,3 @@
+from .ml.parsers import *
+
+__version__ = "0.0.1"

depthai_nodes/ml/messages/__init__.py

@@ -0,0 +1,21 @@
+from .classification import Classifications
+from .img_detections import (
+    ImgDetectionExtended,
+    ImgDetectionsExtended,
+)
+from .keypoints import HandKeypoints, Keypoints
+from .lines import Line, Lines
+from .misc import AgeGender
+from .segmentation import SegmentationMasks
+
+__all__ = [
+    "ImgDetectionExtended",
+    "ImgDetectionsExtended",
+    "HandKeypoints",
+    "Keypoints",
+    "Line",
+    "Lines",
+    "Classifications",
+    "SegmentationMasks",
+    "AgeGender",
+]

depthai_nodes/ml/messages/classification.py

@@ -0,0 +1,48 @@
+from typing import List
+
+import depthai as dai
+
+
+class Classifications(dai.Buffer):
+    """Classification class for storing the class names and their respective scores.
+
+    Attributes
+    ----------
+    classes : list[str]
+        A list of classes.
+    scores : list[float]
+        A list of corresponding probability scores.
+    """
+
+    def __init__(self):
+        """Initializes the Classifications object and sets the classes and scores to
+        empty lists."""
+        dai.Buffer.__init__(self)
+        self._classes: List[str] = []
+        self._scores: List[float] = []
+
+    @property
+    def classes(self) -> List:
+        """Returns the list of classes."""
+        return self._classes
+
+    @property
+    def scores(self) -> List:
+        """Returns the list of scores."""
+        return self._scores
+
+    @classes.setter
+    def classes(self, class_names: List[str]):
+        """Sets the list of classes.
+
+        @param classes: A list of class names.
+        """
+        self._classes = class_names
+
+    @scores.setter
+    def scores(self, scores: List[float]):
+        """Sets the list of scores.
+
+        @param scores: A list of scores.
+        """
+        self._scores = scores

depthai_nodes/ml/messages/creators/__init__.py

@@ -0,0 +1,24 @@
+from .classification import create_classification_message
+from .depth import create_depth_message
+from .detection import create_detection_message, create_line_detection_message
+from .image import create_image_message
+from .keypoints import create_hand_keypoints_message, create_keypoints_message
+from .misc import create_age_gender_message
+from .segmentation import create_sam_message, create_segmentation_message
+from .thermal import create_thermal_message
+from .tracked_features import create_tracked_features_message
+
+__all__ = [
+    "create_image_message",
+    "create_segmentation_message",
+    "create_hand_keypoints_message",
+    "create_detection_message",
+    "create_depth_message",
+    "create_line_detection_message",
+    "create_tracked_features_message",
+    "create_keypoints_message",
+    "create_thermal_message",
+    "create_classification_message",
+    "create_sam_message",
+    "create_age_gender_message",
+]

depthai_nodes/ml/messages/creators/classification.py

@@ -0,0 +1,79 @@
+from typing import List, Union
+
+import numpy as np
+
+from ...messages import Classifications
+
+
+def create_classification_message(
+    classes: List, scores: Union[np.ndarray, List]
+) -> Classifications:
+    """Create a message for classification. The message contains the class names and
+    their respective scores, sorted in descending order of scores.
+
+    @param classes: A list containing class names.
+    @type classes: List
+    @param scores: A numpy array of shape (n_classes,) containing the probability score of each class.
+    @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
+
+    @raises ValueError: If the provided classes are None.
+    @raises ValueError: If the provided classes are not a list.
+    @raises ValueError: If the provided classes are empty.
+    @raises ValueError: If the provided scores are None.
+    @raises ValueError: If the provided scores are not a list or a numpy array.
+    @raises ValueError: If the provided scores are empty.
+    @raises ValueError: If the provided scores are not a 1D array.
+    @raises ValueError: If the provided scores are not of type float.
+    @raises ValueError: If the provided scores do not sum to 1.
+    @raises ValueError: If the number of labels and scores mismatch.
+    """
+    if isinstance(classes, type(None)):
+        raise ValueError("Classes should not be None.")
+
+    if not isinstance(classes, list):
+        raise ValueError(f"Classes should be a list, got {type(classes)}.")
+
+    if len(classes) == 0:
+        raise ValueError("Classes should not be empty.")
+
+    if type(scores) == type(None):
+        raise ValueError("Scores should not be None.")
+
+    if not isinstance(scores, np.ndarray) and not isinstance(scores, list):
+        raise ValueError(
+            f"Scores should be a list or a numpy array, got {type(scores)}."
+        )
+
+    if isinstance(scores, list):
+        scores = np.array(scores)
+
+    if len(scores) == 0:
+        raise ValueError("Scores should not be empty.")
+
+    if len(scores) != len(scores.flatten()):
+        raise ValueError(f"Scores should be a 1D array, got {scores.shape}.")
+
+    scores = scores.flatten()
+
+    if not np.issubdtype(scores.dtype, np.floating):
+        raise ValueError(f"Scores should be of type float, got {scores.dtype}.")
+
+    if not np.isclose(np.sum(scores), 1.0, atol=1e-2):
+        raise ValueError(f"Scores should sum to 1, got {np.sum(scores)}.")
+
+    if len(scores) != len(classes):
+        raise ValueError(
+            f"Number of labels and scores mismatch. Provided {len(scores)} scores and {len(classes)} class names."
+        )
+
+    classification_msg = Classifications()
+    sorted_args = np.argsort(-scores, kind="stable")
+    scores = scores[sorted_args]
+
+    classification_msg.classes = [classes[i] for i in sorted_args]
+    classification_msg.scores = scores.tolist()
+
+    return classification_msg

depthai_nodes/ml/messages/creators/depth.py

@@ -0,0 +1,93 @@
+from typing import Literal
+
+import depthai as dai
+import numpy as np
+
+UINT16_MAX_VALUE = 65535
+
+
+def create_depth_message(
+    depth_map: np.ndarray,
+    depth_type: Literal["relative", "metric"],
+    depth_limit: float = 0.0,
+) -> dai.ImgFrame:
+    """Create a DepthAI message for a depth map.
+
+    @param depth_map: A NumPy array representing the depth map with shape HW or NHW/HWN.
+        Here N stands for batch dimension.
+    @type depth_map: np.array
+    @param depth_type: A string indicating the type of depth map. It can either be
+        'relative' or 'metric'.
+    @type depth_type: Literal['relative', 'metric']
+    @param depth_limit: The maximum depth value (in meters) to be used in the depth map.
+        The default value is 0, which means no limit.
+    @type depth_limit: float
+    @return: An ImgFrame object containing the depth information.
+    @rtype: dai.ImgFrame
+    @raise ValueError: If the depth map is not a NumPy array.
+    @raise ValueError: If the depth map is not 2D or 3D.
+    @raise ValueError: If the depth map shape is not NHW or HWN.
+    @raise ValueError: If the depth type is not 'relative' or 'metric'.
+    @raise ValueError: If the depth limit is not 0 and the depth type is 'relative'.
+    @raise ValueError: If the depth limit is 0 and the depth type is 'metric'.
+    @raise ValueError: If the depth limit is negative.
+    """
+
+    if not isinstance(depth_map, np.ndarray):
+        raise ValueError(f"Expected numpy array, got {type(depth_map)}.")
+
+    if len(depth_map.shape) == 3:
+        if depth_map.shape[0] == 1:
+            depth_map = depth_map[0, :, :]  # NHW to HW
+        elif depth_map.shape[2] == 1:
+            depth_map = depth_map[:, :, 0]  # HWN to HW
+        else:
+            raise ValueError(
+                f"Unexpected image shape. Expected NHW or HWN, got {depth_map.shape}."
+            )
+
+    if len(depth_map.shape) != 2:
+        raise ValueError(f"Expected 2D or 3D input, got {len(depth_map.shape)}D input.")
+
+    if not (depth_type == "relative" or depth_type == "metric"):
+        raise ValueError(
+            f"Invalid depth type: {depth_type}. Only 'relative' and 'metric' are supported."
+        )
+
+    if depth_type == "relative" and depth_limit != 0:
+        raise ValueError(
+            f"Invalid depth limit: {depth_limit}. For relative depth, depth limit must be equal to 0."
+        )
+
+    if depth_type == "metric" and depth_limit == 0:
+        raise ValueError(
+            f"Invalid depth limit: {depth_limit}. For metric depth, depth limit must be bigger than 0."
+        )
+
+    if depth_limit < 0:
+        raise ValueError(
+            f"Invalid depth limit: {depth_limit}. Depth limit must be bigger than 0."
+        )
+
+    data_type = dai.ImgFrame.Type.RAW16
+
+    min_val = depth_map.min() if depth_type == "relative" else 0
+    max_val = depth_map.max() if depth_type == "relative" else depth_limit
+
+    # clip values bigger than max_val
+    depth_map = np.clip(depth_map, a_min=None, a_max=max_val)
+
+    # normalize depth map to UINT16 range [0, UINT16_MAX_VALUE]
+    if min_val == max_val:  # avoid division by zero
+        depth_map = np.zeros_like(depth_map)
+    else:
+        depth_map = (depth_map - min_val) / (max_val - min_val) * UINT16_MAX_VALUE
+    depth_map = depth_map.astype(np.uint16)
+
+    imgFrame = dai.ImgFrame()
+    imgFrame.setFrame(depth_map)
+    imgFrame.setWidth(depth_map.shape[1])
+    imgFrame.setHeight(depth_map.shape[0])
+    imgFrame.setType(data_type)
+
+    return imgFrame

depthai_nodes/ml/messages/creators/detection.py

@@ -0,0 +1,235 @@
+from typing import List, Tuple, Union
+
+import depthai as dai
+import numpy as np
+
+from ...messages import (
+    ImgDetectionExtended,
+    ImgDetectionsExtended,
+    Line,
+    Lines,
+)
+
+
+def create_detection_message(
+    bboxes: np.ndarray,
+    scores: np.ndarray,
+    labels: List[int] = None,
+    keypoints: Union[
+        List[Tuple[float, float]], List[Tuple[float, float, float]]
+    ] = None,
+    masks: List[np.ndarray] = None,
+) -> dai.ImgDetections:
+    """Create a DepthAI message for an object detection.
+
+    @param bbox: Bounding boxes of detected objects of shape (N,4) meaning [...,[x_min, y_min, x_max, y_max],...].
+    @type bbox: np.ndarray
+    @param scores: Confidence scores of detected objects of shape (N,).
+    @type scores: np.ndarray
+    @param labels: Labels of detected objects of shape (N,).
+    @type labels: List[int]
+    @param keypoints: Keypoints of detected objects of shape (N,2) or (N,3).
+    @type keypoints: Optional[Union[List[Tuple[float, float]], List[Tuple[float, float, float]]]]
+    @param masks: Masks of detected objects of shape (N, H, W).
+    @type masks: List[np.ndarray]
+
+    @return: Message containing the bounding boxes, labels, confidence scores, and keypoints of detected objects.
+    @rtype: dai.ImgDetections OR ImgDetectionsExtended
+
+    @raise ValueError: If the bboxes are not a numpy array.
+    @raise ValueError: If the bboxes are not of shape (N,4).
+    @raise ValueError: If the bboxes 2nd dimension is not of size 4.
+    @raise ValueError: If the bboxes are not in format [x_min, y_min, x_max, y_max] where xmin < xmax and ymin < ymax.
+    @raise ValueError: If the scores are not a numpy array.
+    @raise ValueError: If the scores are not of shape (N,).
+    @raise ValueError: If the scores do not have the same length as bboxes.
+    @raise ValueError: If the labels are not a list.
+    @raise ValueError: If each label is not an integer.
+    @raise ValueError: If the labels do not have the same length as bboxes.
+    @raise ValueError: If the keypoints are not a list.
+    @raise ValueError: If each keypoint pair is not a tuple of two or three floats.
+    @raise ValueError: If the keypoints do not have the same length as bboxes.
+    @raise ValueError: If the masks are not a list.
+    @raise ValueError: If each mask is not a 2D numpy array.
+    """
+
+    # checks for bboxes
+    if not isinstance(bboxes, np.ndarray):
+        raise ValueError(f"Bounding boxes should be numpy array, got {type(bboxes)}.")
+
+    if len(bboxes) != 0:
+        if len(bboxes.shape) != 2:
+            raise ValueError(
+                f"Bounding boxes should be of shape (N,4) meaning [...,[x_min, y_min, x_max, y_max],...], got {bboxes.shape}."
+            )
+        if bboxes.shape[1] != 4:
+            raise ValueError(
+                f"Bounding boxes 2nd dimension should be of size 4 e.g. [x_min, y_min, x_max, y_max] got {bboxes.shape[1]}."
+            )
+
+        x_valid = bboxes[:, 0] < bboxes[:, 2]
+        y_valid = bboxes[:, 1] < bboxes[:, 3]
+        if not (np.all(x_valid) and np.all(y_valid)):
+            raise ValueError(
+                "Bounding boxes should be in format [x_min, y_min, x_max, y_max] where xmin < xmax and ymin < ymax."
+            )
+
+    # checks for scores
+    if not isinstance(scores, np.ndarray):
+        raise ValueError(f"Scores should be numpy array, got {type(scores)}.")
+
+    if len(scores) != 0:
+        if len(scores.shape) != 1:
+            raise ValueError(
+                f"Scores should be of shape (N,) meaning, got {scores.shape}."
+            )
+        if scores.shape[0] != bboxes.shape[0]:
+            raise ValueError(
+                f"Scores should have same length as bboxes, got {scores.shape[0]} and {bboxes.shape[0]}."
+            )
+
+    # checks for labels
+    if labels is not None:
+        if not isinstance(labels, List):
+            raise ValueError(f"Labels should be list, got {type(labels)}.")
+        for label in labels:
+            if not isinstance(label, int):
+                raise ValueError(
+                    f"Labels should be list of integers, got {type(label)}."
+                )
+        if len(labels) != bboxes.shape[0]:
+            raise ValueError(
+                f"Labels should have same length as bboxes, got {len(labels)} and {bboxes.shape[0]}."
+            )
+
+    # checks for keypoints
+    if keypoints is not None:
+        if not isinstance(keypoints, List):
+            raise ValueError(f"Keypoints should be list, got {type(keypoints)}.")
+        if len(keypoints) != bboxes.shape[0]:
+            raise ValueError(
+                f"Keypoints should have same length as bboxes, got {len(keypoints)} and {bboxes.shape[0]}."
+            )
+
+        for object_keypoints in keypoints:
+            dim = len(object_keypoints[0]) if len(object_keypoints) != 0 else 0
+            for point in object_keypoints:
+                if not isinstance(point, (Tuple, List)):
+                    raise ValueError(
+                        f"Keypoint pairs should be list of tuples, got {type(point)}."
+                    )
+                if len(point) not in [2, 3]:
+                    raise ValueError(
+                        f"Keypoint pairs should be list of tuples of length 2 or 3, got {len(point)}."
+                    )
+                if len(point) != dim:
+                    raise ValueError(
+                        "All keypoints should be of same dimension e.g. [x, y] or [x, y, z], got mixed inner dimensions."
+                    )
+
+    if masks is not None and len(masks) != 0:
+        if not isinstance(masks, List):
+            raise ValueError(f"masks should be list, got {type(masks)}.")
+        for mask in masks:
+            if not isinstance(mask, np.ndarray):
+                raise ValueError(f"mask should be numpy array, got {type(mask)}.")
+            if len(mask.shape) != 2:
+                raise ValueError(f"mask should be a 2D matrix, got {mask.shape}.")
+
+        if len(masks) != bboxes.shape[0]:
+            raise ValueError(
+                f"masks should have same length as bboxes, got {len(masks)} and {bboxes.shape[0]}."
+            )
+
+    if keypoints is not None or masks is not None:
+        img_detection = ImgDetectionExtended
+        img_detections = ImgDetectionsExtended
+    else:
+        img_detection = dai.ImgDetection
+        img_detections = dai.ImgDetections
+
+    detections = []
+    for i in range(bboxes.shape[0]):
+        detection = img_detection()
+        detection.xmin = bboxes[i][0]
+        detection.ymin = bboxes[i][1]
+        detection.xmax = bboxes[i][2]
+        detection.ymax = bboxes[i][3]
+        detection.confidence = scores[i]
+        if labels is None:
+            detection.label = 0
+        else:
+            detection.label = labels[i]
+        if keypoints is not None:
+            detection.keypoints = keypoints[i]
+        if masks is not None:
+            detection.mask = masks[i]
+        detections.append(detection)
+
+    detections_msg = img_detections()
+    detections_msg.detections = detections
+    return detections_msg
+
+
+def create_line_detection_message(lines: np.ndarray, scores: np.ndarray):
+    """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],...].
+    @type lines: np.ndarray
+    @param scores: Confidence scores of detected lines of shape (N,).
+    @type scores: np.ndarray
+
+    @return: Message containing the lines and confidence scores of detected lines.
+    @rtype: Lines
+
+    @raise ValueError: If the lines are not a numpy array.
+    @raise ValueError: If the lines are not of shape (N,4).
+    @raise ValueError: If the lines 2nd dimension is not of size E{4}.
+    @raise ValueError: If the scores are not a numpy array.
+    @raise ValueError: If the scores are not of shape (N,).
+    @raise ValueError: If the scores do not have the same length as lines.
+    """
+
+    # checks for lines
+    if not isinstance(lines, np.ndarray):
+        raise ValueError(f"Lines should be numpy array, got {type(lines)}.")
+    if len(lines) != 0:
+        if len(lines.shape) != 2:
+            raise ValueError(
+                f"Lines should be of shape (N,4) meaning [...,[x_start, y_start, x_end, y_end],...], got {lines.shape}."
+            )
+        if lines.shape[1] != 4:
+            raise ValueError(
+                f"Lines 2nd dimension should be of size 4 e.g. [x_start, y_start, x_end, y_end] got {lines.shape[1]}."
+            )
+
+    # checks for scores
+    if not isinstance(scores, np.ndarray):
+        raise ValueError(f"Scores should be numpy array, got {type(scores)}.")
+
+    if len(scores) != 0:
+        if len(scores.shape) != 1:
+            raise ValueError(
+                f"Scores should be of shape (N,) meaning, got {scores.shape}."
+            )
+
+        for score in scores:
+            if not isinstance(score, (float, np.floating)):
+                raise ValueError(f"Scores should be of type float, got {type(score)}.")
+
+    if scores.shape[0] != lines.shape[0]:
+        raise ValueError(
+            f"Scores should have same length as lines, got {scores.shape[0]} and {lines.shape[0]}."
+        )
+
+    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.confidence = float(scores[i])
+        line_detections.append(line_detection)
+
+    lines_msg = Lines()
+    lines_msg.lines = line_detections
+    return lines_msg

depthai_nodes/ml/messages/creators/image.py

@@ -0,0 +1,53 @@
+import cv2
+import depthai as dai
+import numpy as np
+
+
+def create_image_message(
+    image: np.ndarray,
+    is_bgr: bool = True,
+) -> dai.ImgFrame:
+    """Create a DepthAI message for an image array.
+
+    @param image: Image array in HWC or CHW format.
+    @type image: np.array
+    @param is_bgr: If True, the image is in BGR format. If False, the image is in RGB
+        format. Defaults to True.
+    @type is_bgr: bool
+    @return: dai.ImgFrame object containing the image information.
+    @rtype: dai.ImgFrame
+    @raise ValueError: If the image shape is not CHW or HWC.
+    """
+
+    if image.shape[0] in [1, 3]:
+        hwc = False
+    elif image.shape[2] in [1, 3]:
+        hwc = True
+    else:
+        raise ValueError(
+            f"Unexpected image shape. Expected CHW or HWC, got {image.shape}"
+        )
+
+    if not hwc:
+        image = np.transpose(image, (1, 2, 0))
+
+    if isinstance(image[0, 0, 0], (float, np.floating)):
+        raise ValueError(f"Expected int type, got {type(image[0, 0, 0])}.")
+
+    if image.shape[2] == 1:  # grayscale
+        image = image[:, :, 0]
+        img_frame_type = dai.ImgFrame.Type.GRAY8  # HW image
+        height, width = image.shape
+    else:
+        if not is_bgr:
+            image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+        img_frame_type = dai.ImgFrame.Type.BGR888i  # HWC BGR image
+        height, width, _ = image.shape
+
+    imgFrame = dai.ImgFrame()
+    imgFrame.setFrame(image)
+    imgFrame.setWidth(width)
+    imgFrame.setHeight(height)
+    imgFrame.setType(img_frame_type)
+
+    return imgFrame