快捷導(dǎo)航

Pytorch中Tensor與各種圖像格式的相互轉(zhuǎn)化詳解

更新時間：2019年12月26日 14:34:00 作者：Oldpan

這篇文章主要介紹了Pytorch中Tensor與各種圖像格式的相互轉(zhuǎn)化詳解，文中通過示例代碼介紹的非常詳細(xì)，對大家的學(xué)習(xí)或者工作具有一定的參考學(xué)習(xí)價值，需要的朋友們下面隨著小編來一起學(xué)習(xí)學(xué)習(xí)吧

前言

在pytorch中經(jīng)常會遇到圖像格式的轉(zhuǎn)化，例如將PIL庫讀取出來的圖片轉(zhuǎn)化為Tensor，亦或者將Tensor轉(zhuǎn)化為numpy格式的圖片。而且使用不同圖像處理庫讀取出來的圖片格式也不相同，因此，如何在pytorch中正確轉(zhuǎn)化各種圖片格式(PIL、numpy、Tensor)是一個在調(diào)試中比較重要的問題。

本文主要說明在pytorch中如何正確將圖片格式在各種圖像庫讀取格式以及tensor向量之間轉(zhuǎn)化的問題。以下代碼經(jīng)過測試都可以在Pytorch-0.4.0或0.3.0版本直接使用。

對python不同的圖像庫讀取格式有疑問可以看這里：http://www.dbjr.com.cn/article/177288.htm

格式轉(zhuǎn)換

我們一般在pytorch或者python中處理的圖像無非這幾種格式：

PIL：使用python自帶圖像處理庫讀取出來的圖片格式
numpy：使用python-opencv庫讀取出來的圖片格式
tensor：pytorch中訓(xùn)練時所采取的向量格式（當(dāng)然也可以說圖片）

注意，之后的講解圖片格式皆為RGB三通道，24-bit真彩色，也就是我們平常使用的圖片形式。

PIL與Tensor

PIL與Tensor的轉(zhuǎn)換相對容易些，因為pytorch已經(jīng)提供了相關(guān)的代碼，我們只需要搭配使用即可：

所有代碼都已經(jīng)引用了（之后的代碼省略引用部分）：

import torch
from PIL import Image
import matplotlib.pyplot as plt

# loader使用torchvision中自帶的transforms函數(shù)
loader = transforms.Compose([
  transforms.ToTensor()]) 

unloader = transforms.ToPILImage()

1 PIL讀取圖片轉(zhuǎn)化為Tensor

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

2 將PIL圖片轉(zhuǎn)化為Tensor

# 輸入PIL格式圖片
# 返回tensor變量
def PIL_to_tensor(image):
  image = loader(image).unsqueeze(0)
  return image.to(device, torch.float)

3 Tensor轉(zhuǎn)化為PIL圖片

# 輸入tensor變量
# 輸出PIL格式圖片
def tensor_to_PIL(tensor):
  image = tensor.cpu().clone()
  image = image.squeeze(0)
  image = unloader(image)
  return image

4 直接展示tensor格式圖片

def imshow(tensor, title=None):
  image = tensor.cpu().clone() # we clone the tensor to not do changes on it
  image = image.squeeze(0) # remove the fake batch dimension
  image = unloader(image)
  plt.imshow(image)
  if title is not None:
    plt.title(title)
  plt.pause(0.001) # pause a bit so that plots are updated

5 直接保存tensor格式圖片

def save_image(tensor, **para):
  dir = 'results'
  image = tensor.cpu().clone() # we clone the tensor to not do changes on it
  image = image.squeeze(0) # remove the fake batch dimension
  image = unloader(image)
  if not osp.exists(dir):
    os.makedirs(dir)
  image.save('results_{}/s{}-c{}-l{}-e{}-sl{:4f}-cl{:4f}.jpg'
        .format(num, para['style_weight'], para['content_weight'], para['lr'], para['epoch'],
            para['style_loss'], para['content_loss']))

numpy與Tensor

numpy格式是使用cv2，也就是python-opencv庫讀取出來的圖片格式，需要注意的是用python-opencv讀取出來的圖片和使用PIL讀取出來的圖片數(shù)據(jù)略微不同，經(jīng)測試用python-opencv讀取出來的圖片在訓(xùn)練時的效果比使用PIL讀取出來的略差一些(詳細(xì)過程之后發(fā)布)。

之后所有代碼引用：

import cv2
import torch
import matplotlib.pyplot as plt

numpy轉(zhuǎn)化為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

tensor轉(zhuǎn)化為numpy

def tensor_to_np(tensor):
  img = tensor.mul(255).byte()
  img = img.cpu().numpy().squeeze(0).transpose((1, 2, 0))
  return img

展示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)

展示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)

注意

上面介紹的都是一張圖片的轉(zhuǎn)化，如果是n張圖片一起的話，只需要修改一下相應(yīng)代碼即可。

舉個例子，將之前說過的修改略微修改一下即可：

# 將 N x H x W X C 的numpy格式圖片轉(zhuǎn)化為相應(yīng)的tensor格式
def toTensor(img):
  img = torch.from_numpy(img.transpose((0, 3, 1, 2)))
  return img.float().div(255).unsqueeze(0)

以上就是本文的全部內(nèi)容，希望對大家的學(xué)習(xí)有所幫助，也希望大家多多支持腳本之家。

您可能感興趣的文章: