make bbox_area_threshold configurable

Author: JSabadin

luxonis_ml/data/augmentations/albumentations_engine.py

@@ -282,11 +282,13 @@ def __init__(
         is_validation_pipeline: bool = False,
         min_bbox_visibility: float = 0.0,
         seed: int | None = None,
+        bbox_area_threshold: float = 0.0004,
     ):
         self.targets: dict[str, TargetType] = {}
         self.target_names_to_tasks = {}
         self.n_classes = n_classes
         self.image_size = (height, width)
+        self.bbox_area_threshold = bbox_area_threshold
 
         for task, task_type in targets.items():
             target_name = self.task_to_target_name(task)
@@ -572,7 +574,9 @@ def postprocess(
             task_name = get_task_name(task)
 
             if target_type == "bboxes":
-                out_labels[task], index = postprocess_bboxes(array)
+                out_labels[task], index = postprocess_bboxes(
+                    array, self.bbox_area_threshold
+                )
                 bboxes_indices[task_name] = index
 
         for target_name, target_type in self.targets.items():

luxonis_ml/data/augmentations/base_engine.py

@@ -30,6 +30,7 @@ def __init__(
         is_validation_pipeline: bool,
         min_bbox_visibility: float = 0.0,
         seed: int | None = None,
+        bbox_area_threshold: float = 0.0004,
     ):
         """Initialize augmentation pipeline from configuration.
 
@@ -67,6 +68,9 @@ def __init__(
             pipeline (in which case some augmentations are skipped)
         @type min_bbox_visibility: float
         @param min_bbox_visibility: Minimum fraction of the original bounding box that must remain visible after augmentation.
+        @type bbox_area_threshold: float
+        @param bbox_area_threshold: Minimum area threshold for bounding boxes to be considered valid. In the range [0, 1].
+            Default is 0.0004, which corresponds to a small area threshold to remove invalid bboxes and respective keypoints.
         @type seed: Optional[int]
         @param seed: Random seed for reproducibility. If None, a random seed will be used.
             If provided, it will be used to initialize the random number generator.

luxonis_ml/data/augmentations/utils.py

@@ -39,8 +39,9 @@ def postprocess_mask(mask: np.ndarray) -> np.ndarray:
     return mask.transpose(2, 0, 1)
 
 
-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.
+def postprocess_bboxes(
+    bboxes: np.ndarray, area_threshold: float = 0.0004
+) -> tuple[np.ndarray, np.ndarray]:
     if bboxes.size == 0:
         return np.zeros((0, 5)), np.zeros((0,), dtype=np.uint8)
     ordering = bboxes[:, -1]