Pytorch transform.ToTensor() changes image

user_12422410 0
Viewed 19664

I want to convert images to tensor using torchvision.transforms.ToTensor(), after processing I printed the image but the image became so weird. Here is my code:

trans = transforms.Compose([
    transforms.ToTensor()])

demo = Image.open(img) 
demo_img = trans(demo)
demo_array = demo_img.numpy()*255
print(Image.fromarray(demo_array.astype(np.uint8)))

The original image is this

But after processing it is showed like this

Did I write something wrong or miss something?

data-science pytorch python-imaging-library python
1 Answers

It seems that the problem is with the channel axis.

If you look at torchvision.transforms docs, especially on ToTensor()

Converts a PIL Image or numpy.ndarray (H x W x C) in the range [0, 255] to a torch.FloatTensor of shape (C x H x W) in the range [0.0, 1.0]

So once you perform the transformation and return to numpy.array your shape is: (C, H, W) and you should change the positions, you can do the following:

demo_array = np.moveaxis(demo_img.numpy()*255, 0, -1)

This will transform the array to shape (H, W, C) and then when you return to PIL and show it will be the same image.

So in total:

import numpy as np
from PIL import Image
from torchvision import transforms

trans = transforms.Compose([transforms.ToTensor()])

demo = Image.open(img) 
demo_img = trans(demo)
demo_array = np.moveaxis(demo_img.numpy()*255, 0, -1)
print(Image.fromarray(demo_array.astype(np.uint8)))
user_88906040
Related