Generating pytorch's theta from affine transform matrix

user_9897331 0
Viewed 1470

I'm trying to perform a rigid + scale transformation on a 3D volume with pytorch, but I can't seem to understand how the theta required for torch.nn.functional.affine_grid works.

I have a transformation matrix of size (1,4,4) generated by multiplying the matrices Translation * Scale * Rotation. If I use this matrix in, for example, scipy.ndimage.affine_transform, it works with no issues. However, the same matrix (cropped to size (1,3,4)) fails completely with torch.nn.functional.affine_grid.

I have managed to understand how the translation works (range -1 to 1) and I have confirmed that the Translation matrix works by simply normalizing the values to that range. As for the other two, I am lost.

I tried using a basic Scaling matrix alone (below) as a most basic comparison but the results in pytorch are different than that of scipy

Scaling = 
[[0.75, 0, 0, 0],
[[0, 0.75, 0, 0],
[[0, 0, 0.75, 0],
[[0, 0, 0, 1]]

How can I convert the (1,4,4) affine matrix to work the same with torch.nn.functional.affine_grid? Alternatively, is there a way to generate the correct matrix based on the transformation parameters (shift, euler angles, scaling)?

affinetransform pytorch python
1 Answers

To anyone that comes across a similar issue in the future, the problem with scipy vs pytorch affine transforms is that scipy applies the transforms around (0, 0, 0) while pytorch applies it around the middle of the image/volume.

For example, let's take the parameters:

euler_angles = [ea0, ea1, ea2]
translation = [tr0, tr1, tr2]
scale = [sc0, sc1, sc2]

and create the following transformation matrices:

# Rotation matrix
R_x(ea0, ea1, ea2) = np.array([[1, 0, 0, 0],
                              [0, math.cos(ea0), -math.sin(ea0), 0],
                              [0, math.sin(ea0), math.cos(ea0), 0],
                              [0, 0, 0, 1]])

R_y(ea0, ea1, ea2) = np.array([[math.cos(ea1), 0, math.sin(ea1), 0],
                               [0, 1, 0, 0],
                               [-math.sin(ea1), 0, math.cos(ea1)], 0],
                               [0, 0, 0, 1]])

R_z(ea0, ea1, ea2) = np.array([[math.cos(ea2), -math.sin(ea2), 0, 0],
                               [math.sin(ea2), math.cos(ea2), 0, 0],
                               [0, 0, 1, 0],
                               [0, 0, 0, 1]])

R = R_x.dot(R_y).dot(R_z)

# Translation matrix
T(tr0, tr1, tr2) = np.array([[1, 0, 0, -tr0],
                             [0, 1, 0, -tr1],
                             [0, 0, 1, -tr2],
                             [0, 0, 0, 1]])
# Scaling matrix
S(sc0, sc1, sc2) = np.array([[1/sc0, 0, 0, 0],
                             [0, 1/sc1, 0, 0],
                             [0, 0, 1/sc2, 0],
                             [0, 0, 0, 1]])

If you have a volume of size (100, 100, 100), the scipy transform around the centre of the volume requires moving the centre of the volume to (0, 0, 0) first, and then moving it back to (50, 50, 50) after S, T, and R have been applied. Defining:

T_zero = np.array([[1, 0, 0, 50],
                   [0, 1, 0, 50],
                   [0, 0, 1, 50],
                   [0, 0, 0, 1]])

T_centre = np.array([[1, 0, 0, -50],
                     [0, 1, 0, -50],
                     [0, 0, 1, -50],
                     [0, 0, 0, 1]])

The scipy transform around the centre is then:

transform_scipy_centre = T_zero.dot(T).dot(S).dot(R).T_centre

In pytorch, there are some slight differences to the parameters.

The translation is defined between -1 and 1. Their order is also different. Using the same (100, 100, 100) volume as an example, the translation parameters in pytorch are given by:

# Note the order difference
translation_pytorch = =[tr0_p, tr1_p, tr2_p] = [tr0/50, tr2/50, tr1/50] 
T_p = T(tr0_p, tr1_p, tr2_p)

The scale parameters are in a different order: scale_pytorch = [sc0_p, sc1_p, sc2_p] = [sc2, sc0, sc1]
S_p = S(sc0_p, sc1_p, sc2_p)

The euler angles are the biggest difference. To get the equivalent transform, first the parameters are negative and in a different order:

# Note the order difference
euler_angles_pytorch = [ea0_p, ea1_p, ea2_p] = [-ea0, -ea2, -ea1]
R_x_p = R_x(ea0_p, ea1_p, ea2_p)
R_y_p = R_y(ea0_p, ea1_p, ea2_p)
R_z_p = R_z(ea0_p, ea1_p, ea2_p)

The order in which the rotation matrix is calculated is also different: # Note the order difference R_p = R_x_p.dot(R_z_p).dot(R_y_p)

With all these considerations, the scipy transform with:

transform_scipy_centre = T_zero.dot(T).dot(S).dot(R).T_centre

is equivalent to the pytorch transform with:

transform_pytorch = T_p.dot(S_p).dot(R_p)

I hope this helps!

user_98973310
Related