将tensor张量转换成图片格式并保存
生活随笔
收集整理的這篇文章主要介紹了
将tensor张量转换成图片格式并保存
小編覺得挺不錯的,現在分享給大家,幫大家做個參考.
這是一個工具包,功能:反向操作transforms.Normalize和transforms.ToTensor函數,將tensor格式的圖片轉換成.jpg、.png格式的圖片
?(注:這里是我原始的寫法,但是是存在著一些改進空間的【如評論區所言,會改變img_tensor的內容,但我這邊就暫時保持原樣了,詳細的改進方法可見評論區】【感謝評論區大佬!】)
import torch import numpy as np from PIL import Image import torchvision.transforms as transformsdef transform_convert(img_tensor, transform):"""param img_tensor: tensorparam transforms: torchvision.transforms"""if 'Normalize' in str(transform):normal_transform = list(filter(lambda x: isinstance(x, transforms.Normalize), transform.transforms))mean = torch.tensor(normal_transform[0].mean, dtype=img_tensor.dtype, device=img_tensor.device)std = torch.tensor(normal_transform[0].std, dtype=img_tensor.dtype, device=img_tensor.device)img_tensor.mul_(std[:,None,None]).add_(mean[:,None,None])img_tensor = img_tensor.transpose(0,2).transpose(0,1) # C x H x W ---> H x W x Cif 'ToTensor' in str(transform) or img_tensor.max() < 1:img_tensor = img_tensor.detach().numpy()*255if isinstance(img_tensor, torch.Tensor):img_tensor = img_tensor.numpy()if img_tensor.shape[2] == 3:img = Image.fromarray(img_tensor.astype('uint8')).convert('RGB')elif img_tensor.shape[2] == 1:img = Image.fromarray(img_tensor.astype('uint8')).squeeze()else:raise Exception("Invalid img shape, expected 1 or 3 in axis 2, but got {}!".format(img_tensor.shape[2]))return img代碼示例
import matplotlib.pyplot as plt img = Image.open('./lena.jpg') ToTensor_transform = transforms.Compose([transforms.ToTensor()]) img_tensor = ToTensor_transform(img) img = transform_convert(img_tensor, ToTensor_transform) plt.imshow(img) plt.savefig('./lenaa.png')總結
以上是生活随笔為你收集整理的将tensor张量转换成图片格式并保存的全部內容,希望文章能夠幫你解決所遇到的問題。
- 上一篇: torchvision 笔记:trans
- 下一篇: torchvision 笔记:trans