Pytorch中Tensor與各種影象格式的相互轉化詳解
前言
在pytorch中經常會遇到影象格式的轉化,例如將PIL庫讀取出來的圖片轉化為Tensor,亦或者將Tensor轉化為numpy格式的圖片。而且使用不同影象處理庫讀取出來的圖片格式也不相同,因此,如何在pytorch中正確轉化各種圖片格式(PIL、numpy、Tensor)是一個在除錯中比較重要的問題。
本文主要說明在pytorch中如何正確將圖片格式在各種影象庫讀取格式以及tensor向量之間轉化的問題。以下程式碼經過測試都可以在Pytorch-0.4.0或0.3.0版本直接使用。
對python不同的影象庫讀取格式有疑問可以看這裡:https://www.jb51.net/article/177288.htm
格式轉換
我們一般在pytorch或者python中處理的影象無非這幾種格式:
- PIL:使用python自帶影象處理庫讀取出來的圖片格式
- numpy:使用python-opencv庫讀取出來的圖片格式
- tensor:pytorch中訓練時所採取的向量格式(當然也可以說圖片)
注意,之後的講解圖片格式皆為RGB三通道,24-bit真彩色,也就是我們平常使用的圖片形式。
PIL與Tensor
PIL與Tensor的轉換相對容易些,因為pytorch已經提供了相關的程式碼,我們只需要搭配使用即可:
所有程式碼都已經引用了(之後的程式碼省略引用部分):
import torch from PIL import Image import matplotlib.pyplot as plt # loader使用torchvision中自帶的transforms函式 loader = transforms.Compose([ transforms.ToTensor()]) unloader = transforms.ToPILImage()
1 PIL讀取圖片轉化為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圖片轉化為Tensor
# 輸入PIL格式圖片 # 返回tensor變數 def PIL_to_tensor(image): image = loader(image).unsqueeze(0) return image.to(device,torch.float)
3 Tensor轉化為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讀取出來的圖片資料略微不同,經測試用python-opencv讀取出來的圖片在訓練時的效果比使用PIL讀取出來的略差一些(詳細過程之後釋出)。
之後所有程式碼引用:
import cv2 import torch import matplotlib.pyplot as plt
numpy轉化為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,1))) return img.float().div(255).unsqueeze(0) # 255也可以改為256
tensor轉化為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)
注意
上面介紹的都是一張圖片的轉化,如果是n張圖片一起的話,只需要修改一下相應程式碼即可。
舉個例子,將之前說過的修改略微修改一下即可:
# 將 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)
以上就是本文的全部內容,希望對大家的學習有所幫助,也希望大家多多支援我們。