ToPILImage() does not work for batches>1 and dimension more than 3D, is it a bug?

vision
1

Im using the ToPILImage() to conver an image from tensor to numpy and possibly show it or save it etc
I noticed that I need to do squeeze or manually make the input to be 3D, other wise the ToPILImage does not works.

trainloader = torch.utils.data.DataLoader(TrainSet, batch_size=1,
                                          shuffle=True, num_workers=2)

dataiter = iter(trainloader)
images, labels = dataiter.next()

TensorToNumpy = ToPILImage()

img = TensorToNumpy(images)
plt.imshow(img)

wornt work and i have to do:

img = TensorToNumpy(images[0,:,:,:])
plt.imshow(img)

Can you please fix this bug?

2

As a workaround you could create a list with your images:

x = torch.randn(10, 3, 24, 24)
images = [transforms.ToPILImage()(x_) for x_ in x]
2 Likes