feat: Add documentation examples for image preprocessing augmentation layers #21978

Abhinavexists · 2026-01-04T10:16:51Z

In reference to #21954
Added documentation examples across 20 image preprocessing augmentation layers.

Updated Argumentation Layers

gemini-code-assist · 2026-01-04T10:17:25Z

Summary of Changes

Hello @Abhinavexists, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly improves the documentation for Keras image preprocessing augmentation layers by embedding practical code examples directly within the docstrings of 20 different layers. This enhancement provides developers with immediate, runnable demonstrations of how to configure and apply various image augmentation techniques, streamlining the learning curve and accelerating the adoption of these powerful tools for building robust computer vision models.

Highlights

Documentation Examples Added: Comprehensive documentation examples have been added for 20 Keras image preprocessing augmentation layers, demonstrating their usage with sample code.
Improved Usability: The addition of practical code examples directly within the docstrings significantly enhances the clarity and ease of use for developers integrating these augmentation techniques into their models.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This pull request adds valuable documentation examples to 20 image preprocessing augmentation layers, which significantly improves the usability of these layers. The examples are generally clear and demonstrate the functionality well.

However, there are a few areas for improvement:

Consistency in Examples: Many examples are missing necessary imports (numpy, keras) and use placeholders for data (e.g., image = [...]), making them not directly runnable. It would be great to make all examples self-contained and runnable.
Bug Fix: The PR description mentions a bug fix in AugMix.get_config(), but the bug appears to be present in the provided file content. This is a critical issue that needs to be addressed.
Explicitness: For augmentation layers, it's good practice to explicitly pass training=True in the examples to make it clear that the augmentation is applied during training.

I've left specific comments with suggestions for each of these points. Addressing them will make this contribution even more impactful.

keras/src/layers/preprocessing/image_preprocessing/aug_mix.py

gemini-code-assist · 2026-01-04T10:22:50Z

keras/src/layers/preprocessing/image_preprocessing/aug_mix.py

+    # Create an AugMix layer
+    augmix = keras.layers.AugMix(
+        value_range=(0, 255),
+        num_chains=3,  # Creates 3 different augmentation chains
+        chain_depth=3, # Each chain applies up to 3 random augmentations
+        factor=0.3,    # Controls the strength of augmentations
+        all_ops=True
+    )
+
+    # Sample images
+    images = np.random.randint(0, 255, (8, 224, 224, 3), dtype='uint8')
+
+    # Each image is augmented in 3 different ways (chains) and then mixed
+    augmented_images = augmix(images, training=True)
+
+    # At inference, no augmentation is applied
+    output = augmix(images, training=False)


The code example is not runnable as it's missing the necessary imports for numpy and keras. To improve documentation quality and user experience, please make the example self-contained and runnable by adding the imports.

import numpy as np import keras # Create an AugMix layer augmix = keras.layers.AugMix( value_range=(0, 255), num_chains=3, # Creates 3 different augmentation chains chain_depth=3, # Each chain applies up to 3 random augmentations factor=0.3, # Controls the strength of augmentations all_ops=True ) # Sample images images = np.random.randint(0, 255, (8, 224, 224, 3), dtype='uint8') # Each image is augmented in 3 different ways (chains) and then mixed augmented_images = augmix(images, training=True) # At inference, no augmentation is applied output = augmix(images, training=False)

gemini-code-assist · 2026-01-04T10:22:50Z

keras/src/layers/preprocessing/image_preprocessing/cut_mix.py

+    # Create a CutMix layer in which factor controls the patch size variability
+    cutmix = keras.layers.CutMix(factor=0.8)
+
+    # Generate sample images and one-hot encoded labels
+    images = np.random.randint(0, 255, (8, 224, 224, 3), dtype='uint8')
+    labels = keras.ops.one_hot(
+        np.array([0, 1, 2, 3, 0, 1, 2, 3]),
+        num_classes=4
+    )
+
+    # Random rectangular patches are cut from one image and pasted into another
+    # Labels are also mixed proportionally to the patch area
+    output = cutmix(
+        {"images": images, "labels": labels},
+        training=True
+    )
+
+    # At inference, no augmentation is applied
+    output_inference = cutmix(
+        {"images": images, "labels": labels},
+        training=False
+    )


The code example is not runnable as it's missing the necessary imports for numpy and keras. To improve documentation quality and user experience, please make the example self-contained and runnable by adding the imports.

import numpy as np import keras # Create a CutMix layer in which factor controls the patch size variability cutmix = keras.layers.CutMix(factor=0.8) # Generate sample images and one-hot encoded labels images = np.random.randint(0, 255, (8, 224, 224, 3), dtype='uint8') labels = keras.ops.one_hot( np.array([0, 1, 2, 3, 0, 1, 2, 3]), num_classes=4 ) # Random rectangular patches are cut from one image and pasted into another # Labels are also mixed proportionally to the patch area output = cutmix( {"images": images, "labels": labels}, training=True ) # At inference, no augmentation is applied output_inference = cutmix( {"images": images, "labels": labels}, training=False )

gemini-code-assist · 2026-01-04T10:22:50Z

keras/src/layers/preprocessing/image_preprocessing/max_num_bounding_box.py

+    # Create a MaxNumBoundingBoxes layer to ensure max 10 boxes
+    max_boxes_layer = keras.layers.MaxNumBoundingBoxes(
+        max_number=10,
+        fill_value=-1
+    )
+
+    # Sample bounding boxes dict
+    bounding_boxes = {
+        "boxes": np.array([
+            [[10, 20, 100, 150], [50, 60, 200, 250], [0, 0, 50, 50]],
+        ]),
+        "labels": np.array([[1, 2, 3]])
+    }
+
+    # Ensure max 10 boxes per image
+    # If fewer than 10 boxes, pad with fill_value (-1)
+    # If more than 10 boxes, truncate to 10
+    result = max_boxes_layer({"bounding_boxes": bounding_boxes})


The code example is not runnable as it's missing the necessary imports for numpy and keras. To improve documentation quality and user experience, please make the example self-contained and runnable by adding the imports.

import numpy as np import keras # Create a MaxNumBoundingBoxes layer to ensure max 10 boxes max_boxes_layer = keras.layers.MaxNumBoundingBoxes( max_number=10, fill_value=-1 ) # Sample bounding boxes dict bounding_boxes = { "boxes": np.array([ [[10, 20, 100, 150], [50, 60, 200, 250], [0, 0, 50, 50]], ]), "labels": np.array([[1, 2, 3]]) } # Ensure max 10 boxes per image # If fewer than 10 boxes, pad with fill_value (-1) # If more than 10 boxes, truncate to 10 result = max_boxes_layer({"bounding_boxes": bounding_boxes})

keras/src/layers/preprocessing/image_preprocessing/rand_augment.py

gemini-code-assist · 2026-01-04T10:22:51Z

keras/src/layers/preprocessing/image_preprocessing/random_perspective.py

+    # Create a RandomPerspective layer with scale factor
+    # This simulates a 3D-like viewing angle shift
+    perspective_layer = keras.layers.RandomPerspective(
+        factor=1.0,
+        scale=0.3  # Control how extreme the perspective shift is
+    )
+
+    # Sample image
+    image = np.random.randint(0, 255, (224, 224, 3), dtype='uint8')
+
+    # Apply perspective transformation
+    # Different corners of the image will be shifted randomly
+    output = perspective_layer(image, training=True)


The code example is not runnable as it's missing the necessary imports for numpy and keras. To improve documentation quality and user experience, please make the example self-contained and runnable by adding the imports.

import numpy as np import keras # Create a RandomPerspective layer with scale factor # This simulates a 3D-like viewing angle shift perspective_layer = keras.layers.RandomPerspective( factor=1.0, scale=0.3 # Control how extreme the perspective shift is ) # Sample image image = np.random.randint(0, 255, (224, 224, 3), dtype='uint8') # Apply perspective transformation # Different corners of the image will be shifted randomly output = perspective_layer(image, training=True)

gemini-code-assist · 2026-01-04T10:22:51Z

keras/src/layers/preprocessing/image_preprocessing/random_posterization.py

+    # Create a RandomPosterization layer with 4 bits
+    random_posterization = keras.layers.RandomPosterization(factor=4)
+
+    # Your input image
+    image = [...]  # your input image
+
+    # Apply posterization
+    output = random_posterization(image)
+
+    # For more extreme posterization with 2 bits
+    extreme_posterization = keras.layers.RandomPosterization(
+        factor=2,
+        value_range=[0.0, 1.0]
+    )
+    output_extreme = extreme_posterization(image)


The code example uses a placeholder image = [...] and doesn't explicitly set training=True when calling the layer. For clarity and consistency with other augmentation layers, it's better to provide a runnable example and be explicit about the training mode.

import numpy as np import keras # Create a RandomPosterization layer with 4 bits random_posterization = keras.layers.RandomPosterization(factor=4) # Your input image image = np.random.randint(0, 255, (224, 224, 3), dtype='uint8') # Apply posterization output = random_posterization(image, training=True) # For more extreme posterization with 2 bits extreme_posterization = keras.layers.RandomPosterization( factor=2, value_range=[0.0, 1.0] ) output_extreme = extreme_posterization(image, training=True)

gemini-code-assist · 2026-01-04T10:22:51Z

keras/src/layers/preprocessing/image_preprocessing/random_sharpness.py

+    # Create a RandomSharpness layer
+    # factor can be sampled between 0.0 (full blur) and 0.5 (no change)
+    sharpness_layer = keras.layers.RandomSharpness(
+        factor=0.5,
+        value_range=(0, 255)
+    )
+
+    # Sample image
+    image = np.array([[[100, 150, 200], [50, 75, 100]]])
+
+    # Apply sharpness adjustment
+    output = sharpness_layer(image)


The code example is not runnable as it's missing the necessary imports for numpy and keras. Also, for clarity and consistency with other augmentation layers, it's better to be explicit about the training mode by passing training=True.

import numpy as np import keras # Create a RandomSharpness layer # factor can be sampled between 0.0 (full blur) and 0.5 (no change) sharpness_layer = keras.layers.RandomSharpness( factor=0.5, value_range=(0, 255) ) # Sample image image = np.random.randint(0, 255, (224, 224, 3), dtype='uint8') # Apply sharpness adjustment output = sharpness_layer(image, training=True)

gemini-code-assist · 2026-01-04T10:22:51Z

keras/src/layers/preprocessing/image_preprocessing/random_shear.py

+    shear_layer = keras.layers.RandomShear(x_factor=0.2, y_factor=0.2)
+
+    images = [...]  # your input image
+
+    # Apply random shear transformation
+    output = shear_layer(images, training=True)


The code example uses a placeholder images = [...] which makes it not runnable. To improve the documentation, please provide a complete, runnable example. This can be done by creating a sample image tensor using numpy.

import numpy as np import keras shear_layer = keras.layers.RandomShear(x_factor=0.2, y_factor=0.2) images = np.random.randint(0, 255, (2, 224, 224, 3), dtype='uint8') # Apply random shear transformation output = shear_layer(images, training=True)

gemini-code-assist · 2026-01-04T10:22:51Z

keras/src/layers/preprocessing/image_preprocessing/random_translation.py

+    translation_layer = keras.layers.RandomTranslation(
+        height_factor=0.2, width_factor=0.2
+    )
+
+    images = [...]  # your input image
+
+    # Apply random translation
+    output = translation_layer(images, training=True)


The code example uses a placeholder images = [...] which makes it not runnable. To improve the documentation, please provide a complete, runnable example. This can be done by creating a sample image tensor using numpy.

import numpy as np import keras translation_layer = keras.layers.RandomTranslation( height_factor=0.2, width_factor=0.2 ) images = np.random.randint(0, 255, (2, 224, 224, 3), dtype='uint8') # Apply random translation output = translation_layer(images, training=True)

codecov-commenter · 2026-01-04T10:25:13Z

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 82.69%. Comparing base (c67eddb) to head (2846fd5).
⚠️ Report is 24 commits behind head on master.

Additional details and impacted files

@@           Coverage Diff            @@
##           master   #21978    +/-   ##
========================================
  Coverage   82.69%   82.69%            
========================================
  Files         588      592     +4     
  Lines       61448    62072   +624     
  Branches     9622     9723   +101     
========================================
+ Hits        50812    51332   +520     
- Misses       8147     8215    +68     
- Partials     2489     2525    +36

Flag	Coverage Δ
keras	`82.52% <ø> (+<0.01%)`	⬆️
keras-jax	`61.55% <ø> (+<0.01%)`	⬆️
keras-numpy	`56.56% <ø> (-0.25%)`	⬇️
keras-openvino	`37.42% <ø> (+0.06%)`	⬆️
keras-tensorflow	`63.70% <ø> (-0.02%)`	⬇️
keras-torch	`62.47% <ø> (+0.01%)`	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

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

SamanehSaadat

Thanks for the PR!

Could you please review the comments in the examples and keep only the ones that are necessary? I think if the code is clear, it doesn't need a comment, e.g. "# Apply layer X" because it makes the example harder to read.

SamanehSaadat · 2026-01-15T00:05:53Z

keras/src/layers/preprocessing/image_preprocessing/rand_augment.py

+    (images, labels), _ = keras.datasets.cifar10.load_data()
+    images = images.astype("float32")
+
+    # Create a RandAugment layer


To make the examples more concise, we can remove comments if the code is clear. For example, here the comment doesn't add much value as it's already clear what the code is doing.

SamanehSaadat · 2026-01-15T01:00:22Z

keras/src/layers/preprocessing/image_preprocessing/rand_augment.py

+        factor=0.5  # Control the strength of augmentations (0-1 normalized)
+    )
+
+    # Apply augmentation during training


same here regarding comment

SamanehSaadat · 2026-01-15T01:39:14Z

keras/src/layers/preprocessing/image_preprocessing/max_num_bounding_box.py

+
+    # Ensure max 10 boxes per image
+    # If fewer than 10 boxes, pad with fill_value (-1)
+    # If more than 10 boxes, truncate to 10


Could you make these comments a paragraph?

SamanehSaadat · 2026-01-15T01:58:07Z

keras/src/layers/preprocessing/image_preprocessing/random_erasing.py

+    random_erasing = keras.layers.RandomErasing(factor=1.0)
+
+    # Your input image
+    image = [...]  # your input image


Is it possible to create random images similar to previous examples?

SamanehSaadat · 2026-01-15T01:59:22Z

keras/src/layers/preprocessing/image_preprocessing/random_invert.py

+    output = random_invert(image, training=True)
+
+    # For always inverting colors with custom value range
+    invert_always = keras.layers.RandomInvert(


Remove always from the variable name?

Abhinavexists · 2026-01-15T11:39:14Z

Hi @SamanehSaadat i have made the changes as suggested by you, i hopes codes are much cleaner now while maintaining important comments.

feat: Add documentation examples for image preprocessing layers

9314f59

google-ml-butler bot added the size:L label Jan 4, 2026

google-ml-butler bot assigned gbaned Jan 4, 2026

gemini-code-assist bot reviewed Jan 4, 2026

View reviewed changes

Abhinavexists mentioned this pull request Jan 4, 2026

Usages of Keras Augmentaiton Layers #21954

Open

SamanehSaadat suggested changes Jan 15, 2026

View reviewed changes

Abhinavexists added 4 commits January 15, 2026 09:25

removed comments for simplification

4efe3a6

revert few comments -> usefull ones

fcdfb8e

remove few more comments

711e1c1

sample image in random_erasing.py

0b5a671

lint fix

2846fd5

Abhinavexists mentioned this pull request Jan 15, 2026

Fix typo in contrast adjustment method #22012

Merged

feat: Add documentation examples for image preprocessing augmentation layers #21978

Are you sure you want to change the base?

feat: Add documentation examples for image preprocessing augmentation layers #21978

Uh oh!

Conversation

Abhinavexists commented Jan 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Updated Argumentation Layers

Uh oh!

gemini-code-assist bot commented Jan 4, 2026

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

gemini-code-assist bot Jan 4, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jan 4, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jan 4, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

gemini-code-assist bot Jan 4, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jan 4, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jan 4, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jan 4, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jan 4, 2026

Choose a reason for hiding this comment

Uh oh!

codecov-commenter commented Jan 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

SamanehSaadat left a comment

Choose a reason for hiding this comment

Uh oh!

SamanehSaadat Jan 15, 2026

Choose a reason for hiding this comment

Uh oh!

SamanehSaadat Jan 15, 2026

Choose a reason for hiding this comment

Uh oh!

SamanehSaadat Jan 15, 2026

Choose a reason for hiding this comment

Uh oh!

SamanehSaadat Jan 15, 2026

Choose a reason for hiding this comment

Uh oh!

SamanehSaadat Jan 15, 2026

Choose a reason for hiding this comment

Uh oh!

Abhinavexists commented Jan 15, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Abhinavexists commented Jan 4, 2026 •

edited

Loading

codecov-commenter commented Jan 4, 2026 •

edited

Loading