Display examples of augmented images in PyTorch
I want to display some samples of augmented training images.

My transform includes the standard ImageNet transforms.Normalize like this:

train_transforms = transforms.Compose([transforms.RandomRotation(30),
                                       transforms.RandomResizedCrop(224),
                                       transforms.RandomHorizontalFlip(),
                                       transforms.ToTensor(),
                                       transforms.Normalize([0.485, 0.456, 0.406],
                                                            [0.229, 0.224, 0.225])])


However, because of the Normalise, the images display in weird colours.

This answer says I'd need access to the original image, which is difficult when the transforms are applied at load time:

image_datasets['train'] = datasets.ImageFolder(train_dir, transform=train_transforms)


How would I go about displaying a few sample augmented images in their usual colours while using the normalised ones for calculation?

To answer my own question, I came up with the following:



# Undo transforms.Normalize
def denormalise(image):
    image = image.numpy().transpose(1, 2, 0)  # PIL images have channel last
    mean = [0.485, 0.456, 0.406]
    stdd = [0.229, 0.224, 0.225]
    image = (image * stdd + mean).clip(0, 1)
    return image


example_rows = 2
example_cols = 5

sampler = torch.utils.data.RandomSampler(image_datasets['train'],
                                         num_samples=example_rows * example_cols)

# Get a batch of images and labels  
images, indices = next(iter(sampler)) 

plt.rcParams['figure.dpi'] = 120  # Increase size of pyplot plots

# Show a grid of example images    
fig, axes = plt.subplots(example_rows, example_cols, figsize=(9, 5)) #  sharex=True, sharey=True)
axes = axes.flatten()
for ax, image, index in zip(axes, images, indices):
    ax.imshow(denormalise(image))
    ax.set_axis_off()
    ax.set_title(class_names[index], fontsize=7)

fig.subplots_adjust(wspace=0.02, hspace=0)
fig.suptitle('Augmented training set images', fontsize=20)
plt.show()


This is based on PyTorch's Transfer Learning Tutorial's code but displays the title above each image and generally looks much nicer.
