Search code examples
pythonnumpyarray-broadcasting

Element-wise multiplication of a 3D array with a 2D array


I have a portion of a RGB image as numpy array, the shape of which is (height, width, channel) = (5, 5, 3).

What I want to do with this is to get the sum of element-wise multiplication with 5x5 kernel matrix, channel by channel. So it should yield a vector of size 3.

My current solution is:

print(portion.shape)  # (5, 5, 3)
print(kernel.shape)  # (5, 5)
result = [(kernel * portion[:, :, channel]).sum() for channel in range(3)]
print(result.shape)  # (3,)

How can I achieve the same result in a more efficient way, hopefully without for-loop?


Solution

  • I'll show here two methods of doing this. The first one is basically the "manual" version that relies on broadcasting, which is an important concept to understand for using numpy and similar libraries.

    The second method is basically using the Einstein summation convention, which can be quite fast if used right.

    import numpy as np
    portion = np.zeros((5, 5, 3))
    kernel = np.zeros((5, 5))
    # alternative
    result = np.sum(kernel[..., None] * portion, axis=(0,1))
    print(result.shape)
    # einsum method:
    result = np.einsum('ij,ijk->k', kernel, portion)
    print(result.shape)
    
    

    Try it online!