⚡️ Speed up function preprocess by 21%

[codeflash-ai]

📄 21% (0.21x) speedup for preprocess in unstructured_inference/models/yolox.py

⏱️ Runtime : 68.0 milliseconds → 56.1 milliseconds (best of 17 runs)

📝 Explanation and details

The optimized code achieves a 21% speedup by eliminating redundant operations and reducing memory allocation overhead in the image preprocessing pipeline.

Key Optimizations

1. Efficient buffer initialization with np.full: Replacing np.ones(...) * 114 with np.full(..., 114) directly creates the padded buffer with the fill value in a single operation, avoiding the allocation-then-multiply pattern. This reduces the padded image creation time from ~13.9% to ~7.6% of total runtime.

2. Precomputed dimensions eliminate redundant calculations: The original code called int(img.shape[0] * r) and int(img.shape[1] * r) multiple times across different operations (resize parameters, slicing). The optimized version computes img_h, img_w, resized_h, and resized_w once and reuses them, eliminating repeated attribute lookups and float-to-int conversions.

3. Avoided unnecessary dtype cast: The original code unconditionally called .astype(np.uint8) after cv2.resize(), but cv2.resize() already returns uint8 when the input is uint8. The optimized version conditionally checks the dtype first, avoiding a redundant copy operation in the common case. This reduces the resize operation overhead from ~25.7% to ~20.7% of runtime.

4. Combined transpose with contiguous array conversion: By passing the transposed array directly to np.ascontiguousarray() instead of first assigning and then converting, we reduce intermediate assignments, though both versions spend similar time (~55-66%) on this final memory layout operation.

Performance Context

Based on function_references, this preprocess function is called in the hot path of image_processing(), which runs YOLOX layout detection on every image. The function is invoked once per image before model inference, making these micro-optimizations worthwhile especially for batch processing scenarios.

Test Case Performance

The optimizations show consistent gains across different workloads:

• Small images (3×5, 5×5): 20-26% faster - benefits most from reduced overhead
• Medium images (200×200, 300×300): 17-21% faster - balanced improvement across all optimizations
• Test cases with varying aspect ratios all benefit, indicating the optimizations are robust to different resize scenarios

The speedup is most pronounced in the allocation phase (np.full) and when avoiding the redundant astype cast, which together account for the majority of the 21% overall improvement.

✅ Correctness verification report:

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

🌀 Click to see Generated Regression Tests

import numba  # required because the function uses numba.njit
import numpy as np  # used to construct inputs and inspect outputs
# imports
import pytest  # used for our unit tests
from unstructured_inference.models.yolox import preprocess

def test_preprocess_rgb_constant_fill_basic():
    # Create a small RGB image where each channel is constant:
    # channel 0 = 50, channel 1 = 100, channel 2 = 200
    img = np.zeros((3, 5, 3), dtype=np.uint8)
    img[..., 0] = 50
    img[..., 1] = 100
    img[..., 2] = 200

    # Target input canvas is larger so padding will be visible
    input_size = (6, 6)

    # Call preprocess (numba compiled function) with default swap
    out, r = preprocess(img, input_size) # 41.5μs -> 34.1μs (21.9% faster)

def test_preprocess_custom_swap_for_rgb():
    # Ensure that supplying a custom swap rearranges axes as expected.
    img = np.ones((5, 5, 3), dtype=np.uint8)
    img[..., 0] = 7
    img[..., 1] = 14
    img[..., 2] = 21

    input_size = (5, 5)

    # Use swap that results in shape (H,W,C) -> swap=(0,1,2) (identity)
    out_identity, r_id = preprocess(img, input_size, swap=(0, 1, 2)) # 43.0μs -> 35.8μs (20.1% faster)

    # Default swap produces (C,H,W); verify it differs from identity layout
    out_default, r_def = preprocess(img, input_size) # 19.6μs -> 15.6μs (26.1% faster)

def test_preprocess_medium_scale_performance_and_correctness():
    # Create a moderate-size image to test scalability but keep total pixels < 1000 per channel
    # Choose new_h * new_w <= 900 to avoid heavy loops per instructions.
    src_h, src_w = 15, 12  # source dimensions
    img = np.zeros((src_h, src_w, 3), dtype=np.uint8)

    # Fill channels with linear ramps so interpolation has to do non-trivial work
    for y in range(src_h):
        for x in range(src_w):
            img[y, x, 0] = (x % 256)  # varying with x
            img[y, x, 1] = (y % 256)  # varying with y
            img[y, x, 2] = ((x + y) % 256)  # combined

    # Choose input_size so that new_h * new_w <= 900 (e.g., scale up slightly)
    input_size = (20, 20)

    out, r = preprocess(img, input_size) # 47.9μs -> 39.9μs (20.0% faster)

    # Verify ratio computed is consistent with dimensions
    expected_r = min(input_size[0] / src_h, input_size[1] / src_w)

    # Check that padded areas (if any) are filled with 114.0
    # Determine new_h and new_w as function computes them
    new_h = int(src_h * r)
    new_w = int(src_w * r)
    if new_h <= 0:
        new_h = 1
    if new_w <= 0:
        new_w = 1

    # Rows from new_h .. input_size[0]-1 should be padding (114) for channels
    if new_h < input_size[0]:
        pass

    # Columns from new_w .. input_size[1]-1 should be padding (114) for channels
    if new_w < input_size[1]:
        pass
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import numpy as np
# imports
import pytest
from unstructured_inference.models.yolox import preprocess

class TestPreprocessBasicFunctionality:
    """Basic test cases for preprocess function with normal inputs."""

    def test_3d_image_basic_resize_down(self):
        """Test basic downsampling of a 3D color image to smaller input size."""
        # Create a simple 300x300x3 RGB image with uniform color
        img = np.ones((300, 300, 3), dtype=np.uint8) * 50
        input_size = (416, 416)
        
        # Call preprocess
        out, r = preprocess(img, input_size) # 880μs -> 742μs (18.6% faster)
        
        # Verify resize ratio is computed correctly (smaller dimension determines ratio)
        expected_r = 416 / 300

    def test_3d_image_basic_resize_up(self):
        """Test basic upsampling of a 3D color image to larger input size."""
        # Create a simple 200x200x3 RGB image
        img = np.ones((200, 200, 3), dtype=np.uint8) * 100
        input_size = (416, 416)
        
        # Call preprocess
        out, r = preprocess(img, input_size) # 762μs -> 649μs (17.4% faster)
        
        # Verify resize ratio
        expected_r = 416 / 200

    

To edit these changes git checkout codeflash/optimize-preprocess-mkovdz0h and push.