Search code examples
pythonnumpymatrixtransformation

Optimize this linear transformation for images with Numpy


Good evening,

I'm trying to learn NumPy and have written a simple Linear transformation that applies to an image using for loops:

import numpy as np

M = np.array([
    [width, 0],
    [0, height]
])

T = np.array([
    [1, 3],
    [0, 1]
])

def transform_image(M, T):
    T_rel_M = abs(M @ T)
    new_img = np.zeros(T_rel_M.sum(axis=1).astype("int")).T
    
    for i in range(0, 440):
        for j in range(0, 440):
            x = np.array([j, i])
            coords = (T @ x)
            x = coords[0]
            y = coords[1]
            new_img[y, -x] = image[i, -j]
    
    return new_img

plt.imshow(transform_image(M, T))

It does what I want and spits out a transformation that is correct, except that I think there is a way to do this without the loops.

I tried doing some stuff with meshgrid but I couldn't figure out how to get the pixels from the image in the same way I do it in the loop (using i and j). I think I figured out how to apply the transformation but then getting the pixels from the image in the correct spots wouldn't work.

Any ideas?

EDIT: Great help with below solutions, lezaf's solution was very similar to what I tried before, the only step missing that I couldn't figure out was assigning the pixels from the old to the new image. I made some changes to the code to exclude transposing, and also added a astype("int") so it works with float values in the T matrix:

def transform_image(M, T):
    T_rel_M = abs(M @ T)
    new_img = np.zeros(T_rel_M.sum(axis=1).astype("int")).T
    x_combs = np.array(np.meshgrid(np.arange(width), np.arange(height))).reshape(2,-1)
    coords = (T @ x_combs).astype("int")
    new_img[coords[1, :], -coords[0, :]] = image[x_combs[1, :], -x_combs[0, :]]
    
    return new_img

Solution

  • A more efficient solution is the following:

    def transform_image(M, T):
        T_rel_M = abs(M @ T)
        new_img = np.zeros(T_rel_M.sum(axis=1).astype("int")).T
    
        # This one replaces the double for-loop
        x_combs = np.array(np.meshgrid(np.arange(440), np.arange(440))).T.reshape(-1,2)
        # Calculate the new coordinates
        coords = (T@x_combs.T)
        # Apply changes to new_img
        new_img[coords[1, :], -coords[0, :]] = image[x_combs[:, 1], -x_combs[:,0]]
    

    I updated my solution removing the for-loop, so now is a lot more straightforward.

    After this change, the time of the optimized code is 50 ms compared to the initial 3.06 s of the code in question.