Adding data into new dimension of numpy array-CodePudding

I have a numpy array of size 55 x 10 x 10, which represents 55 10 x 10 grayscale images. I'm trying to make them RGB by duplicating the 10 x 10 images 3 times.

From what I've understood, I first need to add a new dimension to house the duplicated data. I've done this using:

array_4d = np.expand_dims(array_3d, 1),

so I now have a 55 x 1 x 10 x 10 array. How do I now duplicate the 10 x 10 images and add them back into this array?

Quick edit: In the end I want a 55 x 3 x 10 x 10 array

CodePudding user response：

Let us first create a 3d array of size 55x10x10

from matplotlib          import pyplot as plt
import numpy as np
original_array = np.random.randint(10,255, (55,10,10))
print(original_array.shape)
>>>(55, 10, 10)

Visual of first image in array:

first_img = original_array[0,:,:]
print(first_img.shape)
plt.imshow(first_img, cmap='gray')
>>>(10, 10)

Now you can get your desired array in just one single step.

stacked_img = np.stack(3*(original_array,), axis=1)
print(stacked_img.shape)
>>>(55, 3, 10, 10)

Use axis=-1 if you want channel last

Now let us verify that the value are correct by extracting the first image from this array and taking average of 3 channels:

new_img = stacked_img[0,:,:,:]
print(new_img.shape)
>>> (3, 10, 10)

new_img_mean = new_img.mean(axis=0)
print(new_img_mean.shape)
>>> (10, 10)

np.allclose(new_img_mean, first_img) # If this is True then the two arrays are same
>>> True

For visual verification, you'll have to move the channel to last because that is what matplotlib needs. This is a 3 channel image, so we are not using cmap='gray' here

print(np.moveaxis(new_img, 0, -1).shape)
plt.imshow(np.moveaxis(new_img, 0, -1))
>>> (10, 10, 3)