【图像处理】tensor 与 PIL Image/ numpy 的相互转换

numpy 与 tensor 的相互转换

首先以 OpenCV读取图像，返回的是一个代表图像的NumPy数组，采用的格式是 H×W×C，通道顺序为BGR。

import cv2
path = "D:/Deeplearning/img/dog.jpg"
image = cv2.imread(path)  # 返回的是 numpy 数组

print(image.shape)  # (433, 650, 3), 即 (H, W, C)

numpy 转换成 tensor

以torch.from_numpy(image)的方式，将 numpy 转换成 tensor。

import torch

# 先将图像从BGR转换成RGB
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

image_tensor = torch.from_numpy(image)  
# image_tensor = torch.tensor(image) 这种方式也可以

print(type(image_tensor))  # <class 'torch.Tensor'> 
print(image_tensor.shape)  # torch.Size([433, 650, 3]), 即 (H,W,C)

使用matplotlib绘图时，可以读取tesnor和numpy数据类型，图片通道顺序为RGB，且尺寸为(H,W,C)。

使用 matplotlib 画出转换成 tensor 后的图像，检查是否能够正常显示。

import matplotlib.pyplot as plt
# 在 jupyter 中需要加上下面一行代码
# %matplotlib inline

plt.imshow(image_tensor)  # 正常显示图片
plt.axis("off")  # 不显示坐标轴
plt.show()

tensor 转换成 numpy

以image_tensor.numpy()的方式，将 tensor 转换成 numpy。

image_np = image_tensor.numpy()  

print(type(image_np))  # <class 'numpy.ndarray'>
print(image_np.shape)  # (433, 650, 3), 即 (H,W,C)

可以使用matplotlib绘图。同样，可以使用 OpenCV 显示转换回 numpy 数组的图片，检验是否能够正常显示。

# 使用OpenCV显示图像需要将RGB图像转换回BGR
image_np = image_np[:, :, ::-1]

cv2.imshow("Image_np", image_np)
cv2.waitKey(0)

PIL Image 与 tensor 的相互转换

首先使用 PIL 读取一张图像。

from PIL import Image

path = "D:/Deeplearning/img/dog.jpg"
image = Image.open(path)  # PIL Image 类型

print(image.size)  # (650, 433), 即 (W, H)

Image 转换成 tensor

使用transforms.ToTensor()将 PIL Image 转换成 Tensor 类型。

from torchvision import transforms

trans_tensor = transforms.ToTensor()
image_tensor = trans_tensor(image)

print(type(image_tensor))  # <class 'torch.Tensor'>
print(image_tensor.shape)  # torch.Size([3, 433, 650]),即(C, H, W)

使用 matplotlib 画出转换成 tensor 后的图像，检查是否能够正常显示。

注意：现在的 image_tensor 为 C×H×W，而 matplotlib 需要是 H×W×C，需要先使用image_tensor.permute(1, 2, 0)调换一下顺序。

import matplotlib.pyplot as plt

image_hwc = image_tensor.permute(1, 2, 0)  # 将通道(C,H,W)转换成(H,W,C)
plt.imshow(image_hwc)  # 显示正常
plt.axis("off")  # 不显示坐标轴
plt.show() 

tensor 转换成 Image

使用transforms.ToPILImage()将 PIL Image 转换成 Tensor 类型。

trans_pil = transforms.ToPILImage()
image_pil = trans_pil(image_tensor)  # 这里是image_tensor,不是image_hwc

print(type(image_pil))  # <class 'PIL.Image.Image'>
print(image_pil.size)  # (650, 433), 即 (W, H)

使用 PIL 检查能否正常显示转换回 Image 的图像。

image_pil.show()  # 显示正常