Pytorch中保存圖片（tensor，cv2，pillow）

tensor直接保存

#!/usr/bin/env python
# _*_ coding:utf-8 _*_
import torch
from torchvision import utils as vutils
 
 
def save_image_tensor(input_tensor: torch.Tensor, filename):
    """
    將tensor保存為圖片
    :param input_tensor: 要保存的tensor
    :param filename: 保存的文件名
    """
    assert (len(input_tensor.shape) == 4 and input_tensor.shape[0] == 1)
    # 復制一份
    input_tensor = input_tensor.clone().detach()
    # 到cpu
    input_tensor = input_tensor.to(torch.device('cpu'))
    # 反歸一化
    # input_tensor = unnormalize(input_tensor)
    vutils.save_image(input_tensor, filename)

 

 

tensor轉cv2保存

如果你是先轉numpy，再交換維度，一定用transpose,而不是swapaxes，不然顏色會出問題= =

就像下面這張圖

 

 

原圖

 

 

tensor轉cv2保存 正確的代碼

#!/usr/bin/env python
# _*_ coding:utf-8 _*_
import torch
import cv2
 
 
def save_image_tensor2cv2(input_tensor: torch.Tensor, filename):
    """
    將tensor保存為cv2格式
    :param input_tensor: 要保存的tensor
    :param filename: 保存的文件名
    """
    assert (len(input_tensor.shape) == 4 and input_tensor.shape[0] == 1)
    # 復制一份
    input_tensor = input_tensor.clone().detach()
    # 到cpu
    input_tensor = input_tensor.to(torch.device('cpu'))
    # 反歸一化
    # input_tensor = unnormalize(input_tensor)
    # 去掉批次維度
    input_tensor = input_tensor.squeeze()
    # 從[0,1]轉化為[0,255]，再從CHW轉為HWC，最后轉為cv2
    input_tensor = input_tensor.mul_(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).type(torch.uint8).numpy()
    # RGB轉BRG
    input_tensor = cv2.cvtColor(input_tensor, cv2.COLOR_RGB2BGR)
    cv2.imwrite(filename, input_tensor)

 

tensor轉pillow保存

def save_image_tensor2pillow(input_tensor: torch.Tensor, filename):
    """
    將tensor保存為pillow
    :param input_tensor: 要保存的tensor
    :param filename: 保存的文件名
    """
    assert (len(input_tensor.shape) == 4 and input_tensor.shape[0] == 1)
    # 復制一份
    input_tensor = input_tensor.clone().detach()
    # 到cpu
    input_tensor = input_tensor.to(torch.device('cpu'))
    # 反歸一化
    # input_tensor = unnormalize(input_tensor)
    # 去掉批次維度
    input_tensor = input_tensor.squeeze()
    # 從[0,1]轉化為[0,255]，再從CHW轉為HWC，最后轉為numpy
    input_tensor = input_tensor.mul_(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).type(torch.uint8).numpy()
    # 轉成pillow
    im = Image.fromarray(input_tensor)
    im.save(filename)