From ffedd1b0f57be9716f22175fc414128c6367fadc Mon Sep 17 00:00:00 2001 From: Anas Date: Fri, 29 May 2026 10:48:19 -0400 Subject: [PATCH 1/2] fix: respect channel_axis in pad_to_size for CHW images pad_to_size hardcoded pad_width as ((top, bottom), (left, right), (0, 0)) which assumes channels-last (HWC) layout regardless of the channel_axis argument. For channels-first (CHW) images, this causes the padding to be applied along the channel dimension and leaves the width dimension unpadded, producing a wrong output shape (e.g. (5, 6, 4) instead of (3, 6, 6)). Fix: branch on _channels_last(image, channel_axis) and produce ((0, 0), (top, bottom), (left, right)) for CHW layout, consistent with how every other function in this module handles channel_axis. Add pad_to_size_test.py covering: - HWC baseline (channel_axis=-1): shape is preserved correctly. - CHW shape (channel_axis=0): must be (C, H', W') not (C+pad, H', W). - CHW values (channel_axis=0): original pixels centered, border is zero. - Batched CHW (channel_axis=1): must be (B, C, H', W'). --- dm_pix/_src/augment.py | 5 ++- dm_pix/_src/pad_to_size_test.py | 68 +++++++++++++++++++++++++++++++++ 2 files changed, 72 insertions(+), 1 deletion(-) create mode 100644 dm_pix/_src/pad_to_size_test.py diff --git a/dm_pix/_src/augment.py b/dm_pix/_src/augment.py index 0de0cb4..4a50e8a 100644 --- a/dm_pix/_src/augment.py +++ b/dm_pix/_src/augment.py @@ -358,7 +358,10 @@ def pad_to_size( top = delta_height // 2 bottom = max(target_height - (top + height), 0) - pad_width = ((top, bottom), (left, right), (0, 0)) + if _channels_last(image, channel_axis): + pad_width = ((top, bottom), (left, right), (0, 0)) + else: + pad_width = ((0, 0), (top, bottom), (left, right)) if batch: pad_width = ((0, 0), *pad_width) diff --git a/dm_pix/_src/pad_to_size_test.py b/dm_pix/_src/pad_to_size_test.py new file mode 100644 index 0000000..33da122 --- /dev/null +++ b/dm_pix/_src/pad_to_size_test.py @@ -0,0 +1,68 @@ +# Copyright 2020 DeepMind Technologies Limited. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tests for pad_to_size with channels-first (CHW) layout.""" + +from absl.testing import absltest +from dm_pix._src import augment +import jax.numpy as jnp +import numpy as np + + +class PadToSizeChannelsFirstTest(absltest.TestCase): + """Regression tests for pad_to_size with channel_axis != -1.""" + + def test_pad_to_size_hwc_shape(self): + """Channels-last (HWC) baseline: output shape must be (H', W', C).""" + image = jnp.zeros((4, 4, 3)) # HWC + result = augment.pad_to_size(image, target_height=6, target_width=6, + channel_axis=-1) + self.assertEqual(result.shape, (6, 6, 3)) + + def test_pad_to_size_chw_shape(self): + """Channels-first (CHW) with channel_axis=0: output shape must be (C, H', W'). + + Bug: pad_to_size hardcodes pad_width as ((top,bot), (left,right), (0,0)), + which treats the channel axis as the last dimension regardless of + channel_axis. For CHW input (C, H, W), this pads along C and H but not W, + producing a wrong shape. + """ + image = jnp.zeros((3, 4, 4)) # CHW: C=3, H=4, W=4 + result = augment.pad_to_size(image, target_height=6, target_width=6, + channel_axis=0) + self.assertEqual(result.shape, (3, 6, 6)) + + def test_pad_to_size_chw_values(self): + """CHW padding must center the original image with zeros around it.""" + # 1x1 single-pixel image, CHW with 1 channel + image = jnp.ones((1, 1, 1)) # C=1, H=1, W=1 + result = augment.pad_to_size(image, target_height=3, target_width=3, + channel_axis=0) + self.assertEqual(result.shape, (1, 3, 3)) + # The center pixel should be 1; surrounding pixels should be 0. + np.testing.assert_array_equal(result[0, 1, 1], 1.0) + np.testing.assert_array_equal(result[0, 0, :], 0.0) + np.testing.assert_array_equal(result[0, 2, :], 0.0) + np.testing.assert_array_equal(result[0, :, 0], 0.0) + np.testing.assert_array_equal(result[0, :, 2], 0.0) + + def test_pad_to_size_batch_chw_shape(self): + """Batched CHW (B, C, H, W): output shape must be (B, C, H', W').""" + image = jnp.zeros((2, 3, 4, 4)) # B=2, C=3, H=4, W=4 + result = augment.pad_to_size(image, target_height=6, target_width=6, + channel_axis=1) + self.assertEqual(result.shape, (2, 3, 6, 6)) + + +if __name__ == '__main__': + absltest.main() From 5ab629f7b9fb4c06718179b7b02598609f62eb58 Mon Sep 17 00:00:00 2001 From: Anas Date: Mon, 1 Jun 2026 09:51:32 -0400 Subject: [PATCH 2/2] test: move pad_to_size CHW tests into augment_test.py Address review feedback on PR #112: - Move the channels-first pad_to_size regression tests out of the standalone pad_to_size_test.py and into the existing TestCustom class in augment_test.py, matching that file's conventions. - Drop the explanatory comments, consistent with the surrounding tests. Delete dm_pix/_src/pad_to_size_test.py. --- dm_pix/_src/augment_test.py | 33 ++++++++++++++++ dm_pix/_src/pad_to_size_test.py | 68 --------------------------------- 2 files changed, 33 insertions(+), 68 deletions(-) delete mode 100644 dm_pix/_src/pad_to_size_test.py diff --git a/dm_pix/_src/augment_test.py b/dm_pix/_src/augment_test.py index f584a17..8532e94 100644 --- a/dm_pix/_src/augment_test.py +++ b/dm_pix/_src/augment_test.py @@ -559,6 +559,39 @@ def test_pad_to_size_when_target_size_smaller_than_original( self.assertEqual(output.shape[1], expected_height) self.assertEqual(output.shape[2], expected_width) + def test_pad_to_size_hwc_shape(self): + image = jnp.zeros((4, 4, 3)) + result = augment.pad_to_size( + image, target_height=6, target_width=6, channel_axis=-1 + ) + self.assertEqual(result.shape, (6, 6, 3)) + + def test_pad_to_size_chw_shape(self): + image = jnp.zeros((3, 4, 4)) + result = augment.pad_to_size( + image, target_height=6, target_width=6, channel_axis=0 + ) + self.assertEqual(result.shape, (3, 6, 6)) + + def test_pad_to_size_chw_values(self): + image = jnp.ones((1, 1, 1)) + result = augment.pad_to_size( + image, target_height=3, target_width=3, channel_axis=0 + ) + self.assertEqual(result.shape, (1, 3, 3)) + np.testing.assert_array_equal(result[0, 1, 1], 1.0) + np.testing.assert_array_equal(result[0, 0, :], 0.0) + np.testing.assert_array_equal(result[0, 2, :], 0.0) + np.testing.assert_array_equal(result[0, :, 0], 0.0) + np.testing.assert_array_equal(result[0, :, 2], 0.0) + + def test_pad_to_size_batch_chw_shape(self): + image = jnp.zeros((2, 3, 4, 4)) + result = augment.pad_to_size( + image, target_height=6, target_width=6, channel_axis=1 + ) + self.assertEqual(result.shape, (2, 3, 6, 6)) + if __name__ == "__main__": jax.config.update("jax_default_matmul_precision", "float32") diff --git a/dm_pix/_src/pad_to_size_test.py b/dm_pix/_src/pad_to_size_test.py deleted file mode 100644 index 33da122..0000000 --- a/dm_pix/_src/pad_to_size_test.py +++ /dev/null @@ -1,68 +0,0 @@ -# Copyright 2020 DeepMind Technologies Limited. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Tests for pad_to_size with channels-first (CHW) layout.""" - -from absl.testing import absltest -from dm_pix._src import augment -import jax.numpy as jnp -import numpy as np - - -class PadToSizeChannelsFirstTest(absltest.TestCase): - """Regression tests for pad_to_size with channel_axis != -1.""" - - def test_pad_to_size_hwc_shape(self): - """Channels-last (HWC) baseline: output shape must be (H', W', C).""" - image = jnp.zeros((4, 4, 3)) # HWC - result = augment.pad_to_size(image, target_height=6, target_width=6, - channel_axis=-1) - self.assertEqual(result.shape, (6, 6, 3)) - - def test_pad_to_size_chw_shape(self): - """Channels-first (CHW) with channel_axis=0: output shape must be (C, H', W'). - - Bug: pad_to_size hardcodes pad_width as ((top,bot), (left,right), (0,0)), - which treats the channel axis as the last dimension regardless of - channel_axis. For CHW input (C, H, W), this pads along C and H but not W, - producing a wrong shape. - """ - image = jnp.zeros((3, 4, 4)) # CHW: C=3, H=4, W=4 - result = augment.pad_to_size(image, target_height=6, target_width=6, - channel_axis=0) - self.assertEqual(result.shape, (3, 6, 6)) - - def test_pad_to_size_chw_values(self): - """CHW padding must center the original image with zeros around it.""" - # 1x1 single-pixel image, CHW with 1 channel - image = jnp.ones((1, 1, 1)) # C=1, H=1, W=1 - result = augment.pad_to_size(image, target_height=3, target_width=3, - channel_axis=0) - self.assertEqual(result.shape, (1, 3, 3)) - # The center pixel should be 1; surrounding pixels should be 0. - np.testing.assert_array_equal(result[0, 1, 1], 1.0) - np.testing.assert_array_equal(result[0, 0, :], 0.0) - np.testing.assert_array_equal(result[0, 2, :], 0.0) - np.testing.assert_array_equal(result[0, :, 0], 0.0) - np.testing.assert_array_equal(result[0, :, 2], 0.0) - - def test_pad_to_size_batch_chw_shape(self): - """Batched CHW (B, C, H, W): output shape must be (B, C, H', W').""" - image = jnp.zeros((2, 3, 4, 4)) # B=2, C=3, H=4, W=4 - result = augment.pad_to_size(image, target_height=6, target_width=6, - channel_axis=1) - self.assertEqual(result.shape, (2, 3, 6, 6)) - - -if __name__ == '__main__': - absltest.main()