Pytorch從一個輸入目錄中加載所有的PNG圖像，並將它們存儲在張量中

import os
import imageio
from imageio import imread
import torch

# batch_size = 3
# batch = torch.zeros(batch_size, 3, 256, 256, dtype=torch.uint8)
# batch.shape  #torch.Size([3, 3, 256, 256])[B,C,H,W]

data_dir = 'C:/Users/Dell/Pictures/'
filenames = [name for name in os.listdir(data_dir)
            if os.path.splitext(name)[-1] == '.png'] #選擇指定目錄下的.png圖片

for i, filename in enumerate(filenames):
    img_arr = imageio.imread(os.path.join(data_dir, filename)) #imread讀入為H*W*C
    img_t = torch.from_numpy(img_arr)
    img_t = img_t.permute(2, 0, 1) #交換維度
    img_t = img_t[:3]  #只保留前3個通道
    batch[i] = img_t   #指第i個維度上的所有數據