Fix resize boxes to visible area (#2372)

* Fix

* Fix in resize_boxes_to_visible_area

* Cleanup

* Cleanup

Author: ternaus

.pre-commit-config.yaml

@@ -51,7 +51,7 @@ repos:
         files: setup.py
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.9.7
+    rev: v0.9.8
     hooks:
       # Run the linter.
       - id: ruff

albumentations/augmentations/dropout/functional.py

@@ -242,14 +242,12 @@ def filter_keypoints_in_holes(keypoints: np.ndarray, holes: np.ndarray) -> np.nd
     return keypoints[valid_keypoints]
 
 
+@handle_empty_array("bboxes")
 def resize_boxes_to_visible_area(
     boxes: np.ndarray,
     hole_mask: np.ndarray,
 ) -> np.ndarray:
     """Resize boxes to their largest visible rectangular regions."""
-    if len(boxes) == 0:
-        return boxes
-
     # Extract box coordinates
     x1 = boxes[:, 0].astype(int)
     y1 = boxes[:, 1].astype(int)
@@ -264,11 +262,6 @@ def resize_boxes_to_visible_area(
 
     for i, (visible, box) in enumerate(zip(visible_areas, boxes)):
         if not visible.any():
-            # Box is fully covered - handle directly
-            new_box = box.copy()
-
-            new_box[2:] = new_box[:2]  # collapse to point
-            new_boxes.append(new_box)
             continue
 
         # Find visible coordinates
@@ -278,16 +271,18 @@ def resize_boxes_to_visible_area(
         y_coords = np.nonzero(y_visible)[0]
         x_coords = np.nonzero(x_visible)[0]
 
-        # Create new box
-        new_box = boxes[i].copy()
+        # Update only the coordinate part of the box
+        new_box = box.copy()
         new_box[0] = x1[i] + x_coords[0]  # x_min
         new_box[1] = y1[i] + y_coords[0]  # y_min
         new_box[2] = x1[i] + x_coords[-1] + 1  # x_max
         new_box[3] = y1[i] + y_coords[-1] + 1  # y_max
 
         new_boxes.append(new_box)
 
-    return np.array(new_boxes)
+        # Return empty array with correct shape if all boxes were removed
+
+    return np.array(new_boxes) if new_boxes else np.zeros((0, boxes.shape[1]), dtype=boxes.dtype)
 
 
 def filter_bboxes_by_holes(

tests/functional/test_dropout.py

@@ -225,7 +225,7 @@ def test_cutout_various_types_and_fills(dtype, max_value, shape, fill_type):
         (
             np.array([[5, 5, 8, 8]]),
             np.ones((100, 100), dtype=np.uint8),  # Full image covered
-            np.array([[5, 5, 5, 5]])  # Box collapses to point when fully covered
+            np.zeros((0, 4), dtype=np.int32)  # Empty array with correct shape
         ),
 
         # Test case 3: Box partially covered by hole
@@ -245,7 +245,6 @@ def test_cutout_various_types_and_fills(dtype, max_value, shape, fill_type):
             np.zeros((100, 100), dtype=np.uint8).copy(),  # Start with all visible
             np.array([
                 [0, 0, 10, 10],     # Box 1: unchanged (no hole)
-                [20, 20, 20, 20],   # Box 2: collapsed (fully covered)
                 [40, 40, 45, 50],   # Box 3: width reduced (partially covered)
             ])
         ),

tests/test_bbox.py

@@ -28,6 +28,9 @@
 from albumentations.core.composition import BboxParams, Compose, ReplayCompose
 from albumentations.core.transforms_interface import BasicTransform, NoOp
 
+from albumentations.augmentations.dropout.functional import resize_boxes_to_visible_area
+
+
 
 @pytest.mark.parametrize(
     "bboxes, image_shape, expected",
@@ -2388,3 +2391,82 @@ def test_compose_max_accept_ratio_all_formats(bbox_format, bboxes, shape, max_ac
 
     result = transform(**data)
     np.testing.assert_array_almost_equal(result["bboxes"], expected, decimal=5)
+
+
+def test_resize_boxes_to_visible_area_removes_fully_covered_boxes():
+    """Test that resize_boxes_to_visible_area removes boxes with zero visible area."""
+    # Create bounding boxes with additional columns for encoded labels
+    # Format: [x_min, y_min, x_max, y_max, encoded_label_1, encoded_label_2]
+    boxes = np.array([
+        [10, 10, 30, 30, 1, 2],  # Box 1 with encoded labels 1, 2 - will be fully covered
+        [40, 40, 60, 60, 3, 4],  # Box 2 with encoded labels 3, 4 - will remain visible
+    ], dtype=np.float32)
+
+    # Create a hole mask that completely covers the first box
+    hole_mask = np.zeros((100, 100), dtype=np.uint8)
+    hole_mask[10:30, 10:30] = 1  # Fully cover first box
+
+    # Update boxes using the function
+    updated_boxes = resize_boxes_to_visible_area(boxes, hole_mask)
+
+    # Assertions
+    assert len(updated_boxes) == 1, "Should remove fully covered boxes"
+    assert updated_boxes.shape == (1, 6), "Should preserve the shape with correct number of columns"
+
+    # The remaining box should be the second one
+    np.testing.assert_array_equal(updated_boxes[0, :4], boxes[1, :4], "Second box coordinates should be unchanged")
+    assert updated_boxes[0, 4] == 3, "Second box should preserve first encoded label"
+    assert updated_boxes[0, 5] == 4, "Second box should preserve second encoded label"
+
+def test_resize_boxes_to_visible_area_with_partially_covered_boxes():
+    """Test that resize_boxes_to_visible_area correctly handles partially covered boxes."""
+    # Create bounding boxes with additional columns for encoded labels
+    boxes = np.array([
+        [10, 10, 30, 30, 1, 2],  # Box that will be partially covered
+    ], dtype=np.float32)
+
+    # Create a hole mask that covers the left half of the box
+    hole_mask = np.zeros((100, 100), dtype=np.uint8)
+    hole_mask[10:30, 10:20] = 1  # Cover left half of the box
+
+    # Update boxes using the function
+    updated_boxes = resize_boxes_to_visible_area(boxes, hole_mask)
+
+    # Assertions
+    assert len(updated_boxes) == 1, "Should keep partially covered boxes"
+    assert updated_boxes.shape == (1, 6), "Should preserve the shape with correct number of columns"
+    assert updated_boxes[0, 0] > boxes[0, 0], "X min should increase (left part covered)"
+    assert updated_boxes[0, 2] == boxes[0, 2], "X max should remain the same"
+    assert updated_boxes[0, 4] == 1, "Should preserve first encoded label"
+    assert updated_boxes[0, 5] == 2, "Should preserve second encoded label"
+
+def test_resize_boxes_to_visible_area_with_all_boxes_covered():
+    """Test that resize_boxes_to_visible_area returns empty array when all boxes are covered."""
+    # Create bounding boxes with additional columns
+    boxes = np.array([
+        [10, 10, 30, 30, 1, 2],  # Box that will be fully covered
+    ], dtype=np.float32)
+
+    # Create a hole mask that completely covers the box
+    hole_mask = np.zeros((100, 100), dtype=np.uint8)
+    hole_mask[10:30, 10:30] = 1  # Fully cover the box
+
+    # Update boxes using the function
+    updated_boxes = resize_boxes_to_visible_area(boxes, hole_mask)
+
+    # Assertions
+    assert len(updated_boxes) == 0, "Should return empty array when all boxes are covered"
+    assert updated_boxes.shape == (0, 6), "Empty array should have correct shape with all columns"
+
+def test_resize_boxes_to_visible_area_with_empty_input():
+    """Test that resize_boxes_to_visible_area handles empty input correctly."""
+    # Empty boxes array with correct shape (0 boxes, 6 columns)
+    boxes = np.zeros((0, 6), dtype=np.float32)
+    hole_mask = np.zeros((100, 100), dtype=np.uint8)
+
+    # Update boxes using the function
+    updated_boxes = resize_boxes_to_visible_area(boxes, hole_mask)
+
+    # Assertions
+    assert len(updated_boxes) == 0, "Should return empty array"
+    assert updated_boxes.shape == (0, 6), "Should preserve the shape with correct number of columns"