From 1c92a2b0065fcff428ebec5c89e4e55a0e3c239f Mon Sep 17 00:00:00 2001
From: Thomas Coldwell <31568400+thomas-coldwell@users.noreply.github.com>
Date: Thu, 29 Aug 2024 16:29:37 +0200
Subject: [PATCH] feat: Add `minimum_box_area_ratio` to filter out poor,
 partial boxes

---
 .../preprocessing/bounding_box/jittered_resize_demo.py   | 1 +
 keras_cv/src/bounding_box/utils.py                       | 9 ++++++---
 keras_cv/src/layers/preprocessing/jittered_resize.py     | 4 ++++
 3 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/examples/layers/preprocessing/bounding_box/jittered_resize_demo.py b/examples/layers/preprocessing/bounding_box/jittered_resize_demo.py
index a41da9e8bc..d188f6c7e6 100644
--- a/examples/layers/preprocessing/bounding_box/jittered_resize_demo.py
+++ b/examples/layers/preprocessing/bounding_box/jittered_resize_demo.py
@@ -23,6 +23,7 @@ def main():
         target_size=(512, 512),
         scale_factor=(3 / 4, 4 / 3),
         bounding_box_format="xyxy",
+        minimum_box_area_ratio=0.5
     )
     result = dataset.map(jittered_resize, num_parallel_calls=tf.data.AUTOTUNE)
     demo_utils.visualize_data(result, bounding_box_format="xyxy")
diff --git a/keras_cv/src/bounding_box/utils.py b/keras_cv/src/bounding_box/utils.py
index 4f7db46299..684e44ec21 100644
--- a/keras_cv/src/bounding_box/utils.py
+++ b/keras_cv/src/bounding_box/utils.py
@@ -66,7 +66,7 @@ def _relative_area(boxes, bounding_box_format):
 
 @keras_cv_export("keras_cv.bounding_box.clip_to_image")
 def clip_to_image(
-    bounding_boxes, bounding_box_format, images=None, image_shape=None
+    bounding_boxes, bounding_box_format, images=None, image_shape=None, minimum_box_area_ratio=0.0
 ):
     """clips bounding boxes to image boundaries.
 
@@ -92,6 +92,7 @@ class ID set to -1, indicating that there is no object present in them.
         images=images,
         image_shape=image_shape,
     )
+    original_areas = _relative_area(boxes, bounding_box_format="rel_xyxy")
     boxes, classes, images, squeeze = _format_inputs(boxes, classes, images)
     x1, y1, x2, y2 = ops.split(boxes, 4, axis=-1)
     clipped_bounding_boxes = ops.concatenate(
@@ -106,6 +107,7 @@ class ID set to -1, indicating that there is no object present in them.
     areas = _relative_area(
         clipped_bounding_boxes, bounding_box_format="rel_xyxy"
     )
+    area_ratios = ops.divide(areas, original_areas)
     clipped_bounding_boxes = bounding_box.convert_format(
         clipped_bounding_boxes,
         source="rel_xyxy",
@@ -113,10 +115,11 @@ class ID set to -1, indicating that there is no object present in them.
         images=images,
         image_shape=image_shape,
     )
+    passed = ops.logical_and(areas > 0.0, area_ratios > minimum_box_area_ratio)
     clipped_bounding_boxes = ops.where(
-        ops.expand_dims(areas > 0.0, axis=-1), clipped_bounding_boxes, -1.0
+        ops.expand_dims(passed, axis=-1), clipped_bounding_boxes, -1.0
     )
-    classes = ops.where(areas > 0.0, classes, -1)
+    classes = ops.where(passed, classes, -1)
     nan_indices = ops.any(ops.isnan(clipped_bounding_boxes), axis=-1)
     classes = ops.where(nan_indices, -1, classes)
 
diff --git a/keras_cv/src/layers/preprocessing/jittered_resize.py b/keras_cv/src/layers/preprocessing/jittered_resize.py
index 79a75b4f21..a7ca7105f4 100644
--- a/keras_cv/src/layers/preprocessing/jittered_resize.py
+++ b/keras_cv/src/layers/preprocessing/jittered_resize.py
@@ -103,6 +103,7 @@ def __init__(
         crop_size=None,
         bounding_box_format=None,
         interpolation="bilinear",
+        minimum_box_area_ratio=0.0,
         seed=None,
         **kwargs,
     ):
@@ -130,6 +131,8 @@ def __init__(
         self.seed = seed
 
         self.force_output_dense_images = True
+        
+        self.minimum_box_area_ratio = minimum_box_area_ratio
 
     def compute_ragged_image_signature(self, images):
         ragged_spec = tf.RaggedTensorSpec(
@@ -252,6 +255,7 @@ def augment_bounding_boxes(
             result,
             image_shape=self.target_size + (3,),
             bounding_box_format="yxyx",
+            minimum_box_area_ratio=self.minimum_box_area_ratio
         )
         result = bounding_box.convert_format(
             result,