使用本地图像通过 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 具有以下参数，包括转换：（参见 here更多信息）

torchvision.datasets.ImageFolder(root: str, transform: Optional[Callable] = None, target_transform: Optional[Callable] = None, loader: Callable[[str], Any] = , is_valid_file: Optional[Callable[[str], bool]] = None)

解决方法：在实例化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.DataLoader 但 torch.utils.data.DataLoader.

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

Getting image path through a torchvision dataloader using local images

pytorch

pytorch-dataloader