pytorch中 PIL与tensor、numpy与tensor之间相互转化|简记

参考链接-Pytorch中Tensor与各种图像格式的相互转化

下面是整理的 cv、PIL 读取图片，然后PIL2tensor、Tensor2PILImage、tensor2numpy相互转化的代码，建议直接复制运行，观察输出 ：
torch 1.1.0 ，torchvision 0.3.0

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

if __name__ == '__main__':
    # 方法定义
    def denorm(x):
        return x * 0.5 + 0.5

    def tensor2numpy(x):
        return x.detach().cpu().numpy().transpose(1, 2, 0)

    def RGB2BGR(x):
        return cv2.cvtColor(x, cv2.COLOR_RGB2BGR)
    # 实例化
    test_transform = transforms.Compose([
        # transforms.Resize((256, 256)),
        transforms.ToTensor(),
        transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
    ])

    unloader = transforms.ToPILImage()

    print('-------------Start------------------')
    # cv read img
    img_path = 'default.jpg'
    cv2_image = cv2.imread(img_path)  # numpy.ndarray（H，W，C=3），通道顺序（B,G,R)
    print('type(cv2_image): ', type(cv2_image))
    print('cv2_image.size: ', cv2_image.shape)

    print()
    # PIL read img
    image = Image.open(img_path)
    image.save('PIL_default.jpg')
    print('type(image): ', type(image))
    print('PIL image.size: ', image.size)

    print()
    # PIL2tensor
    img_tensor = test_transform(image)  # (C，H, W), 通道顺序（R,G,B)
    print('PIL2tensor type: ', type(img_tensor))
    print('PIL2tensor shape: ', img_tensor.shape)

    print()
    # Tensor2PILImage， 使用 PIL 进行保存
    Tensor2PIL = unloader(img_tensor)
    Tensor2PIL.save('temp.jpg')
    print('Tensor2PIL type: ', type(Tensor2PIL))
    print('Tensor2PIL size: ', Tensor2PIL.size)

    print('---------------------------------------')
    # tensor2numpy.ndarray 使用cv2 来进行保存
    tensor2numpy_img = RGB2BGR(tensor2numpy(denorm(img_tensor)))
    cv2.imwrite(os.path.join('test_cv.jpg'), tensor2numpy_img)
    print('tensor2numpy_img type: ', type(tensor2numpy_img))
    print('tensor2numpy_img shape: ', tensor2numpy_img.shape)

    print()
    # 给 tensor 扩充一个 维度
    # pytorch中，处理图片需要一个batch一个batch的操作，需要准备的数据格式是 [batch_size, n_channels, hight, width]
    img4 = test_transform(image).unsqueeze(0)
    print('add a dimension: ', img4.shape)

  

 

  
1
  
2
  
3
  
4
  
5
  
6
  
7
  
8
  
9
  
10
  
11
  
12
  
13
  
14
  
15
  
16
  
17
  
18
  
19
  
20
  
21
  
22
  
23
  
24
  
25
  
26
  
27
  
28
  
29
  
30
  
31
  
32
  
33
  
34
  
35
  
36
  
37
  
38
  
39
  
40
  
41
  
42
  
43
  
44
  
45
  
46
  
47
  
48
  
49
  
50
  
51
  
52
  
53
  
54
  
55
  
56
  
57
  
58
  
59
  
60
  
61
  
62
  
63
  
64
 

输出如下：

-------------Start------------------
type(cv2_image):  <class 'numpy.ndarray'>
cv2_image.size:  (220, 178, 3)

type(image):  <class 'PIL.JpegImagePlugin.JpegImageFile'>
PIL image.size:  (178, 220)

PIL2tensor type:  <class 'torch.Tensor'>
PIL2tensor shape:  torch.Size([3, 220, 178])

Tensor2PIL type:  <class 'PIL.Image.Image'>
Tensor2PIL size:  (178, 220)
---------------------------------------
tensor2numpy_img type:  <class 'numpy.ndarray'>
tensor2numpy_img shape:  (220, 178, 3)

add a dimension:  torch.Size([1, 3, 220, 178])


  

 

  
1
  
2
  
3
  
4
  
5
  
6
  
7
  
8
  
9
  
10
  
11
  
12
  
13
  
14
  
15
  
16
  
17
  
18
 

知识点：

torch.unsqueeze() 类似于 np.expand_dims() 、tf.expand_dims(input, axis=-1)

tf.expand_dims(input, axis=-1)

pytorch中squeeze()和unsqueeze()函数介绍
torch.unsqueeze() 和 torch.squeeze()

文章来源: positive.blog.csdn.net，作者：墨理学AI，版权归原作者所有，如需转载，请联系作者。

原文链接：positive.blog.csdn.net/article/details/107944405