1

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

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])

4

1 回答 1

1

使用datasets.ImageFolder将使 PyTorch 独立处理每个“乐队”图像，并将文件夹名称（例如img1，img2...）视为“类标签”。
为了将 5 个图像文件加载为同一图像的不同波段/通道，您需要编写自己的自定义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)

请注意，您可能还需要重新实现增强。

于 2021-06-08T13:47:23.340 回答