PyTorch 实现图片预处理（归一化 & 标准化）代码整理

一、写在前面

本来图片预处理（归一化 & 标准化）的原理不难，但我这个人工智能小白被 Tensor 和 array 等几个类型的通道顺序搞得晕头转向，好不容易整理好了，故写一篇 blog 记录一下

 

二、先上代码

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


img_path = "/lena.jpeg"


# guiyihua：[0, 255] -> [0, 1]
guiyihua = transforms.ToTensor()



# biaozhunhua (ImageNet)
biaozhunhua = transforms.Normalize(
    mean=[0.485, 0.456, 0.406],
    std=[0.229, 0.224, 0.225]
)


# zuhe
zuhe_transform = transforms.Compose([
    guiyihua,
    biaozhunhua
])


# open Image
img = Image.open(img_path)

# transformation: [H, W, C] -> [C, H, W]
img_after = zuhe_transform(img)

# tensor -> numpy
img_after = img_after.numpy()

# [C, H, W] -> [H, W, C]
img_after = np.transpose(img_after, (1, 2, 0))

# print img_after
print("img_after = ", img_after)

 

三、实现过程中一些注意的点

1、用 PIL 打开的图片 img 是标准 RGB 通道。如果转成 array【np.array(img)】的话，它的 shape 为 【Height，Width，Channel】

2、用 PIL 打开的图片 img 经过 ToTensor() 操作后，shape 会从【H，W，C】变成【C，H，W】。

3、如果想要变回原来的 shape ，需要先用【tensor.numpy()】将 tensor 类型转回 array 类型。再使用 np.transpose() 函数，带上参数(1, 2, 0)，变回【H，W，C】。至于为什么是（1，2，0），是因为对于【C，H，W】来说，0指的是第一个（即 C）；1，2分别指的是第二、三个（即 H，W）。那按照（1，2，0）的顺序一排，就是【H，W，C】了。