PIL图像转torch的tensor

1、图片转Tensor

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

pic_location = 'dataset/1.png'
img = Image.open(os.path.join(os.getcwd(), pic_location))

# 方法一
img_convert_to_numpy = np.array(img)  # (32, 32, 3)
img_convert_to_tensor1 = torch.tensor(img_convert_to_numpy.transpose(2, 0, 1) / 255)  # torch.Size([3, 32, 32])

# 方法二
transform = transforms.Compose([transforms.ToTensor()])
img_convert_to_tensor2 = transform(img_convert_to_numpy)  # torch.Size([3, 32, 32])

print(torch.equal(img_convert_to_tensor1, img_convert_to_tensor2))  # False
print(img_convert_to_tensor1.dtype)  # torch.float64
print(img_convert_to_tensor2.dtype)  # torch.float32

 

  首先，使用PIL库（Python Imaging Library）读取图片，得到的是PIL图像对象。通过numpy.array(img)或者numpy.asarray(img)，转化为uint8的数值数组形式，格式为 H x W x C ，数值范围在0-255之间。（主要区别在于当数据源是ndarray时，array仍然会copy出一个副本，占用新的内存，但asarray不会）
  由于在Pytorch中，图像的格式为 C x H x W（想象一下，就是卷积核要卷积图片的形式），所以需要用transpose进行转置。这篇文章使用的图片和我前一篇博文通过cifar-10数据集理解numpy数组的高（H）、宽（W）、通道数（C）中选取的图片一样，通过实例图片和代码解析能更好地帮助你理解。

2、细节

  tensor除法会使输出结果的精度高一级，可能会导致后面计算类型不匹配，如float32 / float32 = float64。在上面的代码中，torch.equal(img_convert_to_tensor1, img_convert_to_tensor2)是等于False的。Tensor默认的dtype是float32，所以当Tensor的类型为float32时，打印Tensor是不会显示的。
  所以，我们要进行这样的处理：img_convert_to_tensor1 = torch.tensor(img_convert_to_numpy.transpose(2, 0, 1) / 255, dtype=torch.float32)，结果就等于True了。

3、拓展：将Tensor转换为PIL

  有batch维度的Tensor一定要通过torch.squeeze(image,dim=0)降维，然后img = transforms.ToPILImage()(image)一步搞定。