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...）视为“类标签”。
为了将 5 个图像文件加载为相同图像的不同波段/通道，您需要编写自己的 custom Dataset。

此自定义 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 以小批量加载数据

1 个答案: