fix kpts and bbox edge cases (#243)

Author: JSabadin

luxonis_ml/data/augmentations/base_engine.py

@@ -24,7 +24,7 @@ def __init__(
         config: Iterable[Params],
         keep_aspect_ratio: bool,
         is_validation_pipeline: bool,
-        min_bbox_visibility: float = 0,
+        min_bbox_visibility: float = 0.0,
     ):
         """Initialize augmentation pipeline from configuration.
 
@@ -52,8 +52,7 @@ def __init__(
         @param is_validation_pipeline: Whether this is a validation
             pipeline (in which case some augmentations are skipped)
         @type min_bbox_visibility: float
-        @param min_bbox_visibility: Minimum area of a bounding box to be
-            considered visible.
+        @param min_bbox_visibility: Minimum fraction of the original bounding box that must remain visible after augmentation.
         """
         ...
 

luxonis_ml/data/augmentations/utils.py

@@ -39,15 +39,24 @@ def postprocess_mask(mask: np.ndarray) -> np.ndarray:
 
 
 def postprocess_bboxes(bboxes: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+    area_threshold = 0.0004  # 0.02 * 0.02 Small area threshold to remove invalid bboxes and respective keypoints.
     if bboxes.size == 0:
-        return np.zeros((0, 6)), np.zeros((0, 1), dtype=np.uint8)
-
+        return np.zeros((0, 5)), np.zeros((0,), dtype=np.uint8)
     ordering = bboxes[:, -1]
-    out_bboxes = bboxes[:, :-1]
-    out_bboxes[:, 2] -= out_bboxes[:, 0]
-    out_bboxes[:, 3] -= out_bboxes[:, 1]
+    raw_bboxes = bboxes[:, :-1]
+    raw_bboxes[:, 2] -= raw_bboxes[:, 0]
+    raw_bboxes[:, 3] -= raw_bboxes[:, 1]
+    widths = raw_bboxes[:, 2]
+    heights = raw_bboxes[:, 3]
+    areas = widths * heights
+
+    valid_mask = areas >= area_threshold
+    raw_bboxes = raw_bboxes[valid_mask]
+    refined_ordering = ordering[valid_mask]
+
+    out_bboxes = raw_bboxes[:, [4, 0, 1, 2, 3]]
 
-    return out_bboxes[:, [4, 0, 1, 2, 3]], ordering.astype(np.uint8)
+    return out_bboxes, refined_ordering.astype(np.uint8)
 
 
 def postprocess_keypoints(
@@ -57,12 +66,31 @@ def postprocess_keypoints(
     image_width: int,
     n_keypoints: int,
 ) -> np.ndarray:
-    keypoints = np.reshape(keypoints[:, :3], (-1, n_keypoints * 3))[
-        bboxes_ordering
-    ]
-    np.maximum(keypoints, 0, out=keypoints)
-    keypoints[..., ::3] /= image_width
-    keypoints[..., 1::3] /= image_height
+    keypoints = keypoints[:, : (n_keypoints * 3)]
+    keypoints = keypoints.reshape(-1, n_keypoints, 3)
+
+    keypoints = keypoints[bboxes_ordering]
+
+    x = keypoints[..., 0]
+    y = keypoints[..., 1]
+    v = keypoints[..., 2]
+
+    in_bounds = (x >= 0) & (x < image_width) & (y >= 0) & (y < image_height)
+
+    v[~in_bounds] = 0
+
+    x = np.clip(x, 0, image_width)
+    y = np.clip(y, 0, image_height)
+
+    x /= image_width
+    y /= image_height
+
+    keypoints[..., 0] = x
+    keypoints[..., 1] = y
+    keypoints[..., 2] = v
+
+    keypoints = keypoints.reshape(-1, n_keypoints * 3)
+
     return keypoints