Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Added filtering on accept ratio #2311

Merged
merged 3 commits into from
Jan 28, 2025
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 29 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,7 @@ def __init__(
self.check_each_transform = check_each_transform
self.clip = clip
self.filter_invalid_bboxes = filter_invalid_bboxes
self.max_accept_ratio = max_accept_ratio # e.g., 5.0
ternaus marked this conversation as resolved.
Show resolved Hide resolved

def to_dict_private(self) -> dict[str, Any]:
data = super().to_dict_private()
Expand All @@ -115,6 +130,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 +609,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 +625,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 +651,22 @@ 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:
with np.errstate(divide="ignore", invalid="ignore"):
aspect_ratios = np.maximum(clipped_widths / clipped_heights, clipped_heights / clipped_widths)
ternaus marked this conversation as resolved.
Show resolved Hide resolved
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)
Loading