700字范文,内容丰富有趣,生活中的好帮手!
700字范文 > 将tensor张量转换成图片格式并保存

将tensor张量转换成图片格式并保存

时间:2018-12-06 16:28:30

相关推荐

将tensor张量转换成图片格式并保存

这是一个工具包,功能:反向操作transforms.Normalize和transforms.ToTensor函数,将tensor格式的图片转换成.jpg、.png格式的图片

(注:这里是我原始的写法,但是是存在着一些改进空间的【如评论区所言,会改变img_tensor的内容,但我这边就暂时保持原样了,详细的改进方法可见评论区】【感谢评论区大佬!】)

import torchimport numpy as npfrom PIL import Imageimport 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 pltimg = Image.open('./lena.jpg')ToTensor_transform = pose([transforms.ToTensor()])img_tensor = ToTensor_transform(img)img = transform_convert(img_tensor, ToTensor_transform)plt.imshow(img)plt.savefig('./lenaa.png')

本内容不代表本网观点和政治立场,如有侵犯你的权益请联系我们处理。
网友评论
网友评论仅供其表达个人看法,并不表明网站立场。