Merge pull request #33 from luxonis/creator-updates

Made creators more robust.

depthai_nodes/ml/messages/creators/classification.py

@@ -30,8 +30,7 @@ def create_classification_message(
     @raises ValueError: If the provided scores do not sum to 1.
     @raises ValueError: If the number of labels and scores mismatch.
     """
-
-    if type(classes) == type(None):
+    if isinstance(classes, type(None)):
         raise ValueError("Classes should not be None.")
 
     if not isinstance(classes, list):
@@ -62,7 +61,7 @@ def create_classification_message(
     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):
+    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):
@@ -71,8 +70,7 @@ def create_classification_message(
         )
 
     classification_msg = Classifications()
-
-    sorted_args = np.argsort(scores)[::-1]
+    sorted_args = np.argsort(-scores, kind="stable")
     scores = scores[sorted_args]
 
     classification_msg.classes = [classes[i] for i in sorted_args]

depthai_nodes/ml/messages/creators/depth.py

@@ -7,7 +7,7 @@
 
 
 def create_depth_message(
-    depth_map: np.array, depth_type: Literal["relative", "metric"]
+    depth_map: np.ndarray, depth_type: Literal["relative", "metric"]
 ) -> dai.ImgFrame:
     """Create a DepthAI message for a depth map.
 
@@ -25,6 +25,19 @@ def create_depth_message(
 
     if not isinstance(depth_map, np.ndarray):
         raise ValueError(f"Expected numpy array, got {type(depth_map)}.")
+
+    if len(depth_map.shape) != 3:
+        raise ValueError(f"Expected 3D input, got {len(depth_map.shape)}D input.")
+
+    if depth_map.shape[0] == 1:
+        depth_map = depth_map[0, :, :]  # CHW to HW
+    elif depth_map.shape[2] == 1:
+        depth_map = depth_map[:, :, 0]  # HWC to HW
+    else:
+        raise ValueError(
+            f"Unexpected image shape. Expected CHW or HWC, got {depth_map.shape}."
+        )
+
     if len(depth_map.shape) != 2:
         raise ValueError(f"Expected 2D input, got {len(depth_map.shape)}D input.")
 

depthai_nodes/ml/messages/creators/detection.py

@@ -48,64 +48,73 @@ def create_detection_message(
 
     # checks for bboxes
     if not isinstance(bboxes, np.ndarray):
-        raise ValueError(f"bboxes should be numpy array, got {type(bboxes)}.")
-    if len(bboxes) != 0:
-        if len(bboxes.shape) != 2:
-            raise ValueError(
-                f"bboxes 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"bboxes 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(
-                "bboxes should be in format [x_min, y_min, x_max, y_max] where xmin < xmax and ymin < ymax."
-            )
+        raise ValueError(f"Bounding boxes should be numpy array, got {type(bboxes)}.")
+    if len(bboxes) == 0:
+        raise ValueError("Bounding boxes should not be empty.")
+
+    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)}.")
+        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}."
+                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]}."
+                f"Scores should have same length as bboxes, got {scores.shape[0]} and {bboxes.shape[0]}."
             )
 
     # checks for labels
-    if labels is not None and len(labels) != 0:
+    if labels is not None:
         if not isinstance(labels, List):
-            raise ValueError(f"labels should be list, got {type(labels)}.")
+            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)}."
+                    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]}."
+                f"Labels should have same length as bboxes, got {len(labels)} and {bboxes.shape[0]}."
             )
 
     # checks for keypoints
-    if keypoints is not None and len(keypoints) != 0:
+    if keypoints is not None:
         if not isinstance(keypoints, List):
-            raise ValueError(f"keypoints should be list, got {type(keypoints)}.")
+            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:
             for point in object_keypoints:
                 if not isinstance(point, Tuple) and not isinstance(point, List):
                     raise ValueError(
-                        f"keypoint pairs should be list of tuples, got {type(point)}."
+                        f"Keypoint pairs should be list of tuples, got {type(point)}."
+                    )
+                if len(point) != 2:
+                    raise ValueError(
+                        f"Keypoint pairs should be list of tuples of length 2, got {len(point)}."
                     )
-        if len(keypoints) != bboxes.shape[0]:
-            raise ValueError(
-                f"keypoints should have same length as bboxes, got {len(keypoints)} and {bboxes.shape[0]}."
-            )
 
     if keypoints is not None:
         img_detection = ImgDetectionWithKeypoints
@@ -156,29 +165,36 @@ def create_line_detection_message(lines: np.ndarray, scores: np.ndarray):
 
     # 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]}."
-            )
+        raise ValueError(f"Lines should be numpy array, got {type(lines)}.")
+    if len(lines) == 0:
+        raise ValueError("Lines should not be empty.")
+
+    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}."
-            )
-        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]}."
-            )
+        raise ValueError(f"Scores should be numpy array, got {type(scores)}.")
+
+    if len(scores) == 0:
+        raise ValueError("Scores should not be empty.")
+
+    if len(scores.shape) != 1:
+        raise ValueError(f"Scores should be of shape (N,) meaning, got {scores.shape}.")
+    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]}."
+        )
+
+    for score in scores:
+        if not isinstance(score, float):
+            raise ValueError(f"Scores should be of type float, got {type(score)}.")
 
     line_detections = []
     for i, line in enumerate(lines):

depthai_nodes/ml/messages/creators/image.py

@@ -4,7 +4,7 @@
 
 
 def create_image_message(
-    image: np.array,
+    image: np.ndarray,
     is_bgr: bool = True,
 ) -> dai.ImgFrame:
     """Create a DepthAI message for an image array.
@@ -25,12 +25,15 @@ def create_image_message(
         hwc = True
     else:
         raise ValueError(
-            "Unexpected image shape. Expected CHW or HWC, got", image.shape
+            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):
+        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

depthai_nodes/ml/messages/creators/keypoints.py

@@ -31,6 +31,7 @@ def create_hand_keypoints_message(
     @raise ValueError: If the hand_keypoints 2nd dimension is not of size E{3}.
     @raise ValueError: If the handedness is not a float.
     @raise ValueError: If the confidence is not a float.
+    @raise ValueError: If the confidence_threshold is not a float.
     """
 
     if not isinstance(hand_keypoints, np.ndarray):
@@ -46,9 +47,27 @@ def create_hand_keypoints_message(
             f"hand_keypoints 2nd dimension should be of size 3 e.g. [x, y, z], got {hand_keypoints.shape[1]}."
         )
     if not isinstance(handedness, float):
-        raise ValueError(f"handedness should be float, got {type(handedness)}.")
+        raise ValueError(f"Handedness should be float, got {type(handedness)}.")
     if not isinstance(confidence, float):
-        raise ValueError(f"confidence should be float, got {type(confidence)}.")
+        raise ValueError(f"Confidence should be float, got {type(confidence)}.")
+    if not isinstance(confidence_threshold, float):
+        raise ValueError(
+            f"Confidence threshold should be float, got {type(confidence_threshold)}."
+        )
+
+    if handedness < 0 or handedness > 1:
+        raise ValueError(
+            f"Handedness should be between 0 and 1, got handedness {handedness}."
+        )
+
+    if confidence < 0 or confidence > 1:
+        raise ValueError(
+            f"Confidence should be between 0 and 1, got confidence {confidence}."
+        )
+    if confidence_threshold < 0 or confidence_threshold > 1:
+        raise ValueError(
+            f"Confidence threshold should be between 0 and 1, got confidence threshold {confidence_threshold}."
+        )
 
     hand_keypoints_msg = HandKeypoints()
     hand_keypoints_msg.handedness = handedness
@@ -93,67 +112,68 @@ def create_keypoints_message(
     @raise ValueError: If the confidence threshold is not provided when scores are provided.
     """
 
-    if not isinstance(keypoints, np.ndarray):
-        if not isinstance(keypoints, list):
-            raise ValueError(
-                f"keypoints should be numpy array or list, got {type(keypoints)}."
+    if not isinstance(keypoints, (np.ndarray, list)):
+        raise ValueError(
+            f"Keypoints should be numpy array or list, got {type(keypoints)}."
+        )
+
+    if not isinstance(scores, (np.ndarray, list)):
+        raise ValueError(f"Scores should be numpy array or list, got {type(scores)}.")
+
+    if len(keypoints) == 0:
+        raise ValueError("Keypoints should not be empty.")
+
+    if len(keypoints) != len(scores):
+        raise ValueError(
+            "Keypoints and scores should have the same length. Got {} keypoints and {} scores.".format(
+                len(keypoints), len(scores)
             )
+        )
+
+    if not all(isinstance(score, float) for score in scores):
+        raise ValueError("Scores should only contain float values.")
+    if not all(0 <= score <= 1 for score in scores):
+        raise ValueError("Scores should only contain values between 0 and 1.")
+
+    if not isinstance(confidence_threshold, float):
+        raise ValueError(
+            f"The confidence_threshold should be float, got {type(confidence_threshold)}."
+        )
+
+    if not (0 <= confidence_threshold <= 1):
+        raise ValueError(
+            f"The confidence_threshold should be between 0 and 1, got confidence_threshold {confidence_threshold}."
+        )
+
+    dimension = len(keypoints[0])
+    if dimension != 2 and dimension != 3:
+        raise ValueError(
+            f"All keypoints should be of dimension 2 or 3, got dimension {dimension}."
+        )
+
+    if isinstance(keypoints, list):
         for keypoint in keypoints:
             if not isinstance(keypoint, list):
                 raise ValueError(
-                    f"keypoints should be list of lists or np.array, got list of {type(keypoint)}."
+                    f"Keypoints should be list of lists or np.array, got list of {type(keypoint)}."
                 )
-            if len(keypoint) not in [2, 3]:
+            if len(keypoint) != dimension:
                 raise ValueError(
-                    f"keypoints inner list should be of size 2 or 3 e.g. [x, y] or [x, y, z], got {len(keypoint)}."
+                    "All keypoints have to be of same dimension e.g. [x, y] or [x, y, z], got mixed inner dimensions."
                 )
             for coord in keypoint:
                 if not isinstance(coord, (float)):
                     raise ValueError(
-                        f"keypoints inner list should contain only float, got {type(coord)}."
+                        f"Keypoints inner list should contain only float, got {type(coord)}."
                     )
-        keypoints = np.array(keypoints)
+
+    keypoints = np.array(keypoints)
+    scores = np.array(scores)
+
     if len(keypoints.shape) != 2:
         raise ValueError(
-            f"keypoints should be of shape (N,2 or 3) got {keypoints.shape}."
-        )
-    if keypoints.shape[1] not in [2, 3]:
-        raise ValueError(
-            f"keypoints 2nd dimension should be of size 2 or 3 e.g. [x, y] or [x, y, z], got {keypoints.shape[1]}."
+            f"Keypoints should be of shape (N,2 or 3) got {keypoints.shape}."
         )
-    if scores is not None:
-        if not isinstance(scores, np.ndarray):
-            if not isinstance(scores, list):
-                raise ValueError(
-                    f"scores should be numpy array or list, got {type(scores)}."
-                )
-            for score in scores:
-                if not isinstance(score, float):
-                    raise ValueError(
-                        f"scores should be list of floats or np.array, got list of {type(score)}."
-                    )
-            scores = np.array(scores)
-        if len(scores.shape) != 1:
-            raise ValueError(
-                f"scores should be of shape (N,) meaning [...,score,...], got {scores.shape}."
-            )
-        if keypoints.shape[0] != scores.shape[0]:
-            raise ValueError(
-                f"keypoints and scores should have the same length, got {keypoints.shape[0]} and {scores.shape[0]}."
-            )
-        if confidence_threshold is None:
-            raise ValueError(
-                "confidence_threshold should be provided when scores are provided."
-            )
-    if confidence_threshold is not None:
-        if not isinstance(confidence_threshold, float):
-            raise ValueError(
-                f"confidence_threshold should be float, got {type(confidence_threshold)}."
-            )
-        if scores is None:
-            raise ValueError(
-                "confidence_threshold should be provided when scores are provided."
-            )
 
     use_3d = keypoints.shape[1] == 3