Repository navigation
Add channel_axis to rotate and affine_transform - #142
krakken190 wants to merge 3 commits into
Conversation
krakken190
changed the title
|
Hi @claudiofantacci , I've resolved the merge conflict with master. Ready for review/merge whenever you get a chance! |
|
@claudiofantacci Sorry, my merge-conflict resolution in the GitHub UI dropped the |
| @parameterized.product( | ||
| mode_and_order=[("nearest", 0), ("nearest", 1), ("mirror", 1), | ||
| ("constant", 1)], | ||
| channel_axis=(0, -3), |
There was a problem hiding this comment.
The use of -3 seems a bit unnecessary here. Can make it explicit?
| ("constant", 1)], | ||
| channel_axis=(0, -3), | ||
| ) | ||
| def test_affine_transform_channel_axis(self, mode_and_order, channel_axis): |
There was a problem hiding this comment.
Can we add also the different ways of passing the dimensions/offset? See how test_affine_transform is done, and we just add channel axis to the product of the parameterized test.
| mode_and_order=[("nearest", 0), ("nearest", 1), ("mirror", 1), | ||
| ("constant", 1)], | ||
| angle=(np.pi / 2, -np.pi / 4), | ||
| channel_axis=(0, -3), |
| angle=(np.pi / 2, -np.pi / 4), | ||
| channel_axis=(0, -3), | ||
| ) | ||
| def test_rotate_channel_axis(self, mode_and_order, angle, channel_axis): |
There was a problem hiding this comment.
Can we test it in the same manner we do test_rotate, we just need to add channel axis to that test probably. Scipy as reference function has axes of which oens you want to rotate.
Fixes #141.
Problem
rotateandaffine_transformsay they support HWC or CHW images, but they always treat the first two array axes as height and width. For CHW inputs this rotates or shifts across the channel axis and silently returns a corrupted image.Fix
channel_axis: int = -1toaffine_transform. When it is not-1, the channel axis is moved to the end before the transform and moved back afterwards. This is the same patternelastic_deformationuses.matrixandoffsetare always interpreted in HWC order, independent ofchannel_axis, so one transform can be applied to either layout. The docstring now says this.channel_axis: int = -1torotate. It computes the rotation centre from the HWC-ordered shape and forwardschannel_axistoaffine_transform.-1keeps the existing behaviour, so current callers are unaffected.Tests
Added to
TestCustom:test_affine_transform_channel_axis: fornearest/0,nearest/1,mirror/1 andconstant/1, withchannel_axis0 and -3, the CHW result equals the HWC result moved to channels-first.test_rotate_channel_axis: the same check for two angles, run underjax.jit.Validation
pytest dm_pix/_src/augment_test.py: 886 passedpytest dm_pix/_src/api_test.py: 80 passedchannel_axis=0, CHW output matches HWC exactly (max diff 0.0) for both functions. Without it, the difference was 0.99997 before and after the change.Tested on CPU with Python 3.12.2, JAX/jaxlib 0.12.0, NumPy 2.5.4, SciPy 1.18.1 and TensorFlow 2.21.0. GPU/TPU not tested.