From 0040861d63ed813e8e2c6b9fbdd20d5a074fdd65 Mon Sep 17 00:00:00 2001 From: v1shay Date: Wed, 12 Aug 2026 13:33:59 -0700 Subject: [PATCH] Add channel axis to random vertical flip --- dm_pix/_src/augment.py | 9 ++++++++- dm_pix/_src/augment_test.py | 6 ++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/dm_pix/_src/augment.py b/dm_pix/_src/augment.py index adad96f..4d62fd5 100644 --- a/dm_pix/_src/augment.py +++ b/dm_pix/_src/augment.py @@ -796,6 +796,7 @@ def random_flip_up_down( image: chex.Array, *, probability: chex.Numeric = 0.5, + channel_axis: int = -1, ) -> chex.Array: """Applies `flip_up_down` with a given probability. @@ -805,6 +806,7 @@ def random_flip_up_down( ...HWC or ...CHW. probability: the probability of applying flip_up_down transform. Must be a value in [0, 1]. + channel_axis: the index of the channel axis. Returns: An up-down flipped image if condition is met, otherwise original image. @@ -812,7 +814,12 @@ def random_flip_up_down( # DO NOT REMOVE - Logging usage. should_transform = jax.random.bernoulli(key=key, p=probability) - return jax.lax.cond(should_transform, flip_up_down, lambda x: x, image) + return jax.lax.cond( + should_transform, + lambda x: flip_up_down(x, channel_axis=channel_axis), + lambda x: x, + image, + ) def random_brightness( diff --git a/dm_pix/_src/augment_test.py b/dm_pix/_src/augment_test.py index 8532e94..83e4e1b 100644 --- a/dm_pix/_src/augment_test.py +++ b/dm_pix/_src/augment_test.py @@ -188,6 +188,12 @@ def test_flip(self, images_list): reference_fn=None, probability=(0., 1.)) + def test_random_flip_up_down_channel_axis(self): + image = jnp.arange(2 * 3 * 4, dtype=jnp.float32).reshape((3, 2, 4)) + result = augment.random_flip_up_down( + jax.random.PRNGKey(0), image, probability=1., channel_axis=0) + np.testing.assert_array_equal(result, jnp.flip(image, axis=1)) + # Due to a bug in scipy we cannot test all available modes, refer to these # issues for more information: https://github.com/jax-ml/jax/issues/11097, # https://github.com/jax-ml/jax/issues/11097