Suppose that I have an image of three channels, and its corresponding numpy array is of the shape (H,W,3). Now I create a function which does something to a 2D numpy array. How can I apply the funtion to each channel of the image without using a for-loop since it is slow?
I'm looking for a function similar to apply in pandas. I've tried np.apply_over_axes and np.apply_along_axis, but both results are not correct.
In the code below, I hope to add a gaussian noise seperately to each channel of the image.
test_img = np.array(train_data[0][0])
test_img.shape
# (198, 180, 3)
def add_gaussian_noise(img_, amplitude=1.0, mean=0.0, variance=1.0):
img_ = img_.astype('float32')
# print(img_.shape)
# assertion failed
# assert img_.shape == (test_img.shape[0], test_img.shape[1])
img_ = img_ + amplitude * np.random.normal(mean, variance, img_.shape)
img_ = np.clip(img_, 0.0, 255.0)
return img_
test_img = np.apply_along_axis(arr= test_img, func1d= add_gaussian_noise, axis = 2)
print(test_img)
I'm anticipating that img_.shape == (test_img.shape[0], test_img.shape[1])
, but it turns out that img_.shape == (3,)
in my code.
In this case, the optimization won't make a big difference, but I am looking for a more generalized solution so that in other cases, e.g. I have to apply the same function to hundreds of grayscale pictures, the code works more efficiently than the for-loop.