Use torchmetric PSNR implementation and argument ordering

GANDLF/cli/generate_metrics.py

@@ -1,3 +1,4 @@
+import sys
 import yaml
 from pprint import pprint
 import pandas as pd
@@ -258,11 +259,14 @@ def __fix_2d_tensor(input_tensor):
  gt_image_infill, output_infill
  ).item()
 
- # PSNR - similar to pytorch PeakSignalNoiseRatio until 4 digits after decimal point
  overall_stats_dict[current_subject_id]["psnr"] = peak_signal_noise_ratio(
  gt_image_infill, output_infill
  ).item()
 
+ overall_stats_dict[current_subject_id]["psnr_eps"] = peak_signal_noise_ratio(
+ gt_image_infill, output_infill, epsilon=sys.float_info.epsilon
+ ).item()
+
  pprint(overall_stats_dict)
  if outputfile is not None:
  with open(outputfile, "w") as outfile:

GANDLF/metrics/synthesis.py

@@ -1,4 +1,3 @@
-import sys
 import SimpleITK as sitk
 import PIL.Image
 import numpy as np
@@ -8,6 +7,7 @@
  MeanSquaredError,
  MeanSquaredLogError,
  MeanAbsoluteError,
+ PeakSignalNoiseRatio,
 )
 from GANDLF.utils import get_image_from_tensor
 
@@ -25,7 +25,7 @@ def structural_similarity_index(target, prediction, mask=None) -> torch.Tensor:
  torch.Tensor: The structural similarity index.
  """
  ssim = StructuralSimilarityIndexMeasure(return_full_image=True)
- _, ssim_idx_full_image = ssim(target, prediction)
+ _, ssim_idx_full_image = ssim(preds=prediction, target=target)
  mask = torch.ones_like(ssim_idx_full_image) if mask is None else mask
  try:
  ssim_idx = ssim_idx_full_image[mask]
@@ -45,23 +45,30 @@ def mean_squared_error(target, prediction) -> torch.Tensor:
  prediction (torch.Tensor): The prediction tensor.
  """
  mse = MeanSquaredError()
- return mse(target, prediction)
+ return mse(preds=prediction, target=target)
 
 
-def peak_signal_noise_ratio(target, prediction) -> torch.Tensor:
+def peak_signal_noise_ratio(target, prediction, data_range=None, epsilon=None) -> torch.Tensor:
  """
  Computes the peak signal to noise ratio between the target and prediction.
 
  Args:
  target (torch.Tensor): The target tensor.
  prediction (torch.Tensor): The prediction tensor.
+ data_range (float, optional): If not None, this data range is used as enumerator instead of computing it from the given data. Defaults to None.
+ epsilon (float, optional): If not None, this epsilon is added to the denominator of the fraction to avoid infinity as output. Defaults to None.
  """
- mse = mean_squared_error(target, prediction)
- return (
- 10.0
- * torch.log10((torch.max(target) - torch.min(target)) ** 2)
- / (mse + sys.float_info.epsilon)
- )
+
+ if epsilon == None:
+ psnr = PeakSignalNoiseRatio(data_range=data_range)
+ return psnr(preds=prediction, target=target)
+ else: # implementation of PSNR that does not give 'inf'/'nan' when 'mse==0'
+ mse = mean_squared_error(target, prediction)
+ if data_range == None: #compute data_range like torchmetrics if not given
+ min_v = 0 if torch.min(target) > 0 else torch.min(target) #look at this line
+ max_v = torch.max(target)
+ data_range = max_v - min_v
+ return 10.0 * torch.log10((data_range ** 2) / (mse + epsilon))
 
 
 def mean_squared_log_error(target, prediction) -> torch.Tensor:
@@ -73,7 +80,7 @@ def mean_squared_log_error(target, prediction) -> torch.Tensor:
  prediction (torch.Tensor): The prediction tensor.
  """
  mle = MeanSquaredLogError()
- return mle(target, prediction)
+ return mle(preds=prediction, target=target)
 
 
 def mean_absolute_error(target, prediction) -> torch.Tensor:
@@ -85,7 +92,7 @@ def mean_absolute_error(target, prediction) -> torch.Tensor:
  prediction (torch.Tensor): The prediction tensor.
  """
  mae = MeanAbsoluteError()
- return mae(target, prediction)
+ return mae(preds=prediction, target=target)
 
 
 def _get_ncc_image(target, prediction) -> sitk.Image: