Visualize image batch from DataLoader

Question

I'm building a classifier for MIT Indoor Scenes dataset (https://web.mit.edu/torralba/www/indoor.html). Downloading the data gives me an 'Images' folder with 67 subfolders (for 67 classes), and different number of images within each subfolder. So far I have

import albumentations as A
from albumentations.pytorch import ToTensorV2
from torchvision.datasets import ImageFolder
from torch.utils.data import Dataset, DataLoader

alb_transform = A.Compose([
    A.Resize(256, 256),
    A.RandomCrop(width=224, height=224),
    A.HorizontalFlip(p=0.5),
    A.RandomBrightnessContrast(p=0.2),
    ToTensorV2()
])

dataset = ImageFolder('Images',transform=alb_transform)
dataloader = DataLoader(dataset, batch_size=64, shuffle=True, num_workers=0)

However, when I do

images, labels = next(iter(dataloader))

I get the error

KeyError: 'You have to pass data to augmentations as named arguments, for example: aug(image=image)'

The solutions I have looked at either used Pytorch Transforms, or wrote their own Dataset class. Is there a way to build the helper function for visualizing images using ImageLoader and Albumentations? Why am I getting that error?

Visualize image batch from DataLoader

0 Answers0