⚡️ Speed up function is_box_near_crop_edge by 11%

[codeflash-ai]

📄 11% (0.11x) speedup for is_box_near_crop_edge in ultralytics/models/sam/amg.py

⏱️ Runtime : 4.17 milliseconds → 3.75 milliseconds (best of 166 runs)

📝 Explanation and details

The optimization achieves an 11% speedup through several targeted improvements that reduce tensor allocations and simplify broadcasting operations:

Key Optimizations Applied:

1. Eliminated unnecessary tensor indexing: Replaced crop_box_torch[None, :] and orig_box_torch[None, :] with direct broadcasting. This removes the need to create new tensors with explicit None indexing since PyTorch's broadcasting handles the dimension expansion automatically.

2. Simplified tensor creation in uncrop_boxes_xyxy: Changed from creating a nested list [[x0, y0, x0, y0]] to a flat list [x0, y0, x0, y0], eliminating one level of tensor construction overhead. Also added proper dtype matching with dtype=boxes.dtype.

3. Replaced torch.as_tensor() with torch.tensor(): While functionally similar for this use case, torch.tensor() can be slightly more efficient when creating new tensors from lists.

4. Used torch.logical_not() instead of ~ operator: This provides a more explicit logical negation that can be better optimized by PyTorch's internal operations.

Performance Impact:

The line profiler shows the most significant improvements in the torch.isclose() operations (39.5% vs 38% and 15.3% vs 16.1% of total time), indicating that removing the [None, :] indexing reduces computational overhead. The uncrop_boxes_xyxy function also shows improvement with the simplified offset tensor creation.

Workload Implications:

Given the function reference showing is_box_near_crop_edge is called within a hot path in SAM's image segmentation pipeline (inside nested loops processing multiple crop regions and point batches), this 11% improvement will compound significantly. The function is used to filter out bounding boxes near crop edges, which happens for every batch of points processed across multiple image crops, making this optimization valuable for real-time segmentation workloads.

Test Case Performance:

The optimization shows consistent 7-18% improvements across all test cases, with larger gains on simpler cases (empty tensors, far-from-edge boxes) and smaller but still meaningful gains on complex cases with many boxes, indicating the optimization is broadly applicable.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import torch
from ultralytics.models.sam.amg import is_box_near_crop_edge

# unit tests

# --------------------------
# BASIC TEST CASES
# --------------------------


def test_single_box_exactly_on_crop_edge():
    # Box is exactly on the left edge of the crop (x0)
    boxes = torch.tensor([[0, 10, 20, 30]], dtype=torch.float)
    crop_box = [0, 0, 100, 100]
    orig_box = [0, 0, 200, 200]
    # Should be near crop edge (left), but not orig edge (since orig_box is bigger)
    codeflash_output = is_box_near_crop_edge(boxes, crop_box, orig_box)
    result = codeflash_output  # 91.7μs -> 83.2μs (10.3% faster)


def test_single_box_far_from_crop_edge():
    # Box is far from any crop edge
    boxes = torch.tensor([[30, 30, 60, 60]], dtype=torch.float)
    crop_box = [0, 0, 100, 100]
    orig_box = [0, 0, 100, 100]
    codeflash_output = is_box_near_crop_edge(boxes, crop_box, orig_box)
    result = codeflash_output  # 90.2μs -> 78.9μs (14.2% faster)


def test_box_near_crop_edge_but_also_on_orig_edge():
    # Box is near crop edge and also on orig edge, so should NOT be flagged
    boxes = torch.tensor([[0, 0, 20, 20]], dtype=torch.float)
    crop_box = [0, 0, 100, 100]
    orig_box = [0, 0, 100, 100]
    codeflash_output = is_box_near_crop_edge(boxes, crop_box, orig_box)
    result = codeflash_output  # 89.9μs -> 78.6μs (14.4% faster)


def test_multiple_boxes_mixed_positions():
    # Some boxes near crop edge, some not
    boxes = torch.tensor(
        [
            [0, 0, 10, 10],  # on crop and orig edge
            [0, 20, 10, 30],  # near crop left edge, not orig edge
            [90, 90, 100, 100],  # on crop and orig edge
            [80, 80, 90, 90],  # not near any edge
        ],
        dtype=torch.float,
    )
    crop_box = [0, 0, 100, 100]
    orig_box = [0, 0, 200, 200]
    codeflash_output = is_box_near_crop_edge(boxes, crop_box, orig_box)
    result = codeflash_output  # 90.8μs -> 81.8μs (10.9% faster)


def test_atol_affects_result():
    # Box is just outside the default atol, should be False, but True with higher atol
    boxes = torch.tensor([[21, 21, 40, 40]], dtype=torch.float)
    crop_box = [0, 0, 100, 100]
    orig_box = [0, 0, 200, 200]
    # atol=20, should be False
    codeflash_output = is_box_near_crop_edge(boxes, crop_box, orig_box, atol=20.0)
    result = codeflash_output  # 89.7μs -> 80.7μs (11.2% faster)
    # atol=22, should be True (since 21 is within 22 of 0)
    codeflash_output = is_box_near_crop_edge(boxes, crop_box, orig_box, atol=22.0)
    result = codeflash_output  # 49.6μs -> 45.0μs (10.3% faster)


# --------------------------
# EDGE TEST CASES
# --------------------------


def test_box_with_negative_coordinates():
    # Box with negative coordinates, crop_box at (10,10,110,110)
    boxes = torch.tensor([[-10, -10, 0, 0]], dtype=torch.float)
    crop_box = [10, 10, 110, 110]
    orig_box = [0, 0, 200, 200]
    # After uncropping: [-10+10, -10+10, 0+10, 0+10] = [0,0,10,10] (on orig edge)
    codeflash_output = is_box_near_crop_edge(boxes, crop_box, orig_box)
    result = codeflash_output  # 84.6μs -> 75.6μs (11.9% faster)


def test_box_with_large_coordinates_outside_crop():
    # Box is far outside crop, but not on orig edge
    boxes = torch.tensor([[200, 200, 210, 210]], dtype=torch.float)
    crop_box = [100, 100, 200, 200]
    orig_box = [0, 0, 300, 300]
    # After uncrop: [300,300,410,410] (well outside orig_box)
    codeflash_output = is_box_near_crop_edge(boxes, crop_box, orig_box)
    result = codeflash_output  # 87.4μs -> 77.1μs (13.3% faster)


def test_empty_boxes_tensor():
    # No boxes at all
    boxes = torch.empty((0, 4), dtype=torch.float)
    crop_box = [0, 0, 100, 100]
    orig_box = [0, 0, 100, 100]
    codeflash_output = is_box_near_crop_edge(boxes, crop_box, orig_box)
    result = codeflash_output  # 91.8μs -> 80.0μs (14.7% faster)


def test_box_with_zero_area():
    # Box with zero area (x0==x1, y0==y1), exactly on crop edge
    boxes = torch.tensor([[0, 0, 0, 0]], dtype=torch.float)
    crop_box = [0, 0, 100, 100]
    orig_box = [0, 0, 200, 200]
    codeflash_output = is_box_near_crop_edge(boxes, crop_box, orig_box)
    result = codeflash_output  # 87.8μs -> 80.2μs (9.54% faster)


def test_box_on_crop_right_edge_but_not_orig_right_edge():
    boxes = torch.tensor([[99, 50, 100, 80]], dtype=torch.float)
    crop_box = [0, 0, 100, 100]
    orig_box = [0, 0, 200, 200]
    codeflash_output = is_box_near_crop_edge(boxes, crop_box, orig_box)
    result = codeflash_output  # 90.7μs -> 77.8μs (16.6% faster)


def test_box_on_orig_right_edge():
    boxes = torch.tensor([[199, 50, 200, 80]], dtype=torch.float)
    crop_box = [100, 0, 200, 100]
    orig_box = [0, 0, 200, 200]
    codeflash_output = is_box_near_crop_edge(boxes, crop_box, orig_box)
    result = codeflash_output  # 89.1μs -> 77.6μs (14.7% faster)


def test_box_on_crop_bottom_edge():
    boxes = torch.tensor([[10, 99, 20, 100]], dtype=torch.float)
    crop_box = [0, 0, 100, 100]
    orig_box = [0, 0, 200, 200]
    codeflash_output = is_box_near_crop_edge(boxes, crop_box, orig_box)
    result = codeflash_output  # 88.2μs -> 74.6μs (18.2% faster)


def test_box_on_crop_top_edge():
    boxes = torch.tensor([[10, 0, 20, 1]], dtype=torch.float)
    crop_box = [0, 0, 100, 100]
    orig_box = [0, 0, 200, 200]
    codeflash_output = is_box_near_crop_edge(boxes, crop_box, orig_box)
    result = codeflash_output  # 87.5μs -> 77.7μs (12.6% faster)


def test_box_on_crop_edge_with_nonzero_crop_offset():
    # Crop box is offset in image
    boxes = torch.tensor([[0, 0, 10, 10]], dtype=torch.float)
    crop_box = [50, 50, 150, 150]
    orig_box = [0, 0, 200, 200]
    # After uncrop: [50,50,60,60], which is on crop edge, not orig edge
    codeflash_output = is_box_near_crop_edge(boxes, crop_box, orig_box)
    result = codeflash_output  # 86.0μs -> 76.7μs (12.2% faster)


def test_box_on_crop_edge_with_nonzero_crop_and_orig_offset():
    boxes = torch.tensor([[0, 0, 10, 10]], dtype=torch.float)
    crop_box = [100, 100, 200, 200]
    orig_box = [100, 100, 200, 200]
    # After uncrop: [100,100,110,110] which is on both crop and orig edge
    codeflash_output = is_box_near_crop_edge(boxes, crop_box, orig_box)
    result = codeflash_output  # 87.5μs -> 79.5μs (10.1% faster)


# --------------------------
# LARGE SCALE TEST CASES
# --------------------------


def test_many_boxes_performance_and_correctness():
    # 1000 boxes, first 10 near crop edge, rest not
    n = 1000
    boxes = torch.zeros((n, 4), dtype=torch.float)
    # First 10 boxes: on crop left edge
    for i in range(10):
        boxes[i] = torch.tensor([0, i, 10, i + 10], dtype=torch.float)
    # Rest: away from edge
    for i in range(10, n):
        boxes[i] = torch.tensor([20, 20, 30, 30], dtype=torch.float)
    crop_box = [0, 0, 100, 100]
    orig_box = [0, 0, 200, 200]
    codeflash_output = is_box_near_crop_edge(boxes, crop_box, orig_box)
    result = codeflash_output  # 149μs -> 136μs (9.15% faster)


def test_large_boxes_tensor_with_varied_edges():
    # 500 boxes on crop edge, 500 not
    n = 1000
    boxes = torch.zeros((n, 4), dtype=torch.float)
    for i in range(500):
        # On crop top edge
        boxes[i] = torch.tensor([i % 100, 0, (i % 100) + 10, 10], dtype=torch.float)
    for i in range(500, n):
        # Away from edge
        boxes[i] = torch.tensor([20, 20, 30, 30], dtype=torch.float)
    crop_box = [0, 0, 100, 100]
    orig_box = [0, 0, 200, 200]
    codeflash_output = is_box_near_crop_edge(boxes, crop_box, orig_box)
    result = codeflash_output  # 157μs -> 146μs (7.58% faster)


def test_large_boxes_tensor_all_on_orig_edge():
    # All boxes are on orig edge, none should be flagged
    n = 1000
    boxes = torch.zeros((n, 4), dtype=torch.float)
    for i in range(n):
        boxes[i] = torch.tensor([0, 0, 10, 10], dtype=torch.float)
    crop_box = [0, 0, 100, 100]
    orig_box = [0, 0, 100, 100]
    codeflash_output = is_box_near_crop_edge(boxes, crop_box, orig_box)
    result = codeflash_output  # 157μs -> 143μs (9.44% faster)


def test_large_boxes_tensor_all_far_from_edges():
    # All boxes far from any edge
    n = 1000
    boxes = torch.full((n, 4), 50.0)
    crop_box = [0, 0, 100, 100]
    orig_box = [0, 0, 100, 100]
    codeflash_output = is_box_near_crop_edge(boxes, crop_box, orig_box)
    result = codeflash_output  # 153μs -> 137μs (12.1% faster)


def test_large_boxes_tensor_with_random_positions():
    # Random positions, check at least some are flagged
    torch.manual_seed(0)
    n = 1000
    boxes = torch.randint(0, 100, (n, 4)).float()
    # Ensure boxes are valid (x0 <= x1, y0 <= y1)
    x0 = torch.min(boxes[:, 0], boxes[:, 2])
    x1 = torch.max(boxes[:, 0], boxes[:, 2])
    y0 = torch.min(boxes[:, 1], boxes[:, 3])
    y1 = torch.max(boxes[:, 1], boxes[:, 3])
    boxes = torch.stack([x0, y0, x1, y1], dim=1)
    crop_box = [0, 0, 100, 100]
    orig_box = [0, 0, 200, 200]
    codeflash_output = is_box_near_crop_edge(boxes, crop_box, orig_box)
    result = codeflash_output  # 141μs -> 134μs (4.81% faster)


# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import torch  # needed for tensors
from ultralytics.models.sam.amg import is_box_near_crop_edge

# unit tests

# ------------------------
# 1. Basic Test Cases
# ------------------------


def test_single_box_near_crop_edge():
    # Box is exactly at the left crop edge, not at image edge
    boxes = torch.tensor([[0, 10, 20, 30]], dtype=torch.float)
    crop_box = [0, 0, 100, 100]
    orig_box = [0, 0, 200, 200]
    codeflash_output = is_box_near_crop_edge(boxes, crop_box, orig_box)
    result = codeflash_output  # 107μs -> 96.7μs (11.3% faster)


def test_single_box_far_from_crop_edge():
    # Box is far from any crop edge
    boxes = torch.tensor([[40, 40, 60, 60]], dtype=torch.float)
    crop_box = [0, 0, 100, 100]
    orig_box = [0, 0, 200, 200]
    codeflash_output = is_box_near_crop_edge(boxes, crop_box, orig_box)
    result = codeflash_output  # 91.4μs -> 81.0μs (12.9% faster)


def test_single_box_near_image_edge_not_crop_edge():
    # Box is at image edge, should not be considered near crop edge
    boxes = torch.tensor([[0, 0, 20, 20]], dtype=torch.float)
    crop_box = [0, 0, 100, 100]
    orig_box = [0, 0, 100, 100]
    codeflash_output = is_box_near_crop_edge(boxes, crop_box, orig_box)
    result = codeflash_output  # 87.6μs -> 78.1μs (12.1% faster)


def test_multiple_boxes_varied_positions():
    # Multiple boxes: one near crop edge, one far, one near image edge
    boxes = torch.tensor(
        [
            [0, 0, 20, 20],  # at crop edge
            [40, 40, 60, 60],  # far from edge
            [100, 100, 120, 120],  # at crop edge (right/bottom)
        ],
        dtype=torch.float,
    )
    crop_box = [0, 0, 100, 100]
    orig_box = [0, 0, 120, 120]
    codeflash_output = is_box_near_crop_edge(boxes, crop_box, orig_box)
    result = codeflash_output  # 93.7μs -> 81.5μs (14.9% faster)


def test_box_near_crop_edge_with_custom_atol():
    # Box is just outside the default atol, but inside a larger atol
    boxes = torch.tensor([[21, 0, 41, 20]], dtype=torch.float)
    crop_box = [0, 0, 100, 100]
    orig_box = [0, 0, 200, 200]
    codeflash_output = is_box_near_crop_edge(boxes, crop_box, orig_box)
    result_default = codeflash_output  # 90.8μs -> 77.3μs (17.5% faster)
    codeflash_output = is_box_near_crop_edge(boxes, crop_box, orig_box, atol=25.0)
    result_large_atol = codeflash_output  # 50.3μs -> 46.3μs (8.61% faster)


# ------------------------
# 2. Edge Test Cases
# ------------------------


def test_empty_boxes_tensor():
    # No boxes provided
    boxes = torch.empty((0, 4), dtype=torch.float)
    crop_box = [0, 0, 100, 100]
    orig_box = [0, 0, 200, 200]
    codeflash_output = is_box_near_crop_edge(boxes, crop_box, orig_box)
    result = codeflash_output  # 86.2μs -> 75.5μs (14.2% faster)


def test_box_exactly_on_crop_and_image_edge():
    # Box at both crop and image edge, should not be considered near crop edge
    boxes = torch.tensor([[0, 0, 100, 100]], dtype=torch.float)
    crop_box = [0, 0, 100, 100]
    orig_box = [0, 0, 100, 100]
    codeflash_output = is_box_near_crop_edge(boxes, crop_box, orig_box)
    result = codeflash_output  # 84.1μs -> 78.0μs (7.86% faster)


def test_box_just_outside_tolerance():
    # Box is just outside the atol threshold
    boxes = torch.tensor([[21, 0, 41, 20]], dtype=torch.float)
    crop_box = [0, 0, 100, 100]
    orig_box = [0, 0, 200, 200]
    codeflash_output = is_box_near_crop_edge(boxes, crop_box, orig_box, atol=20.0)
    result = codeflash_output  # 91.0μs -> 77.0μs (18.3% faster)


def test_box_just_inside_tolerance():
    # Box is just inside the atol threshold
    boxes = torch.tensor([[19.9, 0, 39.9, 20]], dtype=torch.float)
    crop_box = [0, 0, 100, 100]
    orig_box = [0, 0, 200, 200]
    codeflash_output = is_box_near_crop_edge(boxes, crop_box, orig_box, atol=20.0)
    result = codeflash_output  # 87.3μs -> 77.8μs (12.2% faster)


def test_box_with_negative_coordinates():
    # Box with negative coordinates, near crop edge
    boxes = torch.tensor([[-5, -5, 15, 15]], dtype=torch.float)
    crop_box = [0, 0, 100, 100]
    orig_box = [0, 0, 200, 200]
    codeflash_output = is_box_near_crop_edge(boxes, crop_box, orig_box)
    result = codeflash_output  # 88.0μs -> 77.2μs (14.0% faster)


def test_box_with_large_coordinates():
    # Box with coordinates much larger than crop/image box
    boxes = torch.tensor([[500, 500, 600, 600]], dtype=torch.float)
    crop_box = [0, 0, 100, 100]
    orig_box = [0, 0, 200, 200]
    codeflash_output = is_box_near_crop_edge(boxes, crop_box, orig_box)
    result = codeflash_output  # 88.6μs -> 78.8μs (12.4% faster)


def test_box_with_channel_dimension():
    # Test with a 3D tensor (batch size, channels, 4)
    boxes = torch.zeros((2, 1, 4), dtype=torch.float)
    boxes[0, 0] = torch.tensor([0, 0, 20, 20])
    boxes[1, 0] = torch.tensor([40, 40, 60, 60])
    crop_box = [0, 0, 100, 100]
    orig_box = [0, 0, 200, 200]
    # Flatten the channel dimension for testing
    boxes_flat = boxes.view(-1, 4)
    codeflash_output = is_box_near_crop_edge(boxes_flat, crop_box, orig_box)
    result = codeflash_output  # 88.3μs -> 77.0μs (14.8% faster)


def test_crop_box_not_at_origin():
    # Crop box is offset from origin
    boxes = torch.tensor([[0, 0, 20, 20]], dtype=torch.float)
    crop_box = [50, 50, 150, 150]
    orig_box = [0, 0, 200, 200]
    # After uncropping, box should be at (50, 50, 70, 70)
    codeflash_output = is_box_near_crop_edge(boxes, crop_box, orig_box)
    result = codeflash_output  # 88.5μs -> 78.2μs (13.1% faster)


def test_crop_box_same_as_orig_box():
    # Crop box is the same as image box, so no box should be near crop edge
    boxes = torch.tensor([[0, 0, 20, 20]], dtype=torch.float)
    crop_box = [0, 0, 100, 100]
    orig_box = [0, 0, 100, 100]
    codeflash_output = is_box_near_crop_edge(boxes, crop_box, orig_box)
    result = codeflash_output  # 87.5μs -> 76.4μs (14.6% faster)


# ------------------------
# 3. Large Scale Test Cases
# ------------------------


def test_many_boxes_performance():
    # Test with 1000 boxes, some near crop edge, some far
    crop_box = [0, 0, 1000, 1000]
    orig_box = [0, 0, 2000, 2000]
    boxes = torch.zeros((1000, 4), dtype=torch.float)
    # First 10 boxes are near crop edge
    for i in range(10):
        boxes[i] = torch.tensor([0, i, 20, i + 20])
    # Next 990 boxes are far from edge
    for i in range(10, 1000):
        boxes[i] = torch.tensor([500, 500, 520, 520])
    codeflash_output = is_box_near_crop_edge(boxes, crop_box, orig_box)
    result = codeflash_output  # 157μs -> 147μs (6.81% faster)
    expected = [True] * 10 + [False] * 990


def test_large_box_coordinates():
    # Test with large coordinates, all boxes at crop edge
    crop_box = [1000, 1000, 2000, 2000]
    orig_box = [0, 0, 4000, 4000]
    boxes = torch.zeros((1000, 4), dtype=torch.float)
    for i in range(1000):
        boxes[i] = torch.tensor([0, 0, 20, 20])
    codeflash_output = is_box_near_crop_edge(boxes, crop_box, orig_box)
    result = codeflash_output  # 159μs -> 147μs (7.70% faster)


def test_large_scale_varied_atol():
    # Test with many boxes and varied atol
    crop_box = [0, 0, 1000, 1000]
    orig_box = [0, 0, 2000, 2000]
    boxes = torch.zeros((1000, 4), dtype=torch.float)
    # First 500 boxes just inside atol, next 500 just outside
    for i in range(500):
        boxes[i] = torch.tensor([19.9, 0, 39.9, 20])
    for i in range(500, 1000):
        boxes[i] = torch.tensor([20.1, 0, 40.1, 20])
    codeflash_output = is_box_near_crop_edge(boxes, crop_box, orig_box, atol=20.0)
    result_inside = codeflash_output  # 159μs -> 146μs (8.20% faster)
    codeflash_output = is_box_near_crop_edge(boxes, crop_box, orig_box, atol=19.9)
    result_outside = codeflash_output  # 96.3μs -> 91.2μs (5.67% faster)
    expected_inside = [True] * 500 + [False] * 500
    expected_outside = [False] * 1000


def test_large_boxes_tensor_empty():
    # Large batch, but all boxes are far from crop edge
    crop_box = [0, 0, 1000, 1000]
    orig_box = [0, 0, 2000, 2000]
    boxes = torch.full((1000, 4), 5000, dtype=torch.float)
    codeflash_output = is_box_near_crop_edge(boxes, crop_box, orig_box)
    result = codeflash_output  # 145μs -> 132μs (10.1% faster)


# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-is_box_near_crop_edge-mirazvrn and push.