Display Pytorch tensor as image using Matplotlib

Question

I am trying to display an image stored as a pytorch tensor.

trainset = datasets.ImageFolder('data/Cat_Dog_data/train/', transform=transforms)
trainload = torch.utils.data.DataLoader(trainset, batch_size=32, shuffle=True)

images, labels = iter(trainload).next()
image = images[0]
image.shape 

>>> torch.Size([3, 224, 224]) # pyplot doesn't like this, so reshape

image = image.reshape(224,224,3)
plt.imshow(image.numpy())

This method is displaying a 3 by 3 grid of the same image, always in greyscale. For example:

How do I fix this so that the single color image is displayed correctly?

I have the exact same problem. The solution works (thanks Nicolas!) but I would to understand why I get 9 images when I use reshape or view — Jan Pisl
– Jan Pisl, Commented Apr 6, 2022 at 15:06
@Jan Pisl I actually have that question open this very second for my own investigation: stackoverflow.com/questions/51143206/… — rocksNwaves
– rocksNwaves, Commented Apr 6, 2022 at 17:19

Nicolas Gervais · Accepted Answer · 2020-06-23 18:15:49Z

16

That's very odd. Try putting the channels last by permuting rather than reshaping:

image.permute(1, 2, 0)

answered Jun 23, 2020 at 18:15

Nicolas Gervais

36.9k23 gold badges123 silver badges160 bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

Stefano Giannini Over a year ago

This method can be used also for list of frames (video). Just insert 4 arguments

Collectives™ on Stack Overflow

Display Pytorch tensor as image using Matplotlib

1 Answer 1

1 Comment

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

1 Comment

Your Answer

Sign up or log in

Post as a guest

Linked

Related