Я хочу построить некоторые изображения, которые иногда имеют оттенки серого, иногда цветные.
Так как я использую pytorch, преобразованные изображения могут быть (50, 100, 1) или (50, 100, 3) в зависимости от того, являются ли они цвет или нет. Поскольку matplotlib не может взять прежний для изображения, мне нужно вернуть следующие формы:
(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.