Question

我是pytorch的新手。我正在尝试为图像数据集创建一个DataLoader，其中每个图像都有一个对应的地面真实性（相同名称）：

root:
--->RGB:
------>img1.png
------>img2.png
------>...
------>imgN.png
--->GT:
------>img1.png
------>img2.png
------>...
------>imgN.png

当我将根文件夹（包含RGB和GT文件夹）的路径用作torchvision.datasets.ImageFolder的输入时，它将读取所有图像，就好像它们都是要用于输入的（分类为RGB和GT）一样，似乎无法配对RGB-GT图像。我想将RGB-GT图像配对，随机播放，然后将其分为定义大小的批次。怎么做到呢？任何建议将被认真考虑。谢谢。

Answer 1

我认为，一个好的起点是使用VisionDataset类作为基础。我们将在这里使用的是：DatasetFolder source code。因此，我们将创建类似的东西。您会注意到该类依赖于datasets.folder模块中的另外两个函数：default_loader和make_dataset。

我们不会修改default_loader，因为它已经可以了，它只是帮助我们加载图像，因此我们将其导入。

但是我们需要一个新的make_dataset函数，该函数从根文件夹准备了正确的图像对。由于原始make_dataset对图像（如果更精确地说是图像路径）和它们的根文件夹配对为目标类（类索引），因此我们有(path, class_to_idx[target])对的列表，但我们需要(rgb_path, gt_path) 。这是新的make_dataset的代码：

def make_dataset(root: str) -> list:
    """Reads a directory with data.
    Returns a dataset as a list of tuples of paired image paths: (rgb_path, gt_path)
    """
    dataset = []

    # Our dir names
    rgb_dir = 'RGB'
    gt_dir = 'GT'   

    # Get all the filenames from RGB folder
    rgb_fnames = sorted(os.listdir(os.path.join(root, rgb_dir)))

    # Compare file names from GT folder to file names from RGB:
    for gt_fname in sorted(os.listdir(os.path.join(root, gt_dir))):

            if gt_fname in rgb_fnames:
                # if we have a match - create pair of full path to the corresponding images
                rgb_path = os.path.join(root, rgb_dir, gt_fname)
                gt_path = os.path.join(root, gt_dir, gt_fname)

                item = (rgb_path, gt_path)
                # append to the list dataset
                dataset.append(item)
            else:
                continue

    return dataset

我们现在有什么？让我们将我们的功能与原始功能进行比较：

from torchvision.datasets.folder import make_dataset as make_dataset_original


dataset_original = make_dataset_original(root, {'RGB': 0, 'GT': 1}, extensions='png')
dataset = make_dataset(root)

print('Original make_dataset:')
print(*dataset_original, sep='\n')

print('Our make_dataset:')
print(*dataset, sep='\n')

Original make_dataset:
('./data/GT/img1.png', 1)
('./data/GT/img2.png', 1)
...
('./data/RGB/img1.png', 0)
('./data/RGB/img2.png', 0)
...
Our make_dataset:
('./data/RGB/img1.png', './data/GT/img1.png')
('./data/RGB/img2.png', './data/GT/img2.png')
...

我认为效果很好）现在该创建类Dataset了。这里最重要的部分是__getitem__方法，因为它导入图像，应用变换并返回张量，供数据加载器使用。我们需要读取一对图像（rgb和gt）并返回2张量图像的元组：

from torchvision.datasets.folder import default_loader
from torchvision.datasets.vision import VisionDataset


class CustomVisionDataset(VisionDataset):

    def __init__(self,
                 root,
                 loader=default_loader,
                 rgb_transform=None,
                 gt_transform=None):
        super().__init__(root,
                         transform=rgb_transform,
                         target_transform=gt_transform)

        # Prepare dataset
        samples = make_dataset(self.root)

        self.loader = loader
        self.samples = samples
        # list of RGB images
        self.rgb_samples = [s[1] for s in samples]
        # list of GT images
        self.gt_samples = [s[1] for s in samples]

    def __getitem__(self, index):
        """Returns a data sample from our dataset.
        """
        # getting our paths to images
        rgb_path, gt_path = self.samples[index]

        # import each image using loader (by default it's PIL)
        rgb_sample = self.loader(rgb_path)
        gt_sample = self.loader(gt_path)

        # here goes tranforms if needed
        # maybe we need different tranforms for each type of image
        if self.transform is not None:
            rgb_sample = self.transform(rgb_sample)
        if self.target_transform is not None:
            gt_sample = self.target_transform(gt_sample)      

        # now we return the right imported pair of images (tensors)
        return rgb_sample, gt_sample

    def __len__(self):
        return len(self.samples)

让我们测试一下：

from torch.utils.data import DataLoader

from torchvision.transforms import ToTensor
import matplotlib.pyplot as plt


bs=4  # batch size
transforms = ToTensor()  # we need this to convert PIL images to Tensor
shuffle = True

dataset = CustomVisionDataset('./data', rgb_transform=transforms, gt_transform=transforms)
dataloader = DataLoader(dataset, batch_size=bs, shuffle=shuffle)

for i, (rgb, gt) in enumerate(dataloader):
    print(f'batch {i+1}:')
    # some plots
    for i in range(bs):
        plt.figure(figsize=(10, 5))
        plt.subplot(221)
        plt.imshow(rgb[i].squeeze().permute(1, 2, 0))
        plt.title(f'RGB img{i+1}')
        plt.subplot(222)
        plt.imshow(gt[i].squeeze().permute(1, 2, 0))
        plt.title(f'GT img{i+1}')
        plt.show()

出局：

batch 1:

...

Here，您可以找到带有代码和简单虚拟数据集的笔记本。

适用于Image GT数据集的Pytorch Dataloader

1 个答案: