帶有圖像的 numpy vstack

Question

我正在編寫一個函數來從圖像中讀取像素數據並將它們存儲在一個 numpy 數組中以進一步進行訓練/測試拆分。

當我運行這段代碼時，它拋出一個異常，說除了連接軸之外的所有輸入數組維度都必須完全匹配。

我不確定為什么會發生這個問題以及如何解決它。

from PIL import Image
import numpy as np
import os

X = np.array([])
y = []

categories = {
    'A': 1,
    'B': 2
}

root = data_dir + '/cropped_resized(128,128)/'

for path, subdirs, files in os.walk(root):
    for name in files:
        img_path = os.path.join(path,name)
        category = categories[os.path.basename(path)]
        im = Image.open(img_path)
        img_pixels = list(im.getdata())
        width, height = im.size
        X = np.vstack((X, img_pixels))
        #X = np.concatenate((X, img_pixels), axis=0)
        y.append(category)

X_train, X_test, y_train, y_test = train_test_split(X, y)

這是一個失敗的圖片示例

Answer 1

確定您希望圖像為 RGB 還是灰度，並確保它們在加載時如此。

具體來說，改變這一行：

im = Image.open(img_path)

到

im = Image.open(img_path).convert('RGB')

或者

im = Image.open(img_path).convert('L')