⚡️ Speed up function get_obb_size_category by 40%

[codeflash-ai]

📄 40% (0.40x) speedup for get_obb_size_category in supervision/metrics/utils/object_size.py

⏱️ Runtime : 243 microseconds → 173 microseconds (best of 541 runs)

📝 Explanation and details

Here are some optimizations for the given code.

1. Calculate the area in a vectorized manner using numpy operations.
2. Optimize the comparison operations using numpy's boolean indexing for assigning category values.
3. Avoid extra space used by the result array initialization with np.full.

This should make the function run faster by leveraging more numpy vectorized operations and avoiding unnecessary initializations.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 30 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Generated Regression Tests Details

from __future__ import annotations

from enum import Enum

import numpy as np
import numpy.typing as npt
# imports
import pytest  # used for our unit tests
from supervision.metrics.utils.object_size import get_obb_size_category


class ObjectSizeCategory(Enum):
    ANY = -1
    SMALL = 1
    MEDIUM = 2
    LARGE = 3

SIZE_THRESHOLDS = (32**2, 96**2)
from supervision.metrics.utils.object_size import get_obb_size_category

# unit tests

def test_single_small_bounding_box():
    # Single small bounding box
    bbox = np.array([[[0, 0], [1, 0], [1, 1], [0, 1]]], dtype=np.float32)
    expected = np.array([ObjectSizeCategory.SMALL.value])

def test_single_medium_bounding_box():
    # Single medium bounding box
    bbox = np.array([[[0, 0], [50, 0], [50, 50], [0, 50]]], dtype=np.float32)
    expected = np.array([ObjectSizeCategory.MEDIUM.value])

def test_single_large_bounding_box():
    # Single large bounding box
    bbox = np.array([[[0, 0], [100, 0], [100, 100], [0, 100]]], dtype=np.float32)
    expected = np.array([ObjectSizeCategory.LARGE.value])

def test_mixed_sizes():
    # Mixed sizes
    bboxes = np.array([
        [[0, 0], [1, 0], [1, 1], [0, 1]],
        [[0, 0], [50, 0], [50, 50], [0, 50]],
        [[0, 0], [100, 0], [100, 100], [0, 100]]
    ], dtype=np.float32)
    expected = np.array([
        ObjectSizeCategory.SMALL.value,
        ObjectSizeCategory.MEDIUM.value,
        ObjectSizeCategory.LARGE.value
    ])

def test_all_small_bounding_boxes():
    # All small bounding boxes
    bboxes = np.array([[[0, 0], [1, 0], [1, 1], [0, 1]]] * 5, dtype=np.float32)
    expected = np.array([ObjectSizeCategory.SMALL.value] * 5)

def test_all_medium_bounding_boxes():
    # All medium bounding boxes
    bboxes = np.array([[[0, 0], [50, 0], [50, 50], [0, 50]]] * 5, dtype=np.float32)
    expected = np.array([ObjectSizeCategory.MEDIUM.value] * 5)

def test_all_large_bounding_boxes():
    # All large bounding boxes
    bboxes = np.array([[[0, 0], [100, 0], [100, 100], [0, 100]]] * 5, dtype=np.float32)
    expected = np.array([ObjectSizeCategory.LARGE.value] * 5)

def test_minimum_area():
    # Minimum area
    bbox = np.array([[[0, 0], [0, 0], [0, 0], [0, 0]]], dtype=np.float32)
    expected = np.array([ObjectSizeCategory.SMALL.value])

def test_maximum_area_just_below_threshold():
    # Maximum area just below the small threshold
    bbox = np.array([[[0, 0], [31.9, 0], [31.9, 31.9], [0, 31.9]]], dtype=np.float32)
    expected = np.array([ObjectSizeCategory.SMALL.value])

def test_maximum_area_just_above_threshold():
    # Maximum area just above the small threshold
    bbox = np.array([[[0, 0], [32.1, 0], [32.1, 32.1], [0, 32.1]]], dtype=np.float32)
    expected = np.array([ObjectSizeCategory.MEDIUM.value])

def test_incorrect_shape():
    # Incorrect shape
    bbox = np.array([[[0, 0], [1, 0], [1, 1]]], dtype=np.float32)  # Shape (1, 3, 2)
    with pytest.raises(ValueError):
        get_obb_size_category(bbox)


def test_empty_array():
    # Empty array
    bbox = np.array([], dtype=np.float32).reshape(0, 4, 2)
    expected = np.array([], dtype=np.int_)

def test_large_number_of_small_bounding_boxes():
    # Large number of small bounding boxes
    bboxes = np.array([[[0, 0], [1, 0], [1, 1], [0, 1]]] * 1000, dtype=np.float32)
    expected = np.array([ObjectSizeCategory.SMALL.value] * 1000)

def test_large_number_of_mixed_bounding_boxes():
    # Large number of mixed bounding boxes
    bboxes = np.array([[[0, 0], [1, 0], [1, 1], [0, 1]],
                       [[0, 0], [50, 0], [50, 50], [0, 50]],
                       [[0, 0], [100, 0], [100, 100], [0, 100]]] * 333, dtype=np.float32)
    expected = np.array([ObjectSizeCategory.SMALL.value,
                         ObjectSizeCategory.MEDIUM.value,
                         ObjectSizeCategory.LARGE.value] * 333)

def test_large_number_of_large_bounding_boxes():
    # Large number of large bounding boxes
    bboxes = np.array([[[0, 0], [100, 0], [100, 100], [0, 100]]] * 1000, dtype=np.float32)
    expected = np.array([ObjectSizeCategory.LARGE.value] * 1000)

def test_exactly_at_small_threshold():
    # Exactly at small threshold
    bbox = np.array([[[0, 0], [32, 0], [32, 32], [0, 32]]], dtype=np.float32)
    expected = np.array([ObjectSizeCategory.MEDIUM.value])

def test_exactly_at_large_threshold():
    # Exactly at large threshold
    bbox = np.array([[[0, 0], [96, 0], [96, 96], [0, 96]]], dtype=np.float32)
    expected = np.array([ObjectSizeCategory.LARGE.value])
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from __future__ import annotations

from enum import Enum

import numpy as np
import numpy.typing as npt
# imports
import pytest  # used for our unit tests
from supervision.metrics.utils.object_size import get_obb_size_category


class ObjectSizeCategory(Enum):
    ANY = -1
    SMALL = 1
    MEDIUM = 2
    LARGE = 3

SIZE_THRESHOLDS = (32**2, 96**2)
from supervision.metrics.utils.object_size import get_obb_size_category

# unit tests

def test_single_small_box():
    # Single small box with area less than 1024 (32^2)
    boxes = np.array([[[0, 0], [0, 1], [1, 1], [1, 0]]], dtype=np.float32)
    expected = np.array([ObjectSizeCategory.SMALL.value])

def test_single_medium_box():
    # Single medium box with area between 1024 and 9216 (96^2)
    boxes = np.array([[[0, 0], [0, 10], [10, 10], [10, 0]]], dtype=np.float32)
    expected = np.array([ObjectSizeCategory.MEDIUM.value])

def test_single_large_box():
    # Single large box with area greater than or equal to 9216
    boxes = np.array([[[0, 0], [0, 100], [100, 100], [100, 0]]], dtype=np.float32)
    expected = np.array([ObjectSizeCategory.LARGE.value])

def test_mixed_boxes():
    # Mixed small, medium, and large boxes
    boxes = np.array([
        [[0, 0], [0, 1], [1, 1], [1, 0]],
        [[0, 0], [0, 10], [10, 10], [10, 0]],
        [[0, 0], [0, 100], [100, 100], [100, 0]]
    ], dtype=np.float32)
    expected = np.array([
        ObjectSizeCategory.SMALL.value,
        ObjectSizeCategory.MEDIUM.value,
        ObjectSizeCategory.LARGE.value
    ])

def test_minimum_area_box():
    # Minimum area box (degenerate case)
    boxes = np.array([[[0, 0], [0, 0], [0, 0], [0, 0]]], dtype=np.float32)
    expected = np.array([ObjectSizeCategory.SMALL.value])

def test_non_oriented_box():
    # Non-oriented bounding box (vertices do not form a valid polygon)
    boxes = np.array([[[0, 0], [1, 1], [2, 2], [3, 3]]], dtype=np.float32)
    expected = np.array([ObjectSizeCategory.SMALL.value])

def test_threshold_boundary_cases():
    # Bounding boxes with areas exactly at the small and large thresholds
    boxes = np.array([
        [[0, 0], [0, 32], [32, 32], [32, 0]],
        [[0, 0], [0, 96], [96, 96], [96, 0]]
    ], dtype=np.float32)
    expected = np.array([
        ObjectSizeCategory.MEDIUM.value,
        ObjectSizeCategory.LARGE.value
    ])

def test_incorrect_shape():
    # Inputs that do not conform to the expected shape (N, 4, 2)
    with pytest.raises(ValueError):
        get_obb_size_category(np.array([[[0, 0], [0, 1], [1, 1]]], dtype=np.float32))
    with pytest.raises(ValueError):
        get_obb_size_category(np.array([[[0, 0], [0, 1], [1, 1], [1, 0], [0, 0]]], dtype=np.float32))
    with pytest.raises(ValueError):
        get_obb_size_category(np.array([[0, 0], [0, 1], [1, 1], [1, 0]], dtype=np.float32))


def test_large_number_of_boxes():
    # Large array of bounding boxes to test the function's performance and scalability
    boxes = np.random.rand(10000, 4, 2).astype(np.float32) * 100
    codeflash_output = get_obb_size_category(boxes)

def test_high_resolution_boxes():
    # Bounding boxes with very high coordinate values to test numerical stability
    boxes = np.array([[[0, 0], [0, 10000], [10000, 10000], [10000, 0]]], dtype=np.float32)
    expected = np.array([ObjectSizeCategory.LARGE.value])

def test_stress_test_max_size():
    # Stress test with maximum size of input that can be handled by the system
    max_boxes = 1000  # Adjust based on system capability
    boxes = np.random.rand(max_boxes, 4, 2).astype(np.float32) * 100
    codeflash_output = get_obb_size_category(boxes)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.