Adding inital image augmentation.rs module in burn-datasets

[catch-twenty-two]

Image augmentation and preprocessing utilities.

Augmentations provides randomized image transformations commonly used to boost model performance through data augmentation.

Supported Augmentations

• random_photometric_distortion — Adjusts brightness, contrast and hue with a set probability
• random_zoom_out — Pads and resizes to simulate zooming out with a set probability.
• random_vertical_flip — Flips images vertically with a set probability.

Notes

• Augmentations are applied probabilistically and may vary on each call.
• Designed to be composable and easily integrated into preprocessing pipelines.
• Aims to improve model robustness to lighting, scale, and spatial orientation.

Checklist

• [x] Confirmed that run-checks all script has been executed.
• [X] Made sure the book is up to date with changes in this PR.

Related Issues/PRs

https://github.com/tracel-ai/burn/issues/207

https://github.com/tracel-ai/burn/issues/2361

https://github.com/tracel-ai/burn/pull/2426

Changes

This adds an initial module for image augmentation using a the image-rs library and the tensor backend where it was easily implementable.

Testing

Tested with actual images, also tested using unit tests.

[codecov[bot]]

Codecov Report

Attention: Patch coverage is 95.89041% with 15 lines in your changes missing coverage. Please review.

Project coverage is 81.22%. Comparing base (84e947d) to head (fe5746a). Report is 163 commits behind head on main.

Files with missing lines	Patch %	Lines
...ates/burn-dataset/src/vision/image_augmentation.rs	92.13%	14 Missing :warning:
crates/burn-dataset/src/vision/image_ops.rs	99.46%	1 Missing :warning:

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2995      +/-   ##
==========================================
+ Coverage   81.07%   81.22%   +0.14%     
==========================================
  Files         872      816      -56     
  Lines      121129   117481    -3648     
==========================================
- Hits        98209    95426    -2783     
+ Misses      22920    22055     -865     

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.

:rocket: New features to boost your workflow:

• :snowflake: Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[antimora]

Related:

https://github.com/tracel-ai/burn/issues/207

https://github.com/tracel-ai/burn/issues/2361

https://github.com/tracel-ai/burn/pull/2426

The previous augmentation related code are in dataset

[catch-twenty-two]

Hi @antimora I cleaned up the file a bunch and included your suggestions. I also have moved everything to burn-dataset as image_augmentation.rs. Thanks for the input!

[catch-twenty-two]

I decided to add bounding boxes to this rather than come back and do it later. Will update the pr soon.

[catch-twenty-two]

Hi @antimora I've integrated bounding box transformations, and also everything is done using tensors now. Looking for some more feedback. I haven't worked at all really on getting the integration tests running, mostly looking for feedback on code structure and any glaring errors. Thanks!

[catch-twenty-two]

should be ready to merge if everything looks good

[catch-twenty-two]

found some issues with integration, doing some more work on this

[catch-twenty-two]

@antimora and @laggui I fixed the issues I was having and have tested adding the module to a training batch by modifying this example:

https://github.com/tracel-ai/burn/blob/main/examples/custom-image-dataset/src/data.rs#L96

        let images = items
            .into_iter()
            .map(|item| TensorData::new(image_as_vec_u8(item), Shape::new([32, 32, 3])))
            .map(|data| {
                Tensor::<B, 3>::from_data(data.convert::<B::FloatElem>(), device)
                    // permute(2, 0, 1)
                    .swap_dims(2, 1) // [H, C, W]
                    .swap_dims(1, 0) // [C, H, W]
            })
            .map(|tensor| tensor / 255) // normalize between [0, 1]
            .collect();

Becomes:

        let mut aug = Augmentations::default();
        let images = items
            .into_iter()
            .map(|item| TensorData::new(image_as_vec_u8(item), Shape::new([32, 32, 3])))
            .map(|data| {
                Tensor::<B, 3>::from_data(data.convert::<B::FloatElem>(), device)
                    // permute(2, 0, 1)
                    .swap_dims(2, 1) // [H, C, W]
                    .swap_dims(1, 0) // [C, H, W]
            })
            .map(|tensor| aug.random_vertical_flip::<B>((tensor, None), 0.5).0)
            .map(|tensor| tensor / 255) // normalize between [0, 1]
            .collect();

It works great! Very simple and straight forward change to add as well.

[catch-twenty-two]

(should be ready to merge now)

[antimora]

One of the checks is failing.

[catch-twenty-two]

One of the checks is failing.

Thanks, fixed

[laggui]

Where did copilot come from and why are we asking it reviews? 🤣

A little busy right now but will take some time to review this properly early next week! I have some general thoughts mostly regarding flexibility / extensibility but will wait to do a proper review.

[antimora]

Where did copilot come from and why are we asking it reviews? 🤣

A little busy right now but will take some time to review this properly early next week! I have some general thoughts mostly regarding flexibility / extensibility but will wait to do a proper review.

I added it to see if there was any value. But I haven't seen useful results yet. Eventually maybe. I'll keep trying it out. It doesn't seem to hurt. Let me know if this annoys anyone.

[catch-twenty-two]

Where did copilot come from and why are we asking it reviews? 🤣

A little busy right now but will take some time to review this properly early next week! I have some general thoughts mostly regarding flexibility / extensibility but will wait to do a proper review.

Thanks, I updated the documentation as well

[laggui]

I added it to see if there was any value. But I haven't seen useful results yet. Eventually maybe. I'll keep trying it out. It doesn't seem to hurt. Let me know if this annoys anyone.

Not annoying, was just surprised to see this 😄 Must be a new thing they're trying to push to users. Doubt it will provide any actual value in reviews.

[github-actions[bot]]

This PR has been marked as stale because it has not been updated for over a month

[laggui]

We can probably close this in favor of #3254 @catch-twenty-two ?

[catch-twenty-two]

We can probably close this in favor of #3254 @catch-twenty-two ?

Of course, sorry I was going to do this, but it slipped my mind.

[laggui]

No worries! Didn't want to close without checking if you had other plans in mind for this