如果有 3 个通道，则返回所有维度，否则返回 0

Question

我想绘制一些有时是灰度、有时是彩色的图像。

由于我使用的是pytorch，因此转换后的图像是(50, 100, 1) 或(50, 100, 3)，具体取决于它们是否为彩色。由于matplotlib 无法将former 用于拍照，我需要返回这些形状：

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

这些是matplotlib 可以处理的形状。以下是我的工作流程：

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)

使用同一行，我需要能够绘制彩色或灰度图像。

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

但这不正确的python 语法。

Answer 1

你就不能写

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