Make bbox_area_threshold and min_bbox_visibility configurable (#348)

Co-authored-by: Nikita Sokovnin <[email protected]>

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]

luxonis_ml/data/loaders/luxonis_loader.py

@@ -48,6 +48,8 @@ def __init__(
         exclude_empty_annotations: bool = False,
         color_space: Literal["RGB", "BGR", "GRAY"] = "RGB",
         seed: int | None = None,
+        min_bbox_visibility: float = 0.0,
+        bbox_area_threshold: float = 0.0004,
         *,
         keep_categorical_as_strings: bool = False,
         update_mode: UpdateMode | Literal["all", "missing"] = UpdateMode.ALL,
@@ -95,6 +97,11 @@ def __init__(
             to C{"RGB"}.
         @type seed: Optional[int]
         @param seed: The random seed to use for the augmentations.
+        @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 exclude_empty_annotations: bool
         @param exclude_empty_annotations: Whether to exclude
             empty annotations from the final label dictionary.
@@ -218,6 +225,8 @@ def __init__(
             width,
             keep_aspect_ratio,
             seed,
+            min_bbox_visibility,
+            bbox_area_threshold,
         )
 
     @override
@@ -434,6 +443,8 @@ def _init_augmentations(
         width: int | None,
         keep_aspect_ratio: bool,
         seed: int | None = None,
+        min_bbox_visibility: float = 0.0,
+        bbox_area_threshold: float = 0.0004,
     ) -> AugmentationEngine | None:
         if isinstance(augmentation_config, PathType):
             with open(augmentation_config) as file:
@@ -475,6 +486,8 @@ def _init_augmentations(
             keep_aspect_ratio=keep_aspect_ratio,
             is_validation_pipeline="train" not in self.view,
             seed=seed,
+            min_bbox_visibility=min_bbox_visibility,
+            bbox_area_threshold=bbox_area_threshold,
         )
 
     def _precompute_image_paths(self) -> None: