Search code examples
pythonnumpymatplotlibimage-processingscipy

Rotate 2D numpy array about a specific line

I am trying to rotate an image about a specific line drawn within the image. Think of the image of a random 2D numpy array. I have two specific points in the array through which I would like to plot a line. I want to rotate the array/image in such a way that this line forms the new x-axis. How do I do it?

Some code explanation:

import numpy as np
import matplotlib.pyplot as plt
from scipy import stats

x = np.linspace(-5,5,8)
y = np.linspace(-25,25,7)

z = np.random.randn(7,8)*0.5
z = z/np.max(z)

z[5,6] = 2; z[1,1] = 2

plt.pcolormesh(x,y,z); plt.colorbar()

Output image

Now I want to draw a line between the the points with maximum z-values and rotate the array such that this line is the new x-axis.

I try this by first getting the end points of the line. These are usually centroids of large z-values and then get the slope of the line of best fit through these points.

# get endpoints which lie at the very end of the gaussian distribution
line_endpoints = np.where(z>z.mean()+3*z.std(), 1, 0)
yline, xline = np.nonzero(line_points)
fitline = stats.linregress(x=xline, y=yline)

Now, fitline gives me the slope and intercept of the best-fit line. How can I rotate the original array z using this slope such that the selected end-points form the new x-axis? I want to look at such an image. I checked out scipy docs but couldn't find anything suitable.

Solution

  • scipy.ndimage.rotate() can do that.

    Code example:

    points_above_3_std = z>z.mean()+3*z.std()
idx = np.array(np.where(points_above_3_std))
endpoint_diff = idx[:, 1] - idx[:, 0]
rot_angle = np.degrees(np.arctan2(endpoint_diff[0], endpoint_diff[1]))
print(rot_angle)
z_rot = scipy.ndimage.rotate(z, rot_angle)
plt.imshow(z_rot[::-1])  # Note: flip y axis to get same direction as pcolormesh plot

    Plot:

    rotated plot

    Notes:

    • I changed from describing the rotation with slope to describing the rotation with degrees. Slopes don't handle rotations by 90 degrees well (i.e. if it should be rotated by 90 degrees, the slope will be infinite.)

    • There are two rotations which satisfy your question. For example, here a 38 degree rotation puts the two points at the same Y coordinate. However, a rotation by -142 degrees, which would turn the graph upside down, also does this.

      If you wanted to make sure that the rotation angle is in the interval [-90, 90], you could add the following code before rotate().

      rot_angle = (rot_angle + 90) % 180 - 90