【轉載】pytorch中Tensor、PIL、numpy的相互轉化

【轉載】pytorch中Tensor、PIL、numpy的相互轉化

原博客地址:Pytorch中Tensor與各種圖像格式的相互轉化

1.常用圖像格式:

     (1)PIL:使用python自帶圖像處理庫讀取出來的圖片格式;
     (2)numpy:使用python-opencv庫讀取出來的圖片格式;
     (3)tensor:pytorch中訓練時所採取的向量格式(圖片)。
       注意,之後的講解圖片格式皆爲RGB三通道,24-bit真彩色,也就是常用的圖像格式。


2.PIL與Tensor相互轉換:

     以下所有代碼都已經被引用。

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

3.numpy與Tensor相互轉換:

     numpy格式是使用cv2,也就是python-opencv庫讀取出來的圖片格式,需要注意的是用python-opencv讀取出來的圖片和使用PIL讀取出來的圖片數據略有不同,經測試用python-opencv讀取出來的圖片在訓練時的效果比使用PIL讀取出來的略差一些。

     以下所有代碼都已經被引用。

import cv2
import torch
import matplotlib.pyplot as plt

     (1)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, 0, 1)))
    return img.float().div(255).unsqueeze(0)  # 255也可以改爲256

     (2)tensor轉化爲numpy:

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

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

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

4.多張圖片處理舉例:

# 將 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)

 

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章