pytorch中 PIL与tensor、numpy与tensor、base64之间相互转化

参考链接-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)

输出如下：

-------------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])

---

知识点：

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

tf.expand_dims(input, axis=-1)

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

---

补充一个 ：cv2.imread()与base64编码、PIL格式互相转换

声明一下，这个章节内容是一个笔记补充，该部分代码原文可查阅：这位博主文章： Python中cv2.imread()与base64编码、PIL格式互相转换

import base64
import numpy as np
import cv2
from PIL import Image
from io import BytesIO
import time


# cv2转base64
def cv2_to_base64(img):
    img = cv2.imencode('.jpg', img)[1]
    image_code = str(base64.b64encode(img))[2:-1]

    return image_code


# base64转cv2
def base64_to_cv2(base64_code):
    img_data = base64.b64decode(base64_code)
    img_array = np.fromstring(img_data, np.uint8)
    img = cv2.imdecode(img_array, cv2.COLOR_RGB2BGR)

    return img


# base64转PIL
def base64_to_pil(base64_str):
    image = base64.b64decode(base64_str)
    image = BytesIO(image)
    image = Image.open(image)

    return image


# PIL转base64
def pil_to_base64(image):
    img = cv2.cvtColor(np.asarray(image), cv2.COLOR_RGB2BGR)
    base64_str = cv2_to_base64(img)

    return base64_str


# PIL转cv2
def pil_to_cv2(image):
    img = cv2.cvtColor(np.asarray(image), cv2.COLOR_RGB2BGR)

    return img


# cv2转PIL
def cv2_to_pil(img):
    image = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))

    return image


if __name__ == '__main__':
    cv2_img = cv2.imread('./test.jpg')
    start_time =time.time()
    code = cv2_to_base64(cv2_img)
    print(code)

    img = base64_to_cv2(code)
    cv2.imwrite('test.jpg', img)