Skip to content

Commit

Permalink
Added filtering on accept ratio (#2311)
Browse files Browse the repository at this point in the history
* Added filtering on accept ratio

* Added filtering on accept ratio

* FIx
  • Loading branch information
ternaus authored Jan 28, 2025
1 parent 188a917 commit 7a65864
Show file tree
Hide file tree
Showing 2 changed files with 88 additions and 1 deletion.
36 changes: 35 additions & 1 deletion albumentations/core/bbox_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,8 @@ class BboxParams(Params):
"""Parameters for bounding box transforms.
Args:
format (str): Format of bounding boxes. Should be one of:
format (Literal["coco", "pascal_voc", "albumentations", "yolo"]): Format of bounding boxes.
Should be one of:
- 'coco': [x_min, y_min, width, height], e.g. [97, 12, 150, 200].
- 'pascal_voc': [x_min, y_min, x_max, y_max], e.g. [97, 12, 247, 212].
- 'albumentations': like pascal_voc but normalized in [0, 1] range, e.g. [0.2, 0.3, 0.4, 0.5].
Expand Down Expand Up @@ -58,6 +59,12 @@ class BboxParams(Params):
or boxes where x_max < x_min or y_max < y_min) at the beginning of the pipeline. If clip=True, filtering
is applied after clipping. Default: False.
max_accept_ratio (float | None): Maximum allowed aspect ratio for bounding boxes. The aspect ratio is calculated
as max(width/height, height/width), so it's always >= 1. Boxes with aspect ratio greater than this value
will be filtered out. For example, if max_accept_ratio=3.0, boxes with width:height or height:width ratios
greater than 3:1 will be removed. Set to None to disable aspect ratio filtering. Default: None.
Note:
The processing order for bounding boxes is:
1. Convert to albumentations format (normalized pascal_voc)
Expand All @@ -82,6 +89,12 @@ class BboxParams(Params):
... clip=True,
... filter_invalid_bboxes=True
... )
>>> # Create BboxParams that filters extremely elongated boxes
>>> bbox_params = BboxParams(
... format='yolo',
... max_accept_ratio=5.0, # Filter boxes with aspect ratio > 5:1
... clip=True
... )
"""

def __init__(
Expand All @@ -95,6 +108,7 @@ def __init__(
check_each_transform: bool = True,
clip: bool = False,
filter_invalid_bboxes: bool = False,
max_accept_ratio: float | None = None,
):
super().__init__(format, label_fields)
self.min_area = min_area
Expand All @@ -104,6 +118,11 @@ def __init__(
self.check_each_transform = check_each_transform
self.clip = clip
self.filter_invalid_bboxes = filter_invalid_bboxes
if max_accept_ratio is not None and max_accept_ratio < 1.0:
raise ValueError(
"max_accept_ratio must be >= 1.0 when provided, as aspect ratio is calculated as max(w/h, h/w)",
)
self.max_accept_ratio = max_accept_ratio # e.g., 5.0

def to_dict_private(self) -> dict[str, Any]:
data = super().to_dict_private()
Expand All @@ -115,6 +134,7 @@ def to_dict_private(self) -> dict[str, Any]:
"min_height": self.min_height,
"check_each_transform": self.check_each_transform,
"clip": self.clip,
"max_accept_ratio": self.max_accept_ratio,
},
)
return data
Expand Down Expand Up @@ -593,6 +613,7 @@ def filter_bboxes(
min_visibility: float = 0.0,
min_width: float = 1.0,
min_height: float = 1.0,
max_accept_ratio: float | None = None,
) -> np.ndarray:
"""Remove bounding boxes that either lie outside of the visible area by more than min_visibility
or whose area in pixels is under the threshold set by `min_area`. Also crops boxes to final image size.
Expand All @@ -608,6 +629,8 @@ def filter_bboxes(
min_visibility: Minimum fraction of area for a bounding box to remain. Default: 0.0.
min_width: Minimum width of a bounding box in pixels. Default: 0.0.
min_height: Minimum height of a bounding box in pixels. Default: 0.0.
max_accept_ratio: Maximum allowed aspect ratio, calculated as max(width/height, height/width).
Boxes with higher ratios will be filtered out. Default: None.
Returns:
numpy array of filtered bounding boxes.
Expand All @@ -632,13 +655,24 @@ def filter_bboxes(
clipped_widths = denormalized_bboxes[:, 2] - denormalized_bboxes[:, 0]
clipped_heights = denormalized_bboxes[:, 3] - denormalized_bboxes[:, 1]

# Calculate aspect ratios if needed
if max_accept_ratio is not None:
aspect_ratios = np.maximum(
clipped_widths / (clipped_heights + epsilon),
clipped_heights / (clipped_widths + epsilon),
)
valid_ratios = aspect_ratios <= max_accept_ratio
else:
valid_ratios = np.ones_like(denormalized_box_areas, dtype=bool)

# Create a mask for bboxes that meet all criteria
mask = (
(denormalized_box_areas >= epsilon)
& (clipped_box_areas >= min_area - epsilon)
& (clipped_box_areas / (denormalized_box_areas + epsilon) >= min_visibility)
& (clipped_widths >= min_width - epsilon)
& (clipped_heights >= min_height - epsilon)
& valid_ratios
)

# Apply the mask to get the filtered bboxes
Expand Down
53 changes: 53 additions & 0 deletions tests/test_bbox.py
Original file line number Diff line number Diff line change
Expand Up @@ -2214,3 +2214,56 @@ def test_compose_bbox_transform(
assert len(transformed["bboxes"]) == 0
assert len(transformed["classes"]) == 0
assert len(transformed["scores"]) == 0



@pytest.mark.parametrize(
["bboxes", "shape", "max_accept_ratio", "expected"],
[
# Normal aspect ratios (should pass)
(
np.array([[0.1, 0.1, 0.2, 0.2]], dtype=np.float32), # 1:1 ratio
{"height": 100, "width": 100},
2.0,
np.array([[0.1, 0.1, 0.2, 0.2]], dtype=np.float32),
),
# Too wide box (should be filtered)
(
np.array([[0.1, 0.1, 0.9, 0.2]], dtype=np.float32), # 8:1 ratio
{"height": 100, "width": 100},
2.0,
np.zeros((0, 4), dtype=np.float32),
),
# Too tall box (should be filtered)
(
np.array([[0.1, 0.1, 0.2, 0.9]], dtype=np.float32), # 1:8 ratio
{"height": 100, "width": 100},
2.0,
np.zeros((0, 4), dtype=np.float32),
),
# Multiple boxes with mixed ratios
(
np.array([
[0.1, 0.1, 0.2, 0.2], # 1:1 ratio (keep)
[0.3, 0.3, 0.9, 0.4], # 6:1 ratio (filter)
[0.5, 0.5, 0.6, 0.6], # 1:1 ratio (keep)
], dtype=np.float32),
{"height": 100, "width": 100},
2.0,
np.array([
[0.1, 0.1, 0.2, 0.2],
[0.5, 0.5, 0.6, 0.6],
], dtype=np.float32),
),
# None max_ratio (should not filter)
(
np.array([[0.1, 0.1, 0.9, 0.2]], dtype=np.float32),
{"height": 100, "width": 100},
None,
np.array([[0.1, 0.1, 0.9, 0.2]], dtype=np.float32),
),
],
)
def test_filter_bboxes_aspect_ratio(bboxes, shape, max_accept_ratio, expected):
filtered = filter_bboxes(bboxes, shape, max_accept_ratio=max_accept_ratio)
np.testing.assert_array_almost_equal(filtered, expected)

0 comments on commit 7a65864

Please sign in to comment.