Can the output of plt.imshow() be converted to a numpy array?-CodePudding

My intention is to use matplotlib to convert a coloured image into a grayscale image and use colormap to display it in the Viridis scale.

The code for that is as follows:

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.image as mpimg

IMG = mpimg.imread('dog_1.jpg')

def rgb2gray(rgb):
    return np.dot(rgb[...,:3], [0.2989, 0.5870, 0.1140])

   
gray = rgb2gray(IMG)    
plt.imshow(gray, cmap='viridis')
plt.show()

The output displayed is proper and as follows: Output Image

Now, I want to save the output image in a variable as a numpy array to carry out further processing. Can I do it in any way?

CodePudding user response：

plt.imread() returns a 3-dimensional numpy array with RGB layers. Your rgb2gray() function returns a 2-dimensional numpy array with a grayscale image. There is no need to extract a numpy array from the object returned by plt.imshow() when you have two numpy arrays with the image data already. However, if you insist in doing it, you can try the following:

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.image as mpimg

IMG = mpimg.imread('img.jpg')

def rgb2gray(rgb):
    return np.dot(rgb[...,:3], [0.2989, 0.5870, 0.1140])


gray = rgb2gray(IMG)    
aximg = plt.imshow(gray, cmap='viridis')

# an array with RGBA data of the image produced by plt.imshow 
arr = aximg.make_image(renderer=None, unsampled=True)[0]