[pytorch] Tensor, numpy与PIL格式的相互转换

[pytorch] Tensor, numpy与PIL格式的相互转换

我们一般在pytorch或者python中处理的图像无非这几种格式:

  • PIL:使用python自带图像处理库读取出来的图片格式
  • Numpy:使用python-opencv库读取出来的图片格式
  • Tensor:pytorch中训练时所采取的向量格式(
    注意,之后的讲解图片格式皆为RGB三通道,24-bit真彩色,也就是我们平常使用的图片形式。

PIL与Tensor

  1. PIL to Tensor
import torch
from PIL import Image
import matplotlib.pyplot as plt

# loader使用torchvision中自带的transforms函数
loader = transforms.Compose([
    transforms.ToTensor()])  

unloader = transforms.ToPILImage()


# 输入图片地址
# 返回tensor变量
def image_loader(image_name):
    image = Image.open(image_name).convert('RGB')
    image = loader(image).unsqueeze(0)
    return image.to(device, torch.float)

  1. Tensor to PIL
# 输入tensor变量
# 输出PIL格式图片
def tensor_to_PIL(tensor):
    image = tensor.cpu().clone()
    image = image.squeeze(0)
    image = unloader(image)
    return image

Tensor与Numpy

  1. Tensor to Numpy
import cv2
import torch
import matplotlib.pyplot as plt

def tensor_to_np(tensor):
    img = tensor.mul(255).byte()
    img = img.cpu().numpy().squeeze(0).transpose((1, 2, 0))
    return img
  1. Numpy to Tensor
def toTensor(img):
    assert type(img) == np.ndarray,'the img type is {}, but ndarry expected'.format(type(img))
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = torch.from_numpy(img.transpose((2, 0, 1)))
    return img.float().div(255).unsqueeze(0)  # 255也可以改为256

图片展示

  1. 展示Tensor
def show_from_tensor(tensor, title=None):
    img = tensor.clone()
    img = tensor_to_np(img)
    plt.figure()
    plt.imshow(img)
    if title is not None:
        plt.title(title)
    plt.pause(0.001)

  1. 展示Numpy
def show_from_cv(img, title=None):
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    plt.figure()
    plt.imshow(img)
    if title is not None:
        plt.title(title)
    plt.pause(0.001)

多张图片的转换

以上代码都是单张图片格式的转换,若改为多张,则只需要稍微改动代码,以下为例:

# 将 N x H x W X C 的numpy格式图片转化为相应的tensor格式
def toTensor(img):
    img = torch.from_numpy(img.transpose((0, 3, 1, 2)))
    return img.float().div(255).unsqueeze(0)

已标记关键词 清除标记
©️2020 CSDN 皮肤主题: 数字20 设计师:CSDN官方博客 返回首页