图片保存:torchvision.utils.save_image(img, imgPath)

12 篇文章 2 订阅
12 篇文章 1 订阅

torchvision.utils.save_image(img, imgPath)

深度学习模型中,一般使用如下方式进行图像保存(torchvision.utils中的save_image()函数),这种方式只能保存RGB彩色图像,如果网络的输出是单通道灰度图像,则该函数依然会输出三个通道,每个通道的数值都是相同的,即“伪灰度图像”,虽然从视觉效果上看不出区别,但是图像所占内存比正常情况大了两倍。

解决方式

from PIL import Image
import torch
from torchvision.utils import make_grid
import os
import torchvision
import numpy as np

# torchvision.utils.save_image(img, imgPath)

def save_multi_channel_image(tensor, filename, nrow=8, padding=2,
               normalize=False, range=None, scale_each=False, pad_value=0):
    """Save a given Tensor into an image file.

    Args:
        tensor (Tensor or list): Image to be saved. If given a mini-batch tensor,
            saves the tensor as a grid of images by calling ``make_grid``.
        **kwargs: Other arguments are documented in ``make_grid``.
    """
    # 用于把几个图像按照网格排列的方式绘制出来
    grid = make_grid(tensor, nrow=nrow, padding=padding, pad_value=pad_value,
                     normalize=normalize, range=range, scale_each=scale_each)
    # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
    ndarr = grid.mul_(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to('cpu', torch.uint8).numpy()
    im = Image.fromarray(ndarr)
    im.save(filename)


def save_gray_image(img, save_dir, type, name, Gray=False):
    fname, fext = name.split('.')
    imgPath = os.path.join(save_dir, "%s_%s.%s" % (fname, type, fext))
    # torchvision.utils.save_image(img, imgPath)
    # 改写:torchvision.utils.save_image
    grid = torchvision.utils.make_grid(img, nrow=8, padding=2, pad_value=0,
                                       normalize=False, range=None, scale_each=False)
    ndarr = grid.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to('cpu', torch.uint8).numpy()
    im = Image.fromarray(ndarr)
    # im.show()
    if Gray:
        im.convert('L').save(imgPath)  # Gray = 0.29900 * R + 0.58700 * G + 0.11400 * B
    else:
        im.save(imgPath)


# 适用各种图片保存
def save_tensor_image(img, save_dir, type, name, Gray=False):
    fname, fext = name.split('.')
    imgPath = os.path.join(save_dir, "%s_%s.%s" % (fname, type, fext))
    img_array = tensor2array(img.data[0])
    image_pil = Image.fromarray(img_array)
    if Gray:
        image_pil.convert('L').save(imgPath)  # Gray = 0.29900 * R + 0.58700 * G + 0.11400 * B
    else:
        image_pil.save(imgPath)

# Converts a Tensor into a Numpy array
# |imtype|: the desired type of the converted numpy array
def tensor2array(image_tensor, imtype=np.uint8, normalize=True):
    if isinstance(image_tensor, list):
        image_numpy = []
        for i in range(len(image_tensor)):
            image_numpy.append(tensor2array(image_tensor[i], imtype, normalize))
        return image_numpy
    image_numpy = image_tensor.cpu().float().numpy()
    if normalize:
        image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + 1) / 2.0 * 255.0
    else:
        image_numpy = np.transpose(image_numpy, (1, 2, 0)) * 255.0
    image_numpy = np.clip(image_numpy, 0, 255)
    if image_numpy.shape[2] == 1 or image_numpy.shape[2] > 3:
        image_numpy = image_numpy[:, :, 0]
    return image_numpy.astype(imtype)

  • 22
    点赞
  • 68
    收藏
    觉得还不错? 一键收藏
  • 打赏
    打赏
  • 1
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论 1
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

代码小白的成长

计算机网络PPT下载

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值