text: import os
  import torch
  import torchvision
  import numpy as np
  from PIL import Image
  import matplotlib.pyplot as plt
  from collections import OrderedDict
  from torchvision.transforms import ToPILImage
  
  def batch_transform(batch, transform):
      """Applies a transform to a batch of samples.
      Keyword arguments:
      - batch (): a batch os samples
      - transform (callable): A function/transform to apply to ``batch``
      """
      # Convert the single channel label to RGB in tensor form
      # 1. torch.unbind removes the 0-dimension of "labels" and returns a tuple of
      # all slices along that dimension
      # 2. the transform is applied to each slice
      transf_slices = [transform(tensor) for tensor in torch.unbind(batch)]
      return torch.stack(transf_slices)
  
  def imshow_batch(images, labels):
      """Displays two grids of images. The top grid displays ``images``
      and the bottom grid ``labels``
      Keyword arguments:
      - images (``Tensor``): a 4D mini-batch tensor of shape
      (B, C, H, W)
      - labels (``Tensor``): a 4D mini-batch tensor of shape
      (B, C, H, W)
      """
      # Make a grid with the images and labels and convert it to numpy
      for i in range(images.size()[0]):
          img = (torchvision.utils.make_grid(images[i,:,:,:]).numpy()*255).astype(np.uint8)
          lb = (torchvision.utils.make_grid(labels[i,:,:]).numpy()*255).astype(np.uint8)
          img = np.moveaxis(img, 0, -1)
          lb = np.moveaxis(lb, 0, -1)
          fig,(ax1,ax2 ) = plt.subplots(2,1)
          ax1.imshow(img)
          ax2.imshow(lb)
          plt.show()
      
  def save_checkpoint(model, optimizer, epoch, miou, args):
      """Saves the model in a specified directory with a specified name.save
      Keyword arguments:
      - model (``nn.Module``): The model to save.
      - optimizer (``torch.optim``): The optimizer state to save.
      - epoch (``int``): The current epoch for the model.
      - miou (``float``): The mean IoU obtained by the model.
      - args (``ArgumentParser``): An instance of ArgumentParser which contains
      the arguments used to train ``model``. The arguments are written to a text
      file in ``args.save_dir`` named "``args.name``_args.txt".
      """
      name = args.name
      save_dir = args.save_dir
      assert os.path.isdir(
          save_dir), "The directory \"{0}\" doesn't exist.".format(save_dir)
      # Save model
      model_path = os.path.join(save_dir, name)
      checkpoint = {
          'epoch': epoch,
          'miou': miou,
          'state_dict': model.state_dict(),
          'optimizer': optimizer.state_dict()
      }
      torch.save(checkpoint, model_path)
      # Save arguments
      summary_filename = os.path.join(save_dir, name + '_summary.txt')
      with open(summary_filename, 'w') as summary_file:
          sorted_args = sorted(vars(args))
          summary_file.write("ARGUMENTS\n")
          for arg in sorted_args:
              arg_str = "{0}: {1}\n".format(arg, getattr(args, arg))
              summary_file.write(arg_str)
          summary_file.write("\nBEST VALIDATION\n")
          summary_file.write("Epoch: {0}\n". format(epoch))
          summary_file.write("Mean IoU: {0}\n". format(miou))
  
  def load_checkpoint(model, optimizer, folder_dir, filename,reset_optimizer=False):
      """Saves the model in a specified directory with a specified name.save
      Keyword arguments:
      - model (``nn.Module``): The stored model state is copied to this model
      instance.
      - optimizer (``torch.optim``): The stored optimizer state is copied to this
      optimizer instance.
      - folder_dir (``string``): The path to the folder where the saved model
      state is located.
      - filename (``string``): The model filename.
      Returns:
      The epoch, mean IoU, ``model``, and ``optimizer`` loaded from the
      checkpoint.
      """
      assert os.path.isdir(
          folder_dir), "The directory \"{0}\" doesn't exist.".format(folder_dir)
      # Create folder to save model and information
      model_path = os.path.join(folder_dir, filename)
      assert os.path.isfile(
          model_path), "The model file \"{0}\" doesn't exist.".format(filename)
      # Load the stored model parameters to the model instance
      checkpoint = torch.load(model_path)
      model.load_state_dict(checkpoint['state_dict'])
      if not reset_optimizer:
          optimizer.load_state_dict(checkpoint['optimizer'])
      epoch = checkpoint['epoch']
      miou = checkpoint['miou']
      return model, optimizer, epoch, miou
  
  class PILToLongTensor(object):
      """Converts a ``PIL Image`` to a ``torch.LongTensor``.
      Code adapted from: http://pytorch.org/docs/master/torchvision/transforms.html?highlight=totensor
      """
      def __call__(self, pic):
          """Performs the conversion from a ``PIL Image`` to a ``torch.LongTensor``.
          Keyword arguments:
          - pic (``PIL.Image``): the image to convert to ``torch.LongTensor``
          Returns:
          A ``torch.LongTensor``.
          """
          if not isinstance(pic, Image.Image):
              raise TypeError("pic should be PIL Image. Got {}".format(
                  type(pic)))
          # handle numpy array
          if isinstance(pic, np.ndarray):
              img = torch.from_numpy(pic.transpose((2, 0, 1)))
              # backward compatibility
              return img.long()
          # Convert PIL image to ByteTensor
          img = torch.ByteTensor(torch.ByteStorage.from_buffer(pic.tobytes()))
          # Reshape tensor
          nchannel = len(pic.mode)
          if(nchannel != 1):
              print("img is not greyscale")
          img = img.view(pic.size[1], pic.size[0], nchannel)
          # Convert to long and squeeze the channels
          return img.transpose(0, 1).transpose(0,2).contiguous().long().squeeze_()
  
  class LongTensorToRGBPIL(object):
      """Converts a ``torch.LongTensor`` to a ``PIL image``.
      The input is a ``torch.LongTensor`` where each pixel's value identifies the class.
      Keyword arguments:
      - rgb_encoding (``OrderedDict``): An ``OrderedDict`` that relates pixel
      values, class names, and class colors.
      """
      def __init__(self, rgb_encoding):
          self.rgb_encoding = rgb_encoding
      def __call__(self, tensor):
          """Performs the conversion from ``torch.LongTensor`` to a ``PIL image``
          Keyword arguments:
          - tensor (``torch.LongTensor``): the tensor to convert
          Returns:
          A ``PIL.Image``.
          """
          # Check if label_tensor is a LongTensor
          if not isinstance(tensor, torch.LongTensor):
              raise TypeError("label_tensor should be torch.LongTensor. Got {}"
                              .format(type(tensor)))
          # Check if encoding is a ordered dictionary
          if not isinstance(self.rgb_encoding, OrderedDict):
              raise TypeError("encoding should be an OrderedDict. Got {}".format(
                  type(self.rgb_encoding)))
          # label_tensor might be an image without a channel dimension, in this
          # case unsqueeze it
          if len(tensor.size()) == 2:
              tensor.unsqueeze_(0)
          color_tensor = torch.ByteTensor(3, tensor.size(1), tensor.size(2))
          for index, (class_name, color) in enumerate(self.rgb_encoding.items()):
              # Get a mask of elements equal to index
              mask = torch.eq(tensor, index).squeeze_()
              # Fill color_tensor with corresponding colors
              for channel, color_value in enumerate(color):
                  color_tensor[channel].masked_fill_(mask, color_value)
  
          return ToPILImage()(color_tensor)
          