Save Pytorch 4D tensor as image

user_11087305 0
Viewed 4267

I have a 4-d Pytorch tensor that I would like to save to disk as a .jpg

My tensor is the following size:

print(image_tensor.size())

>>>torch.Size([1, 3, 400, 711])

I can view the entire tensor as one image within my IDE:

ax1.imshow(im_convert(image_tensor))

Since I am able to view the entire tensor as one image, I am assuming there is a way to also save it as such. However, when I try to save the image, it looks like it only saves the blue color channel. I would like to save the entire tensor as a single image.

img1 = image_tensor[0]
save_image(img1, 'img1.jpg')
tensor pytorch image python
1 Answers

In PyTorch this snippet is working and saving the image:

from torchvision.utils import save_image
import torch
import torchvision

tensor= torch.rand(2, 3, 400, 711) 

img1 = tensor[0]
save_image(img1, 'img1.png')

Before saving the image can you check the shape of the img1 in any case something happened.

user_59795430
Related