⚡️ Speed up function histogram_equalization by 31,942%

[codeflash-ai]

📄 31,942% (319.42x) speedup for histogram_equalization in src/signal/image.py

⏱️ Runtime : 2.26 seconds → 7.05 milliseconds (best of 250 runs)

📝 Explanation and details

The optimized code achieves a 319x speedup by replacing nested Python loops with vectorized NumPy operations, which are executed in highly optimized C code.

Key Optimizations

1. Vectorized Histogram Computation (74.4% → 27.8% of runtime)

The original code used nested loops to build the histogram:

for y in range(height):
    for x in range(width):
        histogram[image[y, x]] += 1

The optimized version uses np.bincount():

histogram = np.bincount(image.ravel(), minlength=256)[:256]

Why this is faster: bincount is a compiled C function that directly counts occurrences in a single pass, eliminating Python loop overhead and individual array indexing operations.

2. Vectorized Output Generation (74.4% → 13.6% of runtime)

The original code mapped each pixel individually:

for y in range(height):
    for x in range(width):
        equalized[y, x] = np.round(cdf[image[y, x]] * 255)

The optimized version uses fancy indexing:

mapping = np.round(cdf * 255).astype(image.dtype)
equalized = mapping[image]

Why this is faster: Pre-computing the mapping table and using advanced indexing (mapping[image]) allows NumPy to apply the transformation to all pixels in parallel, avoiding per-pixel Python interpretation.

3. Preserved Behavior

• Maintains exact CDF calculation logic using the same iterative approach to match floating-point precision
• Adds bounds checking to preserve IndexError behavior for out-of-range values

Performance Impact by Test Case

• Large images (1000×1000): 498x faster - the vectorization benefit scales dramatically with image size
• Small images (<10×10): 6-60% faster - still beneficial but with more setup overhead
• Single pixel: 3% slower - vectorization overhead exceeds benefit for trivial cases

The optimization is particularly effective for typical image processing workloads where images contain hundreds to millions of pixels, making nested loop elimination critical for performance.

✅ Correctness verification report:

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

🌀 Click to see Generated Regression Tests

import numpy as np

# imports
import pytest
from src.signal.image import histogram_equalization

# unit tests

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


def test_uniform_image():
    # All pixels have the same value, so output should be all 255
    img = np.full((4, 4), 100, dtype=np.uint8)
    codeflash_output = histogram_equalization(img)
    out = codeflash_output  # 100μs -> 71.3μs (41.1% faster)


def test_two_level_image():
    # Image with half zeros, half 255s
    img = np.array(
        [[0, 0, 255, 255], [0, 0, 255, 255], [0, 0, 255, 255], [0, 0, 255, 255]],
        dtype=np.uint8,
    )
    codeflash_output = histogram_equalization(img)
    out = codeflash_output  # 98.8μs -> 71.1μs (39.0% faster)
    # All zeros should map to 127, all 255s to 255
    expected = np.array(
        [
            [127, 127, 255, 255],
            [127, 127, 255, 255],
            [127, 127, 255, 255],
            [127, 127, 255, 255],
        ],
        dtype=np.uint8,
    )


def test_gradient_image():
    # 2x2 image with increasing values
    img = np.array([[0, 64], [128, 255]], dtype=np.uint8)
    codeflash_output = histogram_equalization(img)
    out = codeflash_output  # 76.0μs -> 71.0μs (6.98% faster)
    # Each pixel is unique, so mapping should spread them out
    # The mapping should be [0, 1/4, 2/4, 3/4, 1] * 255, but only 4 pixels
    expected = np.array([[64, 127], [191, 255]], dtype=np.uint8)


def test_small_random_image():
    # 3x3 image with random values
    img = np.array([[10, 20, 30], [40, 50, 60], [70, 80, 90]], dtype=np.uint8)
    codeflash_output = histogram_equalization(img)
    out = codeflash_output  # 84.8μs -> 71.2μs (19.1% faster)
    # Each value is unique, so output should be spread linearly
    expected = np.array(
        [[28, 57, 85], [113, 142, 170], [198, 227, 255]], dtype=np.uint8
    )


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


def test_all_zeros():
    # All pixels are zero
    img = np.zeros((5, 5), dtype=np.uint8)
    codeflash_output = histogram_equalization(img)
    out = codeflash_output  # 114μs -> 71.2μs (60.9% faster)


def test_all_255():
    # All pixels are 255
    img = np.full((3, 3), 255, dtype=np.uint8)
    codeflash_output = histogram_equalization(img)
    out = codeflash_output  # 84.7μs -> 70.7μs (19.8% faster)


def test_single_pixel():
    # 1x1 image
    img = np.array([[77]], dtype=np.uint8)
    codeflash_output = histogram_equalization(img)
    out = codeflash_output  # 69.0μs -> 71.3μs (3.27% slower)


def test_max_min_adjacent():
    # 2x2 image with min and max values
    img = np.array([[0, 255], [0, 255]], dtype=np.uint8)
    codeflash_output = histogram_equalization(img)
    out = codeflash_output  # 75.4μs -> 70.8μs (6.53% faster)
    expected = np.array([[127, 255], [127, 255]], dtype=np.uint8)


def test_non_square_image():
    # 2x4 image
    img = np.array([[10, 20, 30, 40], [50, 60, 70, 80]], dtype=np.uint8)
    codeflash_output = histogram_equalization(img)
    out = codeflash_output  # 82.3μs -> 70.8μs (16.3% faster)
    expected = np.array([[28, 57, 85, 113], [142, 170, 198, 227]], dtype=np.uint8)


def test_image_with_gaps():
    # Image with only values 0, 128, 255
    img = np.array([[0, 128, 255], [0, 128, 255], [0, 128, 255]], dtype=np.uint8)
    codeflash_output = histogram_equalization(img)
    out = codeflash_output  # 84.3μs -> 70.8μs (19.1% faster)
    # Each value appears 3 times, so mapping should be [85, 170, 255]
    expected = np.array(
        [[85, 170, 255], [85, 170, 255], [85, 170, 255]], dtype=np.uint8
    )


def test_large_constant_block():
    # Large block of same value, surrounded by different value
    img = np.full((10, 10), 50, dtype=np.uint8)
    img[0, 0] = 0
    img[-1, -1] = 255
    codeflash_output = histogram_equalization(img)
    out = codeflash_output  # 249μs -> 71.5μs (250% faster)


def test_invalid_input_type():
    # Pass a float array, should still work with integer conversion
    img = np.array([[0.0, 127.5], [255.0, 127.5]], dtype=np.float32)
    codeflash_output = histogram_equalization(img.astype(np.uint8))
    out = codeflash_output  # 75.2μs -> 71.0μs (5.93% faster)
    expected = np.array([[127, 255], [255, 255]], dtype=np.uint8)


def test_large_image_performance():
    # Large image, but within 1000x1000 pixels
    img = np.random.randint(0, 256, size=(1000, 1000), dtype=np.uint8)
    codeflash_output = histogram_equalization(img)
    out = codeflash_output  # 1.79s -> 3.57ms (49881% faster)


def test_large_low_dynamic_range():
    # Large image with only values 100, 101, 102
    img = np.random.choice([100, 101, 102], size=(500, 500)).astype(np.uint8)
    codeflash_output = histogram_equalization(img)
    out = codeflash_output  # 441ms -> 1.32ms (33432% faster)
    # Output should have only 3 unique values, spread between 0 and 255
    unique = np.unique(out)


def test_large_gradient_image():
    # Large gradient image
    img = np.linspace(0, 255, num=1000, dtype=np.uint8).reshape(10, 100)
    codeflash_output = histogram_equalization(img)
    out = codeflash_output  # 1.83ms -> 76.8μs (2286% faster)


# 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  # used for our unit tests
from src.signal.image import histogram_equalization

# unit tests

# ---------------- Basic Test Cases ----------------


def test_single_value_image():
    # Image with all pixels the same (should become all zeros)
    img = np.full((4, 4), 100, dtype=np.uint8)
    codeflash_output = histogram_equalization(img)
    eq = codeflash_output  # 103μs -> 71.3μs (44.7% faster)


def test_two_value_image():
    # Image with half pixels 0, half pixels 255
    img = np.array(
        [[0, 0, 255, 255], [0, 0, 255, 255], [0, 0, 255, 255], [0, 0, 255, 255]],
        dtype=np.uint8,
    )
    codeflash_output = histogram_equalization(img)
    eq = codeflash_output  # 101μs -> 71.2μs (42.3% faster)
    # All 0s should map to 128 (since CDF=0.5), all 255s should map to 255
    expected = np.array(
        [
            [128, 128, 255, 255],
            [128, 128, 255, 255],
            [128, 128, 255, 255],
            [128, 128, 255, 255],
        ],
        dtype=np.uint8,
    )


def test_gradient_image():
    # Image with a linear gradient from 0 to 15
    img = np.arange(16, dtype=np.uint8).reshape((4, 4))
    codeflash_output = histogram_equalization(img)
    eq = codeflash_output  # 99.2μs -> 71.4μs (39.0% faster)
    # Each value appears once, so CDF increases by 1/16 each time
    expected = np.round(np.linspace(255 / 16, 255, 16)).astype(np.uint8).reshape((4, 4))
    # But the function's CDF includes the value itself, so first pixel is 255/16, second is 2*255/16, etc.
    # Let's compute expected using the same logic as the function
    histogram = np.zeros(256, dtype=int)
    for v in img.flatten():
        histogram[v] += 1
    cdf = np.zeros(256, dtype=float)
    cdf[0] = histogram[0] / 16
    for i in range(1, 256):
        cdf[i] = cdf[i - 1] + histogram[i] / 16
    expected = np.zeros_like(img)
    for y in range(4):
        for x in range(4):
            expected[y, x] = np.round(cdf[img[y, x]] * 255)


def test_small_random_image():
    # Random 2x2 image
    img = np.array([[10, 20], [30, 40]], dtype=np.uint8)
    codeflash_output = histogram_equalization(img)
    eq = codeflash_output  # 76.1μs -> 71.4μs (6.60% faster)
    # All values unique, so should be evenly spread
    vals = sorted(img.flatten())
    val_to_cdf = {}
    for i, v in enumerate(vals):
        val_to_cdf[v] = (i + 1) / 4
    for y in range(2):
        for x in range(2):
            expected_val = np.round(val_to_cdf[img[y, x]] * 255)


# ---------------- Edge Test Cases ----------------


def test_max_value_image():
    # All pixels at maximum value
    img = np.full((3, 3), 255, dtype=np.uint8)
    codeflash_output = histogram_equalization(img)
    eq = codeflash_output  # 92.6μs -> 72.7μs (27.4% faster)


def test_min_value_image():
    # All pixels at minimum value
    img = np.zeros((3, 3), dtype=np.uint8)
    codeflash_output = histogram_equalization(img)
    eq = codeflash_output  # 86.5μs -> 71.8μs (20.5% faster)


def test_bimodal_image():
    # Image with two values, but not evenly distributed
    img = np.array(
        [[0, 0, 0, 255], [0, 0, 0, 255], [0, 0, 0, 255], [0, 0, 0, 255]], dtype=np.uint8
    )
    codeflash_output = histogram_equalization(img)
    eq = codeflash_output  # 99.0μs -> 71.2μs (38.9% faster)
    # 12 zeros, 4 255s; so CDF for 0 is 0.75, for 255 is 1.0
    expected = np.array(
        [
            [191, 191, 191, 255],
            [191, 191, 191, 255],
            [191, 191, 191, 255],
            [191, 191, 191, 255],
        ],
        dtype=np.uint8,
    )


def test_non_uint8_image():
    # Image with values outside 0-255 should error
    img = np.array([[0, 300], [400, 255]], dtype=int)
    with pytest.raises(IndexError):
        histogram_equalization(img)  # 5.46μs -> 7.04μs (22.5% slower)


def test_one_row_image():
    # Image with one row
    img = np.array([[0, 128, 255]], dtype=np.uint8)
    codeflash_output = histogram_equalization(img)
    eq = codeflash_output  # 74.2μs -> 71.6μs (3.61% faster)
    # Each value appears once, so CDF increases by 1/3
    expected = np.array([[85, 170, 255]], dtype=np.uint8)


def test_one_column_image():
    # Image with one column
    img = np.array([[0], [128], [255]], dtype=np.uint8)
    codeflash_output = histogram_equalization(img)
    eq = codeflash_output  # 72.8μs -> 71.3μs (2.16% faster)
    expected = np.array([[85], [170], [255]], dtype=np.uint8)


def test_non_square_image():
    # 2x4 image
    img = np.array([[10, 20, 30, 40], [50, 60, 70, 80]], dtype=np.uint8)
    codeflash_output = histogram_equalization(img)
    eq = codeflash_output  # 84.2μs -> 71.4μs (18.0% faster)
    # Each value appears once
    vals = sorted(img.flatten())
    val_to_cdf = {}
    for i, v in enumerate(vals):
        val_to_cdf[v] = (i + 1) / 8
    for y in range(2):
        for x in range(4):
            expected_val = np.round(val_to_cdf[img[y, x]] * 255)


def test_image_with_holes():
    # Image with some values missing in the range
    img = np.array([[0, 0, 100, 100], [200, 200, 255, 255]], dtype=np.uint8)
    codeflash_output = histogram_equalization(img)
    eq = codeflash_output  # 83.6μs -> 71.2μs (17.4% faster)
    # There are 2 of each value, so CDF increments by 0.25
    expected = np.array([[64, 64, 128, 128], [191, 191, 255, 255]], dtype=np.uint8)


# ---------------- Large Scale Test Cases ----------------


def test_large_uniform_image():
    # Large image with all pixels same value
    img = np.full((100, 100), 50, dtype=np.uint8)
    codeflash_output = histogram_equalization(img)
    eq = codeflash_output  # 18.0ms -> 126μs (14103% faster)


def test_large_gradient_image():
    # Large image with a gradient
    img = np.tile(np.linspace(0, 255, 100, dtype=np.uint8), (10, 1))
    codeflash_output = histogram_equalization(img)
    eq = codeflash_output  # 1.85ms -> 75.8μs (2335% faster)
    # Each value appears 10 times, so CDF increments by 10/1000 = 0.01 per value
    for x in range(100):
        expected = np.round(((x + 1) / 100) * 255)


def test_large_random_image():
    # Large random image
    np.random.seed(42)
    img = np.random.randint(0, 256, size=(32, 32), dtype=np.uint8)
    codeflash_output = histogram_equalization(img)
    eq = codeflash_output  # 1.90ms -> 74.7μs (2437% faster)


def test_large_bimodal_image():
    # Large image with two values
    img = np.zeros((50, 20), dtype=np.uint8)
    img[:, 10:] = 255
    codeflash_output = histogram_equalization(img)
    eq = codeflash_output  # 1.86ms -> 76.9μs (2313% faster)


def test_large_sparse_histogram():
    # Large image with only a few unique values
    img = np.random.choice([0, 128, 255], size=(30, 30), p=[0.1, 0.8, 0.1]).astype(
        np.uint8
    )
    codeflash_output = histogram_equalization(img)
    eq = codeflash_output  # 1.68ms -> 76.4μs (2094% faster)
    # Output should only contain a few unique values
    unique = np.unique(eq)
    # Check mapping is monotonic
    prev = -1
    for v in sorted(np.unique(img)):
        mapped = np.unique(eq[img == v])
        prev = mapped[0]


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

from src.signal.image import histogram_equalization

To edit these changes git checkout codeflash/optimize-histogram_equalization-mji18dom and push.