code: """
  paper: Pyramid Channel-based Feature Attention Network for image dehazing 
  file: utils.py
  about: all utilities
  author: Tao Wang
  date: 12/01/2021
  """
  
  # --- Imports --- #
  import time
  import torch
  import torch.nn.functional as F
  import torchvision.utils as utils
  from math import log10
  from skimage import measure
  from torch.autograd import Variable
  import os
  
  def to_psnr(dehaze, gt):
      mse = F.mse_loss(dehaze, gt, reduction='none')
      mse_split = torch.split(mse, 1, dim=0)
      mse_list = [torch.mean(torch.squeeze(mse_split[ind])).item() for ind in range(len(mse_split))]
  
      intensity_max = 1.0
      psnr_list = [10.0 * log10(intensity_max / mse) for mse in mse_list]
      return psnr_list
  
  
  def to_ssim_skimage(dehaze, gt):
      dehaze_list = torch.split(dehaze, 1, dim=0)
      gt_list = torch.split(gt, 1, dim=0)
  
      dehaze_list_np = [dehaze_list[ind].permute(0, 2, 3, 1).data.cpu().numpy().squeeze() for ind in range(len(dehaze_list))]
      gt_list_np = [gt_list[ind].permute(0, 2, 3, 1).data.cpu().numpy().squeeze() for ind in range(len(dehaze_list))]
      ssim_list = [measure.compare_ssim(dehaze_list_np[ind],  gt_list_np[ind], data_range=1, multichannel=True) for ind in range(len(dehaze_list))]
  
      return ssim_list
  
  
  def validation(net, val_data_loader, device, category, save_tag=False):
      """
      :param net: PCFAN
      :param val_data_loader: validation loader
      :param device: The GPU that loads the network
      :param category: indoor or outdoor test dataset
      :param save_tag: tag of saving image or not
      :return: average PSNR value
      """
      psnr_list = []
      ssim_list = []
  
      for batch_id, val_data in enumerate(val_data_loader):
  
          with torch.no_grad():
              haze, gt, image_name = Variable(val_data['hazy_image']), Variable(val_data['clear_image']),val_data['haze_name']
              haze = haze.to(device)
              gt = gt.to(device)
              dehaze = net(haze)
  
          # --- Calculate the average PSNR --- #
          psnr_list.extend(to_psnr(dehaze, gt))
  
          # --- Calculate the average SSIM --- #
          ssim_list.extend(to_ssim_skimage(dehaze, gt))
  
          # --- Save image --- #
          if save_tag:
              path = './results/{}_results'.format(category)
              if not os.path.exists(path):
                  os.makedirs(path)
              save_image(dehaze, image_name, category)
  
      avr_psnr = sum(psnr_list) / len(psnr_list)
      avr_ssim = sum(ssim_list) / len(ssim_list)
      return avr_psnr, avr_ssim
  
  
  def save_image(dehaze, image_name, category):
      dehaze_images = torch.split(dehaze, 1, dim=0)
      batch_num = len(dehaze_images)
      for ind in range(batch_num):
          utils.save_image(dehaze_images[ind], './results/{}_results/{}'.format(category, image_name[ind][:-3] + 'png'))
  
  
  def print_log(epoch, num_epochs, train_psnr, val_psnr, val_ssim, category):
      print('Epoch [{0}/{1}], Train_PSNR:{2:.2f}, Val_PSNR:{3:.2f}, Val_SSIM:{4:.4f}'
            .format(epoch, num_epochs, train_psnr, val_psnr, val_ssim))
  
      # --- Write the training log --- #
      with open('./logs/{}_log.txt'.format(category), 'a') as f:
          print('Date: {0}s, Epoch: [{1}/{2}], Train_PSNR: {3:.2f}, Val_PSNR: {4:.2f}, Val_SSIM: {5:.4f}'
                .format(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()),
                        epoch, num_epochs, train_psnr, val_psnr, val_ssim), file=f)
apis: torch.nn.functional.mse_loss torch.split torch.autograd.Variable torch.no_grad torch.squeeze