pytorch拼接与保存图片

在深度学习训练时，我们常需要对某层的输出或最终模型的预测结果进行输出查看，如featmap、output、predict结果等。在pytorch的设定中有tensor类型变量，不可直接查看和保存图片。

 虽然有 torchvision.utils.save_image 方法可以方便地保存tensor类型图片，调用方法如下：

from torchvision.utils import save_image

save_image(tensor , filename , padding =0)

这个方法不仅可以保存图片，而且可以保存多个图片拼接后输出。但是对于不同通道不同类型，如RGB和灰度图，的图片，仍然有不方便使用的地方，于是我自己写了几个函数，适用于图片的输出保存、拼接。

from PIL import Image

def joint_img_vertical(imgs):
  # print(imgs.ndim)
  if imgs.ndim ==4:     #RGB img
    width, height = imgs[0][0].shape
    result = Image.new('L', (width, height*imgs.shape[0]))
    # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
    imgs = imgs.mul_(255).add_(0.5).clamp_(0, 255).cpu().numpy().transpose((0, 2, 3, 1))
  elif imgs.ndim ==3:   #single channel img
    width, height = imgs[0].shape
    result = Image.new('L', (width, height*imgs.shape[0]))
    imgs =imgs*255
  for i in range(imgs.shape[0]):
    im = Image.fromarray((imgs[i]).astype(np.uint8))
    result.paste(im, box=(0, i * height))

  # 取消注释以下代码可直接保存图片
  # im = Image.fromarray((img * 255).astype(np.uint8))
  # im.save(filename)
  # result.save(filename)
  return result

def joint_img_horizontal(img,output,mask):
  # print(img.size, output.size, mask.size)
  if img.size == output.size == mask.size:
    width, height = img.size
    result = Image.new('L', (width * 3, height))
    result.paste(img, box=(0, 0))
    result.paste(output, box=(width, 0))
    result.paste(mask, box=(2 * width, 0))
  return result

def save_img(img,filename):
  '''
    it is used to save an img in filename directory
  :param img: the img should be Image type .
  :param filename: like "**.png"
  '''
  img.save(filename)

 第一个函数的作用是垂直拼接一个Batch的tensor图片。我做的是一个分割任务，在我的程序中，有RGB图片（原图）和单通道图片（mask和output），对于RGB图片，维度是（batch , channel ,width ,height）。对于单通道图片，维度是（batch , width ,height）。

对于tensor类型，要转化为numpy类型，方法是tensor.cpu().numpy.transpose((0,2,3,1)) 。注意要保存为numpy后转为PIL中的Image类型。对于三通道图，要transpose为（width, height， channel ）类型。

（这里我额外写了一个去Normalization的操作，个人视情况添加，我的代码中如果不添加这一步，会导致保存图片色彩失真，可能是我事先对数据做了数据增强的缘故）

第二个函数的作用是水平拼接三个Image类型的图片，即第一个函数的三个输出。这是为了将Image,output,mask三个图片放在同一个图片中做分割或检测效果对比观察。用到了new函数创建画布，paste方法粘贴图片到画布上。

Tip ：查看tensor或numpy类型数组的大小尺寸。用tensor.shape 、numpy.shape ，而查看Image类型图片的大小，用img.size。

第三个函数的作用是保存Image图片类型到本地，注意只有Image类型的图片才有save方法，tensor和numpy类型的自然是没有的。因为本文对于图片的操作都用PIL来实现，所以保存图片也用的此方法。同样的思想，还可以用opencv库来实现。

同时调用以上三个函数就可以实现对一个Batch类型的图片批量拼接保存。

batch_img = joint_img_vertical(data)        #data是tensor类型
batch_output = joint_img_vertical(output)   #output是array类型
batch_mask = joint_img_vertical(predict)    #predict是array类型
img_for_save = joint_img_horizontal(batch_img, batch_output,batch_mask)
save_img(img_for_save, os.path.join(save_dir,"img_out_mask_%d.png"% i))

 

 

【补充】为什么有时候保存图片会失真？

在make_grid那里，你需要设置参数normalize = True。