使用本地图像通过 torchvision 数据加载器获取图像路径

Question

我想在我的脚本中使用数据加载器。

通常默认的函数调用是这样的。

dataset = ImageFolderWithPaths(
    data_dir,
    transforms.Compose([
            transforms.ColorJitter(0.1, 0.1, 0.1, 0.1),
            transforms.Resize((img_size_XY, img_size_XY)),
            transforms.ToTensor(),
            transforms.Normalize(_mean , _std)
        ])
)

dataloader = torch.utils.data.DataLoader(
    dataset,
    batch_size=batch_size,
    shuffle=False,
    num_workers=2
)

并迭代我使用的这个数据加载器

for inputs, labels , paths in _dataloader:
    break

现在我需要收集每个图像的路径。

我在github中找到了这段代码：（https://gist.github.com/andrewjong/6b02ff237533b3b2c554701fb53d5c4d）

class ImageFolderWithPaths(datasets.ImageFolder):
    """Custom dataset that includes image file paths. Extends
    torchvision.datasets.ImageFolder
    """

    # override the __getitem__ method. this is the method that dataloader calls
    def __getitem__(self, index):
        # this is what ImageFolder normally returns 
        original_tuple = super(ImageFolderWithPaths, self).__getitem__(index)
        # the image file path
        path = self.imgs[index][0]
        # make a new tuple that includes original and the path
        tuple_with_path = (original_tuple + (path,))
        return tuple_with_path

# EXAMPLE USAGE:
# instantiate the dataset and dataloader
data_dir = "your/data_dir/here"
dataset = ImageFolderWithPaths(data_dir) # our custom dataset
dataloader = torch.utils.DataLoader(dataset)

# iterate over data
for inputs, labels, paths in dataloader:
    # use the above variables freely
    print(inputs, labels, paths)

但这段代码没有考虑转换，就像我的原始代码一样。

有人能帮我解决这个问题吗？

Answer 1

由于ImageFolderWithPaths继承datasets.ImageFolder自 GitHub 的代码所示，并datasets.ImageFolder具有以下参数，包括转换：（有关更多信息，请参阅此处）

torchvision.datasets.ImageFolder（根：str，变换：可选[Callable] =无，target_transform：可选[Callable] =无，加载器：Callable[[str]，Any] =，is_valid_file：可选[Callable[[str]，布尔]] = 无）

解决方案：您可以在实例化时直接使用您的转换ImageFolderWithPaths。

import torch
from torchvision import datasets
from torch.utils.data import DataLoader 

class ImageFolderWithPaths(datasets.ImageFolder):

    def __getitem__(self, index):
  
        img, label = super(ImageFolderWithPaths, self).__getitem__(index)
        
        path = self.imgs[index][0]
        
        return (img, label ,path)


# put here your root directory not subfolders directory
# subfolders should be names of classes or encodings
root_dir = "training" 

transform = transforms.Compose([transforms.Resize((32, 32)), 
                                transforms.ToTensor()]) # my transformations.

dataset = ImageFolderWithPaths(root_dir,transform=transform) # add transformation directly

dataloader = DataLoader(dataset)


for inputs, labels, paths in dataloader:
    print(inputs.shape, labels, paths)

# output
torch.Size([1, 3, 32, 32]) tensor([0]) ('training\\0\\1.jpg',)
torch.Size([1, 3, 32, 32]) tensor([0]) ('training\\0\\1000.jpg',)
torch.Size([1, 3, 32, 32]) tensor([0]) ('training\\0\\10005.jpg',)
torch.Size([1, 3, 32, 32]) tensor([0]) ('training\\0\\10010.jpg',)

我还编辑了 github 上的代码，因为没有torch.utils.DataLoaderbut torch.utils.data.DataLoader。