bug with batch_size > 1 ?

[etienne87]

i am testing your code like this:

import torch
import numpy as np

from keymorph.keypoint_aligners import RigidKeypointAligner

bs, num_points = 8, 10
points_m = torch.rand(bs, num_points, 3)
points_f = torch.rand(bs, num_points, 3)

keypoint_aligner = RigidKeypointAligner(
                    points_m=points_m,
                    points_f=points_f,
                    dim=3
                )

snakes) eperot@frbucaw06dl:~/registration/monai-registration> python test_rigid_kp_aligner.py 
Traceback (most recent call last):
  File "/home/eperot/registration/monai-registration/test_rigid_kp_aligner.py", line 9, in <module>
    keypoint_aligner = RigidKeypointAligner(
                       ^^^^^^^^^^^^^^^^^^^^^
  File "/home/eperot/.conda/envs/snakes/lib/python3.11/site-packages/keymorph/keypoint_aligners.py", line 68, in __init__
    inverse_transform_matrix = self._square(
                               ^^^^^^^^^^^^^
  File "/home/eperot/.conda/envs/snakes/lib/python3.11/site-packages/keymorph/transformations.py", line 34, in _square
    square[:, : self.dim, : self.dim + 1] = matrix
    ~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: The expanded size of the tensor (1) must match the existing size (8) at non-singleton dimension 0.  Target sizes: [1, 3, 4].  Tensor sizes: [8, 3, 4]

the "_square" function inside "AffineTransform" is wrong i think

# this code assumes batch size of 1 
 def _square(self, matrix):
        square = torch.eye(self.dim + 1)[None]
        square[:, : self.dim, : self.dim + 1] = matrix
        return square

you can replace with

def _square(self, matrix):
    batch_size = matrix.shape[0]
    square = torch.eye(self.dim + 1).unsqueeze(0).repeat(batch_size, 1, 1)
    square[:, : self.dim, : self.dim + 1] = matrix
    return square

and later if i call affine_grid it also fails.

is the code hardcoded for a batch_size of 1 ? if yes, isn't this a problem a training time ?

Anyway, a simple fix can be made using F.affine_grid and respect your convention (matrix is obtained from ijk points)

def zyx_to_xyz_affine(affine_zyx):
    # Permutation matrix: ZYX -> XYZ (reverse the axes)
    P = torch.Tensor([
        [0, 0, 1, 0],
        [0, 1, 0, 0],
        [1, 0, 0, 0],
        [0, 0, 0, 1]
    ]).to(affine_zyx)

    # Convert: P @ affine_zyx @ P
    affine_xyz = P @ affine_zyx @ P
    return affine_xyz

 def get_flow_field(self, grid_shape):
        matrix = zyx_to_xyz_affine(self.inverse_transform_matrix)
        grid = torch.nn.functional.affine_grid(matrix[:, :3], grid_shape, align_corners=True)
        return grid

EDIT: i found another simpler, less intrusive fix

def get_inverse_transformed_points(self, points):
        """Transforms a set of points in fixed space to moving space using the fitted matrix.
        If align_in_real_world_coords is False, computes:
            p_m = A p_f.
        If align_in_real_world_coords is True, points must be in voxel coordinates, computes:
            p_m = A_m^-1 A^-1 A_f p_f.
        """
        batch_size, num_points, _ = points.shape
        transform_matrix = self.inverse_transform_matrix[:, :-1, :]

        # Transform real-world coordinates to the moving image space using the registration affine
        # Convert to homogeneous coordinates
        ones = torch.ones(batch_size, num_points, 1).to(points.device)
        points = torch.cat([points, ones], dim=2)
        # this is wrong, does not work for batch_size > 1
        # points = torch.bmm(transform_matrix, points.permute(0, 2, 1)).permute(0, 2, 1)
        points = torch.einsum('brc,bpc->bpr', transform_matrix, points)

        return points

[LukaslWeber]

@etienne87 Hey Etienne, I'm working on resolving the same issue. Were you able to fix it?

I also found that the affine_grid method is also hard-coded for a dimension of 1 if i am not mistaken.

def affine_grid(self, grid_shape):
        """Create a grid of affine transformed coordinates with specified shape to be used in F.grid_sample().
        Args:
            grid_shape (bs, 1, H, W) or (bs, 1, H, W, D): Shape of the grid to create
            m_shape: (bs, 1, H, W) or (bs, 1, H, W, D): Resolution of moving image being interpolating
        Returns:
            transformed_grid (bs, 1, H, W, D, dim): Affine grid in [-1, 1] space
        """
        grid = uniform_norm_grid(grid_shape, dim=self.dim)
        # Flatten the grid to (N x dim) array of voxel coordinates
        grid_flat = (
            grid.reshape(-1, self.dim).unsqueeze(0).to(self.transform_matrix)
        )  # Add batch dimension
        moving_voxel_coords = self.get_inverse_transformed_points(grid_flat)
        transformed_grid = moving_voxel_coords.reshape(1, *grid_shape[2:], self.dim)
        return transformed_grid

By swapping out the line
transformed_grid = moving_voxel_coords.reshape(1, *grid_shape[2:], self.dim)
with
batch_size = grid_shape[0] transformed_grid = moving_voxel_coords.reshape(batch_size, *grid_shape[2:], self.dim)
the AffineTransform should be correct now, right?

[etienne87]

Hello @LukaslWeber i did a P.R here.

What i did:

• swap the lines as you say
• modify torch.bmm for torch.einsum in "get_inverse_transformed_points" (it does not work otherwise because the grid is prepared with batch_size = 1) (an alternative is to repeat the grid batch_size times) (another alternative is to prepare the grid directly with torch.nn.functional.affine_grid)

Can you test this version for your problem and tell me if it fixes your usecase? I can add unit testing also.

[LukaslWeber]

Hey @etienne87, thanks for your quick answer! I just had a look at the P.R. and implemented the changes in the AffineTransform and AffineKeypointAligner.
The only difference that I made compared to you is in AffineTransform.affine_grid where I have
batch_size = grid_shape[0] transformed_grid = moving_voxel_coords.reshape(batch_size, *grid_shape[2:], self.dim)
and your version is
transformed_grid = moving_voxel_coords.reshape(-1, *grid_shape[2:], self.dim)
but this difference should be negligible.

For me, the original test_affine_2 doesn't run anymore because it produces a wrong alignment matrix with the AffineKeypointAligner. Does it work for you?
It gives me a result[0] of
tensor([[ 0., 1., 0., 0.], [-1., 0., 0., 0.], [ 0., 0., 0., 0.], [ 0., 0., 0., 1.]])
and true[0] of
tensor([[ 6.1232e-17, 1.0000e+00, 0.0000e+00, 0.0000e+00], [-1.0000e+00, 6.1232e-17, 0.0000e+00, 0.0000e+00], [ 0.0000e+00, 0.0000e+00, 1.0000e+00, 0.0000e+00], [ 0.0000e+00, 0.0000e+00, 0.0000e+00, 1.0000e+00]])
i.e. result[0][2][2] != true[0][2][2]

I also tried out the get_flow_field method and found that it doesn't work:
def test_grid_correctness_B1_and_B8(): """ Tests if our aligner'sget_flow_fieldoutput is identical to PyTorch's built-inF.affine_grid`. This is the true test
of the grid-generation logic.
"""
print("\n--- Running test_grid_correctness (B=1 & B=8) ---")
DIM = 2
H, W = 224, 224 # Use a small square image

# --- Test B=1 ---
print("Testing grid correctness with B=1...")
B1 = 1
grid_shape_b1 = (B1, 1, H, W)

# 1. Get our aligner's identity grid
matrix_b1 = torch.eye(DIM + 1, device=DEVICE).unsqueeze(0)
aligner_b1 = AffineTransform(inverse_matrix=matrix_b1, dim=DIM)
our_grid_b1 = aligner_b1.get_flow_field(grid_shape_b1)

# 2. Get PyTorch's "ground truth" identity grid
#    F.affine_grid expects the *forward* transform (not inverse)
#    and only the [B, 2, 3] part.
true_theta_b1 = torch.tensor([[1., 0., 0.], [0., 1., 0.]], device=DEVICE).unsqueeze(0)
true_grid_b1 = F.affine_grid(true_theta_b1, grid_shape_b1, align_corners=False)

# 3. Assert they are identical
assert_close(our_grid_b1, true_grid_b1, atol=1e-6, msg="Grid B=1")

When align_corners=False in thealign_img` method, the pixel coordinates are normalized so that the centers of the edge pixels map to -1 and +1, not the edges themselves.

As a result, I had to change the uniform_norm_grid.

`

def uniform_norm_grid(grid_shape, dim=3, dtype=torch.float32, device="cpu"):

  if dim == 2:
      # x = torch.linspace(-1, 1, grid_shape[2], dtype=dtype, device=device)
      # y = torch.linspace(-1, 1, grid_shape[3], dtype=dtype, device=device)
      N, H, W = grid_shape[0], grid_shape[2], grid_shape[3]
      y = (torch.linspace(0.5, H - 0.5, H, device=device, dtype=dtype) * 2 / H) - 1
      x = (torch.linspace(0.5, W - 0.5, W, device=device, dtype=dtype) * 2 / W) - 1
      grid = torch.meshgrid(x, y, indexing="ij")

  else:
      # x = torch.linspace(-1, 1, grid_shape[2], dtype=dtype, device=device)
      # y = torch.linspace(-1, 1, grid_shape[3], dtype=dtype, device=device)
      # z = torch.linspace(-1, 1, grid_shape[4], dtype=dtype, device=device)
      N, D, H, W = grid_shape[0], grid_shape[2], grid_shape[3], grid_shape[4]
      z = (torch.linspace(0.5, D - 0.5, D, device=device, dtype=dtype) * 2 / D) - 1
      y = (torch.linspace(0.5, H - 0.5, H, device=device, dtype=dtype) * 2 / H) - 1
      x = (torch.linspace(0.5, W - 0.5, W, device=device, dtype=dtype) * 2 / W) - 1
      grid = torch.meshgrid(x, y, z, indexing="ij")

  grid = torch.stack(grid, dim=-1).float()
  return grid

`

With this change, it runs.

[etienne87]

you are right the test_affine_2 does not pass, but for me it gives me error:

FAILED test/test.py::TestAffineAligner::test_affine_2 - torch._C._LinAlgError: linalg.inv: (Batch element 0): The diagonal element 3 is zero, the inversion could not be completed because the input matrix is singular.

i will investigate and fix this.

[LukaslWeber]

Ah maybe that could be solved by the torch.linalg.pinv?

When using the alignment module in training, I realized that RigidKeypointAligner.fit also needs some adaptions, because if we have a batch size > 1, our weights will also have a batch dimension and thus it needs adaptions. Sorry for being so cryptic but I'm trying to solve this issue too, currently! :)

if w is not None:
   w_exp = w.unsqueeze(1) # <- this is the fix
   # find weighed mean column wise
   p1_c = torch.sum(p1 * w_exp, axis=2, keepdims=True)
   p2_c = torch.sum(p2 * w_exp, axis=2, keepdims=True)
else:
   p1_c = torch.mean(p1, axis=2, keepdim=True)
   p2_c = torch.mean(p2, axis=2, keepdim=True)

...

#dets = torch.unsqueeze(dets, -1)
#dets = torch.stack([torch.ones_like(dets), torch.ones_like(dets), dets], axis=1)
#dets = torch.cat([dets, dets, dets], axis=2)
#V = V * torch.sign(dets)
signs = torch.where(dets < 0, -torch.ones_like(dets), torch.ones_like(dets)) # [B]
V_except_last_col = V[..., :, :-1]
V_last_col_flipped = V[..., :, -1] * signs.unsqueeze(-1)
V_corr = torch.cat([V_except_last_col, V_last_col_flipped.unsqueeze(-1)], dim=-1)

[etienne87]

Yes, i was testing the original keymorph, not my branch!! in my branch that uses pinv the code runs it is the assert_close that does not pass because i get

(Pdb) result
tensor([[[ 0., 1., 0., 0.],
[-1., 0., 0., 0.],
[ 0., 0., 0., 0.],
[ 0., 0., 0., 1.]]])

instead of :
(Pdb) true
tensor([[[ 6.1232e-17, 1.0000e+00, 0.0000e+00, 0.0000e+00],
[-1.0000e+00, 6.1232e-17, 0.0000e+00, 0.0000e+00],
[ 0.0000e+00, 0.0000e+00, 1.0000e+00, 0.0000e+00],
[ 0.0000e+00, 0.0000e+00, 0.0000e+00, 1.0000e+00]]])

the pinv get this 0 instead of 1 in the affine matrix, but i suspect this to be underconstrained no?

[LukaslWeber]

Oh great, this is the exact error that I also got! Yup, I also believe that it's underconstrained, as the z-dimension is always 0.
I just tested this with the following code, which passed!

def test_affine_2_improved():
  input1 = torch.tensor([[1, 0, 0], [0, -1, 0], [-1, 0, 0], [0, 1, 0], [0, 0, 1]], device=DEVICE).float()
    input1 = input1.view(1, 5, 3) # B=1, N=5

    input2 = torch.tensor([[0, -1, 0], [-1, 0, 0], [0, 1, 0], [1, 0, 0], [0, 0, 1]], device=DEVICE).float()
    input2 = input2.view(1, 5, 3) # B=1, N=5

    r = -np.pi / 2
    true_aug = torch.tensor(
        [
            [np.cos(r), -np.sin(r), 0, 0],
            [np.sin(r), np.cos(r), 0, 0],
            [0, 0, 1, 0],
        ],
        device=DEVICE
    ).float()
    true = torch.eye(4, 4, device=DEVICE).unsqueeze(0)
    true[:, :3, :] = true_aug

    affine_aligner = AffineKeypointAligner(points_m=input2, points_f=input1, dim=3)
    result = affine_aligner.inverse_transform_matrix
    assert_close(result, true)

[etienne87]

ok! yes i will modify the P.R code, thank you.