Tensor、Numpy、PIL格式转换以及图像显示

重点

cv.imshow() ：显示图像是 BGR格式的
plt.imshow() ：图像显示是 RGB格式的

Tensor ：存储的数据分布在 [0, 1]
Numpy ：存储的数据分布在 [0, 255]

CIFAR-10 数据集

数据集为 RGB格式的；
在使用 opencv 显示时需要先转换成 BGR格式；
在使用 plt显示时 无需 转换格式；

示例：

dict = unpickle('./dataset/cifar-10-batches-py/test_batch')

img = dict[b'data']
image = img[0]
image = image.reshape(3, 32, 32).transpose(1, 2, 0)
cv_show('image', image)

r, g, b = cv.split(image)
pic = cv.merge([b, g, r])
cv_show('pic', pic)
1
2
3
4
5
6
7
8
9
10

左侧图为原图，右侧图为失真图 ：

格式转换

import torch
import numpy as np

# convert numpy to tensor or vise versa
np_data = np.arange(6).reshape((2, 3))
torch_data = torch.from_numpy(np_data)  # ndarray转换为Tensor
tensor2array = torch_data.numpy()  # Tensor转换为ndarray

test = out.cpu().detach().numpy() # 从CUDA上的Tensor转换为ndarray
1
2
3
4
5
6
7
8
9

Tensor ==> Numpy

import torch
import torchvision
import pickle
import cv2 as cv

transform_tensor = torchvision.transforms.Compose([
    torchvision.transforms.ToTensor(),
])

transform_picture = torchvision.transforms.Compose([
    torchvision.transforms.ToTensor(),                   # 转换成Tensor格式
    torchvision.transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])


def cv_show(name, img):
    cv.imshow(name, img)
    cv.waitKey(0)
    cv.destroyAllWindows()


def unpickle(file):
    with open(file, 'rb') as fo:
        dict = pickle.load(fo, encoding='bytes')
    return dict


dict = unpickle('./dataset/cifar-10-batches-py/test_batch')

img = dict[b'data']
image = img[0]
image = image.reshape(3, 32, 32).transpose(1, 2, 0)
print(image)

image_Tensor = transform_tensor(image).unsqueeze(0)
print(image_Tensor)            # 没有数据归一化操作

image_Tensor_Nor = transform_picture(image).unsqueeze(0)
print(image_Tensor_Nor)        # 有数据归一化操作
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

[[[158 112  49]
  [159 111  47]
  [165 116  51]
  ...
  [ 24  77 124]
  [ 34  84 129]
  [ 21  67 110]]]
  
tensor([[[[0.6196, 0.6235, 0.6471,  ..., 0.5373, 0.4941, 0.4549],
          [0.5961, 0.5922, 0.6235,  ..., 0.5333, 0.4902, 0.4667],
          [0.5922, 0.5922, 0.6196,  ..., 0.5451, 0.5098, 0.4706],
          ...,
          [0.6941, 0.5804, 0.5373,  ..., 0.5725, 0.4235, 0.4980],
          [0.6588, 0.5804, 0.5176,  ..., 0.5098, 0.4941, 0.4196],
          [0.6275, 0.5843, 0.5176,  ..., 0.4863, 0.5059, 0.4314]]]])
          
tensor([[[[ 0.6338,  0.6531,  0.7694,  ...,  0.2267,  0.0134, -0.1804],
          [ 0.5174,  0.4981,  0.6531,  ...,  0.2073, -0.0060, -0.1223],
          [ 0.4981,  0.4981,  0.6338,  ...,  0.2654,  0.0910, -0.1029],
          ...,
          [ 1.2319,  0.6661,  0.4515,  ...,  0.6271, -0.1143,  0.2564],
          [ 1.0563,  0.6661,  0.3540,  ...,  0.3149,  0.2369, -0.1338],
          [ 0.9003,  0.6856,  0.3540,  ...,  0.1979,  0.2954, -0.0753]]]])
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

Tensor 转 Numpy ： （用于显示图像）

mean = (0.4914, 0.4822, 0.4465)
std = (0.2023, 0.1994, 0.2010)

def tensor_numpy(image):
    clean = image.clone().detach().cpu().squeeze(0)        # 去掉batch通道 (batch, C, H, W) --> (C, H, W)
    clean[0] = clean[0] * std[0] + mean[0]                 # 数据去归一化
    clean[1] = clean[1] * std[1] + mean[1]
    clean[2] = clean[2].mul(std[2]) + mean[2]
    clean = np.around(clean.mul(255))                     # 转换到颜色255 [0, 1] --> [0, 255]
    clean = np.uint8(clean).transpose(1, 2, 0)            # 跟换三通道 (C, H, W) --> (H, W, C)
    r, g, b = cv.split(clean)                             # RGB 通道转换
    clean = cv.merge([b, g, r])
    return clean

Num = tensor_numpy(image_Tensor_Nor)
print(Num)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

如果使用 cv.imshow() 需要使用上面的 RGB 通道转换；
如果使用 plt.imshow() 不需要使用上面的 RGB 通道转换；

示例：

plt.imshow('image', image)
plt.show()
1
2

RGB 格式：

[[[158 112  49]
  [159 111  47]
  [165 116  51]
  ...
  [ 24  77 124]
  [ 34  84 129]
  [ 21  67 110]]]
1
2
3
4
5
6
7

RGB 通道转换：

r, g, b = cv.split(image)
pic = cv.merge([b, g, r])
plt.imshow('image', pic)
plt.show()
1
2
3
4

BGR 格式：

[[[ 49 112 158]
  [ 47 111 159]
  [ 51 116 165]
  ...
  [124  77  24]
  [129  84  34]
  [110  67  21]]]
1
2
3
4
5
6
7

PyTorch save_image()

torchvision.utils.save_image()

最后再来看一下 pytorch 自带的函数是如何进行格式转换保存图片的 ：

def save_image(
    tensor: Union[torch.Tensor, List[torch.Tensor]],
    fp: Union[Text, pathlib.Path, BinaryIO],
    format: Optional[str] = None,
    **kwargs
) 

    grid = make_grid(tensor, **kwargs)
    # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
    ndarr = grid.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to('cpu', torch.uint8).numpy()
    im = Image.fromarray(ndarr)
    im.save(fp, format=format)
1
2
3
4
5
6
7
8
9
10
11
12