Pytorch 以小批量加载数据

Question

我有一个这样的图片文件夹

Images
|
|__img1
|   |__img1_b01.tiff
|   |__img1_b02.tiff
|   |__img1_b03.tiff
|   |__img1_b04.tiff
|   |__img1_b05.tiff
|  
|__img2
|   |__img2_b02.tiff
|   |__img2_b02.tiff
|   |__img2_b03.tiff
|   |__img2_b04.tiff
|   |__img2_b05.tiff
|
|.. img1000

每个文件夹代表一张图片。文件夹中的每个文件代表图像的一个波段通道。

因此每张图片都会有一个

我一直在编写 pytorch 自定义数据加载器以分批加载 64 个

所以我可以拥有特征批量形状：torch.Size([64,5, 256, 256])

我试过下面的代码

from torchvision import datasets, transforms
from torch.utils import data

dataset = datasets.ImageFolder(root = Images/, 
            transform = transforms.ToTensor())

loader = data.DataLoader(dataset, batch_size = 64, shuffle = True)

但它没有给出我想要的结果 Feature batch shape: torch.Size([64, 5, 256, 256])

Answer 1

使用 datasets.ImageFolder 将使 PyTorch 独立处理每个“波段”图像并将文件夹名称（例如，img1、img2...）视为“class 标签".
为了加载相同图像的不同 bands/channels 的 5 个图像文件，您需要编写自己的 .

这个习俗 Dataset 可能看起来像这样：

import torch
import os
from PIL import Image
import numpy as np

class MultiBandDataset(torch.utils.data.Dataset):
  def __init__(self, root, num_bands):
    self.root = root
    self.num_bands = num_bands
    self.imgs = os.listdir(root)  # all `imgNN` folders
     
  def __len__(self):
    return len(self.imgs)  # number of images = number of subfolders

  def __getitem__(self, index):
    multi_band = []
    # get the subfolder
    subf = os.path.join(self.root, self.imgs[index])
    for band in range(self.num_bands):
      b = Image.open(os.path.join(subf, f'{self.imgs[index]}_b{band+1:02d}.tiff')).convert("F")  # make sure you are reading a single channel from each image. you need to verify this part.
      multi_band.append(np.array(b).astype(np.float32)[None,...])  # add singleton channel dimension
    return np.concatenate(numti_band, axis=0)

请注意，您可能还需要重新实现扩充。

Pytorch 以小批量加载数据

Pytorch load data in mini batches

python

image-processing

pytorch

satellite-image

pytorch-dataloader