Classification parser

.gitignore

@@ -3,7 +3,6 @@
 __pycache__/
 *.py[cod]
 *$py.class
-
 # C extensions
 *.so
 
@@ -162,4 +161,4 @@ cython_debug/
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
 
-.DS_Store
+.DS_Store

depthai_nodes/ml/messages/__init__.py

@@ -1,3 +1,4 @@
+from .classification import Classifications
 from .img_detections import ImgDetectionsWithKeypoints, ImgDetectionWithKeypoints
 from .keypoints import HandKeypoints, Keypoints
 from .lines import Line, Lines
@@ -9,4 +10,5 @@
     "Keypoints",
     "Line",
     "Lines",
+    "Classifications",
 ]

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

@@ -1,3 +1,4 @@
+from .classification_message 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
@@ -16,4 +17,5 @@
     "create_tracked_features_message",
     "create_keypoints_message",
     "create_thermal_message",
+    "create_classification_message",
 ]

depthai_nodes/ml/messages/creators/classification_message.py

@@ -0,0 +1,68 @@
+import depthai as dai
+import numpy as np
+
+from ...messages import Classifications
+
+
+def create_classification_message(
+    scores: np.ndarray, classes: np.ndarray = None
+) -> dai.Buffer:
+    """Create a message for classification. The message contains the class names and
+    their respective scores, sorted in descending order of scores.
+
+    Parameters
+    ----------
+    scores : np.ndarray
+        A numpy array of shape (n_classes,) containing the probability score of each class.
+
+    classes : np.ndarray = []
+        A numpy array of shape (n_classes, ), containing class names. If not provided, class names are set to [].
+
+
+    Returns
+    --------
+    Classifications : dai.Buffer
+        A message with parameter `classes` which is a list of shape (n_classes, 2)
+        where each item is [class_name, probability_score].
+        If no class names are provided, class_name is set to None.
+    """
+
+    if type(classes) == type(None):
+        classes = np.array([])
+    else:
+        classes = np.array(classes)
+
+    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}.")
+
+    if len(classes) != len(classes.flatten()):
+        raise ValueError(f"Classes should be a 1D array, got {classes.shape}.")
+
+    scores = scores.flatten()
+    classes = classes.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-1):
+        raise ValueError(f"Scores should sum to 1, got {np.sum(scores)}.")
+
+    if len(scores) != len(classes) and len(classes) != 0:
+        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)[::-1]
+    scores = scores[sorted_args]
+
+    if len(classes) != 0:
+        classification_msg.classes = classes[sorted_args].tolist()
+
+    classification_msg.scores = scores.tolist()
+
+    return classification_msg

depthai_nodes/ml/parsers/__init__.py

@@ -1,3 +1,4 @@
+from .classification_parser import ClassificationParser
 from .image_output import ImageOutputParser
 from .keypoints import KeypointParser
 from .mediapipe_hand_landmarker import MPHandLandmarkParser
@@ -24,4 +25,5 @@
     "MLSDParser",
     "XFeatParser",
     "ThermalImageParser",
+    "ClassificationParser",
 ]

depthai_nodes/ml/parsers/classification_parser.py

@@ -0,0 +1,90 @@
+import depthai as dai
+import numpy as np
+
+from ..messages.creators import create_classification_message
+
+
+class ClassificationParser(dai.node.ThreadedHostNode):
+    """Postprocessing logic for Classification model.
+
+    Attributes
+    ----------
+    input : Node.Input
+        Node's input. It is a linking point to which the Neural Network's output is linked. It accepts the output of the Neural Network node.
+    out : Node.Output
+        Parser sends the processed network results to this output in a form of DepthAI message. It is a linking point from which the processed network results are retrieved.
+    classes : list[str]
+        List of class names to be used for linking with their respective scores. Expected to be in the same order as Neural Network's output. If not provided, the message will only return sorted scores.
+    is_softmax : bool = True
+        If False, the scores are converted to probabilities using softmax function.
+    n_classes : int = len(classes)
+        Number of provided classes. This variable is set automatically based on provided classes.
+
+    Output Message/s
+    ----------------
+    **Type** : Classifications(dai.Buffer):
+         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.
+    """
+
+    def __init__(self, classes: list[str] = None, is_softmax: bool = True):
+        """Initializes the ClassificationParser node.
+
+        @param classes: List of class names to be used for linking with their respective
+            scores.
+        @param is_softmax: If False, the scores are converted to probabilities using
+            softmax function.
+        """
+
+        dai.node.ThreadedHostNode.__init__(self)
+        self.out = self.createOutput()
+        self.input = self.createInput()
+        self.classes = classes if classes is not None else []
+        self.n_classes = len(self.classes)
+        self.is_softmax = is_softmax
+
+    def setClasses(self, classes: list[str]):
+        """Sets the class names for the classification model.
+
+        @param classes: List of class names to be used for linking with their respective
+            scores.
+        """
+        self.classes = classes if classes is not None else []
+        self.n_classes = len(self.classes)
+
+    def setSoftmax(self, is_softmax: bool):
+        """Sets the softmax flag for the classification model.
+
+        @param is_softmax: If False, the parser will convert the scores to probabilities
+            using softmax function.
+        """
+        self.is_softmax = is_softmax
+
+    def run(self):
+        while self.isRunning():
+            try:
+                output: dai.NNData = self.input.get()
+            except dai.MessageQueue.QueueException:
+                break  # Pipeline was stopped
+
+            output_layer_names = output.getAllLayerNames()
+            if len(output_layer_names) != 1:
+                raise ValueError(
+                    f"Expected 1 output layer, got {len(output_layer_names)}."
+                )
+
+            scores = output.getTensor(output_layer_names[0])
+            scores = np.array(scores).flatten()
+            classes = np.array(self.classes)
+            if len(scores) != self.n_classes and self.n_classes != 0:
+                raise ValueError(
+                    f"Number of labels and scores mismatch. Provided {self.n_classes} class names and {len(scores)} scores."
+                )
+
+            if not self.is_softmax:
+                ex = np.exp(scores)
+                scores = ex / np.sum(ex)
+
+            msg = create_classification_message(scores, classes)
+            msg.setTimestamp(output.getTimestamp())
+
+            self.out.send(msg)