Skip to content

perf(illumination): vectorize supported batch paths - #424

Open
aswanth-07 wants to merge 2 commits into
albumentations-team:mainfrom
aswanth-07:agent/native-batch-illumination
Open

perf(illumination): vectorize supported batch paths#424
aswanth-07 wants to merge 2 commits into
albumentations-team:mainfrom
aswanth-07:agent/native-batch-illumination

Conversation

@aswanth-07

@aswanth-07 aswanth-07 commented Aug 11, 2026

Copy link
Copy Markdown
Contributor

Fixes #53.

Summary

  • applies the shared corner/Gaussian illumination gradient to the complete float32 batch when that route is faster
  • keeps guarded per-image multiplication for uint8 and clipped low-channel float32 cases that regress in one-thread measurements
  • restores single-image clipping parity for float32 brighten/darken batches and preserves empty-batch behavior
  • adds permanent ASV coverage for direct apply_to_images and Compose(images=...) across all requested modes, sizes, channels, dtypes, and batch sizes

Correctness

The focused parity matrix compares every batch result with the corresponding single-image kernel for:

  • linear, corner, and gaussian
  • positive and negative intensity
  • uint8 and float32
  • 1, 3, and 5 channels

It also covers empty batches in all three modes. The same gradient is created once and reused across the batch.

While expanding the parity matrix, I found that current main does not clip float32 corner/Gaussian batch results in the cases whose multiplier can exceed 1, although the single-image kernels are decorated with @clipped. This patch clips those batch results so the public batch route is bit-exact with the single-image contract.

Performance coverage

The ASV matrix now measures both direct and Compose routes for 256, 512, and 1024 square images; 1, 3, and 5 channels; uint8 and float32; all three modes; and batch sizes 2, 4, 8, and 16. The largest cells are intentionally left to the repository's scheduled or manual ASV workflow rather than being run on my Windows laptop.

A safe one-thread local sample at 256x256, batch size 4 showed the useful float32 native paths improving representative cells by roughly 7-28%, while the guarded uint8 routes remained effectively unchanged. Clipped 1/3-channel float32 stays on per-image multiplication because whole-batch multiplication plus the required clip regressed; 5-channel float32 still improved by roughly 15%.

Validation

  • python -m pytest tests/test_transforms.py -q -k illumination - 51 passed
  • python -m pytest tests/test_benchmark_coverage.py tests/test_select_benchmark_filters.py -q - 28 passed
  • direct and Compose ASV benchmark smoke cases - passed
  • file-scoped pre-commit hooks - all passed, including Ruff, mypy, Pyrefly, and benchmark catalog validation

AI assistance

OpenAI Codex assisted with repository analysis, implementation, tests, benchmarks, and drafting this pull request. I reviewed the resulting diff and validation output before submission.

Summary by Sourcery

Vectorize Illumination batch processing while aligning batch behavior and benchmark coverage with single-image contracts and performance requirements.

New Features:

  • Add parity tests ensuring Illumination batch results match per-image application across modes, dtypes, and channel counts.
  • Add explicit handling for empty Illumination batches to return an identically shaped and typed result.
  • Introduce dedicated ASV benchmarks for Illumination's direct batch route across modes, sizes, dtypes, channels, and batch sizes.

Bug Fixes:

  • Ensure float32 corner and gaussian Illumination batches are clipped consistently with their single-image counterparts in cases where multipliers can exceed 1.
  • Preserve correct behavior for uint8 and clipped low-channel float32 Illumination batches by guarding against performance regressions when choosing the batch path.

Enhancements:

  • Optimize Illumination's batch path by reusing a shared gradient and applying vectorized multiplication for supported float32 cases.
  • Extend benchmark coverage tooling and mappings so Illumination modes are tracked and reported as fully covered for batch performance contracts.

Tests:

  • Expand Illumination transform tests with batch-vs-per-image parity coverage and empty-batch behavior checks.
  • Add benchmark coverage tests ensuring Illumination modes and scenarios are correctly mapped to batch matrix ASV cases.

@sourcery-ai

sourcery-ai Bot commented Aug 11, 2026

Copy link
Copy Markdown

Reviewer's Guide

Vectorizes Illumination’s batch path by reusing a shared gradient and selectively applying whole-batch multiplication and clipping, adds empty-batch handling, and extends ASV benchmark and coverage mappings so Illumination’s batch performance contracts are exercised via both Compose and direct routes.

Flow diagram for vectorized Illumination.apply_to_images batch path

flowchart TD
    A[apply_to_images] --> B{"mode == linear"}
    B -- yes --> C[fpixel.apply_linear_illumination_batch]
    B -- no --> D{"images.shape[0] == 0"}
    D -- yes --> E[return images.copy]
    D -- no --> F[fpixel.create_illumination_gradient]
    F --> G[set clip_required]
    G --> H{"images.dtype == np.uint8<br/>or<br/>images.shape[-1] <= 4 and clip_required"}
    H -- yes --> I[_apply_to_batch_same_shape with albucore.multiply_by_array]
    H -- no --> J[albucore.multiply_by_array on images]
    I --> K[set result]
    J --> K
    K --> L{"images.dtype == np.float32<br/>and clip_required"}
    L -- yes --> M[albucore.clip result inplace]
    L -- no --> N[return result]
    M --> O[return clipped result]
Loading

File-Level Changes

Change Details Files
Vectorize Illumination batch processing with shared gradients, selective batch vs per-image routes, and corrected clipping semantics for float32 batches.
  • Add early return for empty batches in Illumination.apply_to_images to preserve no-op behavior.
  • Create a single illumination gradient per batch and reuse it across all images.
  • Introduce clip_required logic for corner and gaussian modes based on intensity sign.
  • Route uint8 and clipped low-channel float32 batches through per-image multiplication to avoid regressions.
  • Apply whole-batch multiplication when safe and faster, then clip float32 batch results when required to match single-image kernels.
albumentations/augmentations/pixel/color_lighting.py
Extend Illumination tests to enforce batch vs per-image parity and empty-batch behavior.
  • Add parametrized tests comparing apply_to_images against per-image apply for all modes, intensities, dtypes, and channel counts.
  • Add an explicit test for Illumination empty batches ensuring shapes and dtypes are preserved.
tests/test_transforms.py
Add direct-batch ASV benchmarks for Illumination and integrate them into the batch matrix and coverage metadata.
  • Define Illumination-specific batch size, mode, and case-name constants for the batch matrix benchmarks.
  • Register Illumination batch specs in IMAGE_BATCH_TRANSFORMS with size and batch-size matrices.
  • Derive Illumination direct-image case IDs from image batch cases and add a TimeIlluminationDirectBatchMatrix benchmark class that measures apply_to_images over the matrix.
  • Update benchmark coverage tests to assert Illumination’s batch performance contract and scenario matrix (sizes, channels, dtypes, batch sizes, scopes, and representative ASV cases).
  • Map Illumination aliases into the benchmark coverage transform name map and register the new batch_illumination_direct benchmark and its cases in the ASV catalog.
benchmark/benchmarks/test_batch_matrix.py
tests/test_benchmark_coverage.py
tools/benchmark_coverage.py

Assessment against linked issues

Issue Objective Addressed Explanation
#53 Replace the per-image loop in Illumination.apply_to_images with a native broadcasted operation over the full (N, H, W, C) batch while reusing a single gradient and preserving behavior for all modes and supported dtypes/channels.
#53 Ensure batch results are correct and match the existing per-image behavior across modes (linear, corner, gaussian), intensities, dtypes (uint8, float32), and channels (1, 3, 5), including clipping semantics and empty-batch behavior.
#53 Add and integrate performance benchmarks covering direct apply_to_images and Compose(images=...) for Illumination across the specified matrix of sizes, channels, dtypes, modes, and batch sizes, and guard any regressing paths with fallbacks.

Possibly linked issues


Tips and commands

Interacting with Sourcery

  • Trigger a new review: Comment @sourcery-ai review on the pull request.
  • Continue discussions: Reply directly to Sourcery's review comments.
  • Generate a GitHub issue from a review comment: Ask Sourcery to create an
    issue from a review comment by replying to it. You can also reply to a
    review comment with @sourcery-ai issue to create an issue from it.
  • Generate a pull request title: Write @sourcery-ai anywhere in the pull
    request title to generate a title at any time. You can also comment
    @sourcery-ai title on the pull request to (re-)generate the title at any time.
  • Generate a pull request summary: Write @sourcery-ai summary anywhere in
    the pull request body to generate a PR summary at any time exactly where you
    want it. You can also comment @sourcery-ai summary on the pull request to
    (re-)generate the summary at any time.
  • Generate reviewer's guide: Comment @sourcery-ai guide on the pull
    request to (re-)generate the reviewer's guide at any time.
  • Resolve all Sourcery comments: Comment @sourcery-ai resolve on the
    pull request to resolve all Sourcery comments. Useful if you've already
    addressed all the comments and don't want to see them anymore.
  • Dismiss all Sourcery reviews: Comment @sourcery-ai dismiss on the pull
    request to dismiss all existing Sourcery reviews. Especially useful if you
    want to start fresh with a new review - don't forget to comment
    @sourcery-ai review to trigger a new review!

Customizing Your Experience

Access your dashboard to:

  • Enable or disable review features such as the Sourcery-generated pull request
    summary, the reviewer's guide, and others.
  • Change the review language.
  • Add, remove or edit custom review instructions.
  • Adjust other review settings.

Getting Help

@aswanth-07
aswanth-07 marked this pull request as ready for review August 11, 2026 12:30
Copilot AI lite review requested due to automatic review settings August 11, 2026 12:30

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

@sourcery-ai sourcery-ai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hey - I've left some high level feedback:

  • The float32 parity test for batch vs per-image illumination uses assert_array_equal, which may be brittle for future numeric changes; consider assert_allclose with a tight tolerance instead for the floating-point cases.
  • The new benchmark coverage assertions for Illumination rely on hardcoded case IDs and a fixed covered count of 12; you might factor out helpers to derive these expectations from the scenario matrix to reduce maintenance when cases change.
Prompt for AI Agents
Please address the comments from this code review:

## Overall Comments
- The float32 parity test for batch vs per-image illumination uses `assert_array_equal`, which may be brittle for future numeric changes; consider `assert_allclose` with a tight tolerance instead for the floating-point cases.
- The new benchmark coverage assertions for Illumination rely on hardcoded case IDs and a fixed `covered` count of 12; you might factor out helpers to derive these expectations from the scenario matrix to reduce maintenance when cases change.

Sourcery is free for open source - if you like our reviews please consider sharing them ✨
Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.

@aswanth-07
aswanth-07 marked this pull request as draft August 11, 2026 12:33
@aswanth-07
aswanth-07 marked this pull request as ready for review August 11, 2026 12:42

@sourcery-ai sourcery-ai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sorry @aswanth-07, you have reached your weekly rate limit of 500000 diff characters.

Please try again later or upgrade to continue using Sourcery

@aswanth-07
aswanth-07 force-pushed the agent/native-batch-illumination branch from 4de6a45 to c2dcfc1 Compare August 15, 2026 06:34
@aswanth-07

Copy link
Copy Markdown
Contributor Author

The only failed check is the Windows/Python 3.12 compatibility shard, which stopped during environment setup because setup-uv could not fetch the Astral version manifest (fetch failed); no tests ran. All other compatibility shards and code checks passed. I cannot rerun upstream Actions jobs from the fork, so could a maintainer please rerun the failed job when convenient?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Performance] Remove the per-image loop from Illumination.apply_to_images

2 participants