How to check dimensions of a numpy array?

Question

I want to do something like this inside a function:

if image.shape == 2 dimensions
    return image # this image is grayscale
else if image.shape = 3 dimensions
    return image # image is either RGB or YCbCr colorspace

Here, the image is a numpy array. I am not able to define that checking condition. I am really stuck at this point. Can anyone please help?

Answer 1

numpy.array.shape is a tuple of array dimensions. You can compute the length of a tuple and that will give the number of dimensions.

if len(image.shape) == 2:
    return image # this image is grayscale
elif len(image.shape) == 3:
    return image # image is either RGB or YCbCr colorspace

Numpy arrays also have a ndim attribute.

if image.ndim == 2:
    return image # this image is grayscale
elif image.ndim == 3:
    return image # image is either RGB or YCbCr colorspace