pythonnumpyscipy.ndimage

How to rotate numpy 3D array around its center given a rotation matrix?


I have a 3D numpy array with shape=(50, 50, 50) which I would like to rotate around center, using a rotation matrix R.

I am trying to do this by using scipy.ndimage.affine_transform, but I'm open to better suggestions.

This is the code that shows the problem:

import random
import numpy as np
import scipy
from PIL import Image

R = np.array([
    [ 0.299297976, -0.817653322, -0.491796468],
    [-0.425077904, -0.575710904,  0.698473858],
    [-0.854242060,            0, -0.519875469],
])
print('Is R really orthogonal:', np.allclose(np.dot(R, R.T), np.identity(3)))  # True

# initialize some random data:
a = (50 + np.random.rand(50, 50, 50)*200).astype(np.uint8)
print(a.shape)  # (50, 50, 50)
Image.fromarray(a[:, :, 25]).show()

b = scipy.ndimage.affine_transform(a, R, offset=0)
Image.fromarray(b[:, :, 25]).show()

c = scipy.ndimage.affine_transform(a, R, offset=(50, 50, 50))
Image.fromarray(c[:, :, 25]).show()

original

rotated with offset 0

rotated with offset 50

What is unexpected to me is that:

What am I missing?

EDIT: the answer by @Nin17 is correct, this is the version that is adapted to original (3D) problem:

import random
import numpy as np
import scipy
from PIL import Image

R = np.array([
    [ 0.299297976, -0.817653322, -0.491796468, 0.],
    [-0.425077904, -0.575710904,  0.698473858, 0.],
    [-0.854242060,           0., -0.519875469, 0.],
    [          0.,           0.,           0., 1.],
])

N = 50

shift = np.array(
    [
        [1, 0, 0, N / 2],
        [0, 1, 0, N / 2],
        [0, 0, 1, N / 2],
        [0, 0, 0,     1],
    ]
)
unshift = np.array(
    [
        [1, 0, 0, -N / 2],
        [0, 1, 0, -N / 2],
        [0, 0, 1, -N / 2],
        [0, 0, 0,      1],
    ]
)

print('Is R orthogonal:', np.allclose(np.dot(R, R.T), np.identity(4)))  # True

a = (50 + np.random.rand(N, N, N)*200).astype(np.uint8)
print(a.shape)  # (50, 50, 50)
Image.fromarray(a[:, :, N//2]).show()

b = scipy.ndimage.affine_transform(a, shift @ R @ unshift, offset=0)
Image.fromarray(b[:, :, N//2]).show()

enter image description here

enter image description here


Solution

  • In the documentation for affine_transform, it states:

    Given an output image pixel index vector o, the pixel value is determined from the input image at position np.dot(matrix, o) + offset.

    Therefore, a pure rotation matrix will rotate the array around the pixel at index (0, 0) (in 2d). In order to rotate around the centre of the image, you could apply translations before and after the rotation to shift the image coordinates such that the centre is at (0, 0) perform the rotation, and then shift the image back, in 2d this could look like this:

    import numpy as np
    from scipy.ndimage import affine_transform
    import matplotlib.pyplot as plt
    
    
    rng = np.random.default_rng()
    
    theta = rng.random() * np.pi * 2
    rotation2d = np.array(
        [
            [np.cos(theta), -np.sin(theta), 0.0],
            [np.sin(theta), np.cos(theta), 0.0],
            [0.0, 0.0, 1.0],
        ]
    )
    
    N, M = 50, 51
    
    image = rng.random((N, M))
    
    shift = np.array(
        [
            [1, 0, N / 2],
            [0, 1, M / 2],
            [0, 0, 1], # need to also apply shift on third axes for 3d array
        ]
    )
    unshift = np.array(
        [
            [1, 0, -N / 2],
            [0, 1, -M / 2],
            [0, 0, 1], # need to also apply shift on third axes for 3d array
        ]
    )
    
    plt.figure()
    plt.imshow(image)
    
    
    plt.figure()
    plt.imshow(affine_transform(image, rotation2d))
    
    plt.figure()
    plt.imshow(affine_transform(image, shift @ rotation2d @ unshift))
    

    Output:

    enter image description here

    enter image description here

    enter image description here