How to create a dataset containing one datapoint and its augmentations?

Viewed 38

I have an imagenet dataset. I want to extract only one image of the dataset and apply random augmentations on this single data point. So basically the transform I define is as following:

sample_transform = transforms.Compose([
        transforms.RandomResizedCrop(img_size, scale=(0.2, 1.0)),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize(mean=_IMAGENET_RGB_MEANS, std=_IMAGENET_RGB_STDS),
    ])

And assume that my one sample is as follows:

from torch.utils.data import Subset
def extract_sample(data, idx):
    return Subset(data, [idx])
single_sample = extract_sample(data, idx = 0)

Not I want to create a dataset by applying sample_transform, 64 times on single_sample (in order to get 64 different augmentations of single_sample). How can I do that?

Thanks in advance for your help!

1 Answers

I'd utilize torch.utils.data.RandomSampler with replacement:

import torch
from torchvision import transforms
from torchvision.datasets import ImageFolder
from torch.utils.data import Subset, DataLoader, RandomSampler

import matplotlib.pyplot as plt
from PIL import Image

sample_transform = transforms.Compose([
        transforms.RandomResizedCrop(224, scale=(0.2, 1.0)),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize(0.5, 0.5),
    ])

data = ImageFolder('/media/storage/datasets/t', transform=sample_transform)

def extract_sample(data, idx):
    return Subset(data, [idx])

single_sample = extract_sample(data, idx = 0)
sampler = RandomSampler(single_sample, num_samples=64, replacement=True)
loader = DataLoader(single_sample, sampler=sampler)

for img, target in loader:
    plt.imshow(img.squeeze().permute(1, 2, 0))
    plt.show()
Related