Return all dimensions if there are 3 channels, else return 0
I want to plot some images that are sometimes grayscale, sometimes color.

Since I'm using pytorch, the transformed images are either (50, 100, 1) or (50, 100, 3) depending on if they're color or not. Since matplotlib cannot take the former for a picture, I need to return these shapes:

(100, 100, 3) # if it's a color picture
(100, 100) # if it's a grayscale picture


Those are the shapes matplotlib can deal with. Here is what my workflow looks like:

import numpy as np
import matplotlib.pyplot as plt

h, w = 50, 100

grey = np.random.randint(0, 256, (h, w, 1))
color = np.random.randint(0, 256, (h, w, 3)


With the same line, I need to be able to plot images that can either be color or greyscale.

plt.imshow(grey[:, :, 0 if grey.shape[-1] == 1 else :]
plt.imshow(color[:, :, 0 if color.shape[-1] == 1 else :]


But that's not correct python syntax. 

Can't you just write

plt.imshow(grey[:,:,0] if grey.shape[-1] == 1 else grey)
plt.imshow(color[:,:,0] if color.shape[-1] == 1 else color)

