Fix align_corners mismatch in AffineTransform

The to_norm_affine call was using hardcoded align_corners=False while
affine_grid and grid_sample used self.align_corners (default True).
This caused a half-pixel offset between coordinate systems.

- Change to_norm_affine to use self.align_corners for consistency
- Update test expected values to reflect correct behavior
- Add tests for align_corners consistency verification

Signed-off-by: Soumya Snigdha Kundu <soumya_snigdha.kundu@kcl.ac.uk>

Author: aymuos15

monai/networks/layers/spatial_transforms.py

@@ -566,7 +566,7 @@ def forward(
                 affine=theta,
                 src_size=src_size[2:],
                 dst_size=dst_size[2:],
-                align_corners=False,
+                align_corners=self.align_corners,
                 zero_centered=self.zero_centered,
             )
         if self.reverse_indexing:

tests/networks/layers/test_affine_transform.py

@@ -154,21 +154,21 @@ def test_zoom_1(self):
         affine = torch.as_tensor([[2.0, 0.0, 0.0], [0.0, 1.0, 0.0]])
         image = torch.arange(1.0, 13.0).view(1, 1, 3, 4).to(device=torch.device("cpu:0"))
         out = AffineTransform()(image, affine, (1, 4))
-        expected = [[[[2.333333, 3.333333, 4.333333, 5.333333]]]]
+        expected = [[[[5.0, 6.0, 7.0, 8.0]]]]
         np.testing.assert_allclose(out, expected, atol=_rtol)
 
     def test_zoom_2(self):
         affine = torch.as_tensor([[2.0, 0.0, 0.0], [0.0, 2.0, 0.0]], dtype=torch.float32)
         image = torch.arange(1.0, 13.0).view(1, 1, 3, 4).to(device=torch.device("cpu:0"))
         out = AffineTransform((1, 2))(image, affine)
-        expected = [[[[1.458333, 4.958333]]]]
+        expected = [[[[5.0, 7.0]]]]
         np.testing.assert_allclose(out, expected, atol=1e-5, rtol=_rtol)
 
     def test_zoom_zero_center(self):
         affine = torch.as_tensor([[2.0, 0.0, 0.0], [0.0, 2.0, 0.0]], dtype=torch.float32)
         image = torch.arange(1.0, 13.0).view(1, 1, 3, 4).to(device=torch.device("cpu:0"))
         out = AffineTransform((1, 2), zero_centered=True)(image, affine)
-        expected = [[[[5.5, 7.5]]]]
+        expected = [[[[5.0, 8.0]]]]
         np.testing.assert_allclose(out, expected, atol=1e-5, rtol=_rtol)
 
     def test_affine_transform_minimum(self):
@@ -380,6 +380,53 @@ def test_forward_3d(self):
         np.testing.assert_allclose(actual, expected)
         np.testing.assert_allclose(list(theta.shape), [1, 3, 4])
 
+    def test_align_corners_consistency(self):
+        """
+        Test that align_corners is consistently used between to_norm_affine and grid_sample.
+
+        With an identity affine transform, the output should match the input regardless of
+        the align_corners setting. This test verifies that the coordinate normalization
+        in to_norm_affine uses the same align_corners value as affine_grid/grid_sample.
+        """
+        # Create a simple test image
+        image = torch.arange(1.0, 13.0).view(1, 1, 3, 4)
+
+        # Identity affine in pixel space (i, j, k convention with reverse_indexing=True)
+        identity_affine = torch.eye(3).unsqueeze(0)
+
+        # Test with align_corners=True (the default)
+        xform_true = AffineTransform(align_corners=True)
+        out_true = xform_true(image, identity_affine)
+        np.testing.assert_allclose(out_true.numpy(), image.numpy(), atol=1e-5, rtol=_rtol)
+
+        # Test with align_corners=False
+        xform_false = AffineTransform(align_corners=False)
+        out_false = xform_false(image, identity_affine)
+        np.testing.assert_allclose(out_false.numpy(), image.numpy(), atol=1e-5, rtol=_rtol)
+
+    def test_align_corners_true_translation(self):
+        """
+        Test that translation works correctly with align_corners=True.
+
+        This ensures to_norm_affine correctly converts pixel-space translations
+        to normalized coordinates when align_corners=True.
+        """
+        # 4x4 image
+        image = torch.arange(1.0, 17.0).view(1, 1, 4, 4)
+
+        # Translate by +1 pixel in the j direction (column direction)
+        # With reverse_indexing=True (default), this is the last spatial dimension
+        # Positive translation in the affine shifts the sampling grid, resulting in
+        # the output appearing shifted in the opposite direction
+        affine = torch.tensor([[[1.0, 0.0, 0.0], [0.0, 1.0, 1.0], [0.0, 0.0, 1.0]]])
+
+        xform = AffineTransform(align_corners=True, padding_mode="zeros")
+        out = xform(image, affine)
+
+        # Expected: shift columns left by 1, rightmost column becomes 0
+        expected = torch.tensor([[[[2, 3, 4, 0], [6, 7, 8, 0], [10, 11, 12, 0], [14, 15, 16, 0]]]], dtype=torch.float32)
+        np.testing.assert_allclose(out.numpy(), expected.numpy(), atol=1e-4, rtol=_rtol)
+
 
 if __name__ == "__main__":
     unittest.main()