Вывод изображения из массива numpy в colab

Есть .tif изображения считанные и записанные в массив numpy:

train_images = []
for directory_path in glob.glob('/content/gdrive/My Drive/Image_dataset/sandstone/images_train_128_patches_stone'):
    for img_path in glob.glob(os.path.join(directory_path, "*.tif")):
        #Считывание изображения в (RGB-1)
        img = cv2.imread(img_path, 1)       
        img = cv2.resize(img, (SIZE_Y, SIZE_X))
        #img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
        train_images.append(img)

train_images = np.array(train_images)
print(train_images.shape) #(1600, 128, 128, 3)

train_masks = [] 
for directory_path in glob.glob('/content/gdrive/My Drive/Image_dataset/sandstone/masks_train_128_patches_stone'):
    for mask_path in glob.glob(os.path.join(directory_path, "*.tif")):
        #Считывание изображения в оттенках серого
        mask = cv2.imread(mask_path, 0)       
        mask = cv2.resize(mask, (SIZE_Y, SIZE_X), interpolation = cv2.INTER_NEAREST)  #Otherwise ground truth changes due to interpolation
        #img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
        train_masks.append(mask)

train_masks = np.array(train_masks)
print(train_masks.shape) # (1600, 128, 128)

Далее кодирую метки для классов изображений масок:

from sklearn.preprocessing import LabelEncoder
labelencoder = LabelEncoder()
# Возвращает размерности массивов (уровней массивов)
number_images, height, width = train_masks.shape
print("Размерности массивов: ", train_masks.shape)
print("Class values in the dataset are ... ", np.unique(train_masks))
# Изменение представления массива (перестройка уровней)
train_masks_reshaped = train_masks.reshape(-1,1)
print("Размерности массивов: ", train_masks_reshaped.shape)
#Нормализует метки таким образом, чтобы они содержали только значения от 0 до n_classes-1
#Tакже можно использовать для преобразования нечисловых меток (если они хешируемы и сопоставимы) в числовые метки
train_masks_reshaped_encoded = labelencoder.fit_transform(train_masks_reshaped)
train_masks_encoded_original_shape = train_masks_reshaped_encoded.reshape(number_images, height, width)
print("Class values in the dataset are ... ", np.unique(train_masks_encoded_original_shape))

Добавляю в массив с масками измерение с кодом класса маски:

train_masks_input = np.expand_dims(train_masks_encoded_original_shape, axis=3)
print(train_masks_input.shape) # (1600, 128, 128, 1)

Разбиваю на наборы:

from sklearn.model_selection import train_test_split
X1, X_test, y1, y_test = train_test_split(train_images, train_masks_input, test_size = 0.10, random_state = 0)
X_train, X_do_not_use, y_train, y_do_not_use = train_test_split(X1, y1, test_size = 0.5, random_state = 0)

print(X_train.shape) # (720, 128, 128, 3)
print(y_train.shape) # (720, 128, 128, 1)

И пытаюсь просмотреть соответствия между изображением и маской:

import random
import numpy as np
image_number = random.randint(0, len(X_train))
plt.figure(figsize=(12, 6))
plt.subplot(121)
plt.imshow(np.reshape(X_train[image_number], (128, 128)), cmap='BGR')
#plt.imshow(X_train[image_number])
plt.subplot(122)
plt.imshow(np.reshape(y_train[image_number], (128, 128)), cmap='gray')
#plt.imshow(y_train[image_number])
plt.show()

Вот эта строка отрабатывает нормально и выводит изображение:

plt.imshow(np.reshape(y_train[image_number], (128, 128)), cmap='gray')

В этой строке ошибка:

plt.imshow(np.reshape(X_train[image_number], (128, 128)), cmap='BGR')

ValueError: cannot reshape array of size 49152 into shape (128,128)

Подскажите в чем причина? Как это можно исправить?

Мне нужно после присвоения меток классов и разбития убедится, что маска соответствует изображению и вывести это на экран.

Судя по числу в ошибке, 49152=128x128x3