CODEWITHSUNDEEP
pythonnumpy
Dzung Nguyen
Dzung Nguyen

Reputation: 3944

Polymorphic treating of 2D and 3D numpy array

I'm working with images in Python, and need to write a program that can handle both color and grayscale image. They're numpy array

Shape of Color: (512,512,3)

Shape of Grayscale:(512,512)

Now I have to loop through every channel of images, i.e. return :

For Color: im[:,:,0], im[:,:,1], im[:,:,2]

For Grayscale: im[:,:]

How to write them in the same format without any if condition? I tried im[:,:,0] for grayscale but it's out of range for the index.

Upvotes: 0

Views: 145

Answers (1)

mgilson
mgilson

Reputation: 310097

I'm not sure if this is helpful or not, but numpy provides the ability to insert new axes:

im_new = im_old[:,:,np.newaxis]

As I understand it, this makes it so that im_new[i,j,k] is the same as im_old[i,j] for any k.

(also note that np.newaxis is simply an alias for None)

Upvotes: 1

Related Questions