Add the log spectral distance metric to the conformance test script.

Author: felicialim

tests/README.md

@@ -119,7 +119,7 @@ To set up a python environment using pip
 ```
 python3 -m venv venv
 source venv/bin/activate
-pip install scipy protobuf tqdm numpy
+pip install scipy protobuf tqdm numpy librosa
 ```
 
 Run the test suite.

tests/dsp_utils.py

@@ -5,9 +5,10 @@
 import wave
 import numpy as np
 import scipy.io.wavfile as wavfile
+import librosa
 
 
-def calc_average_channel_psnr_pcm(
+def calc_per_channel_psnr_pcm(
     ref_signal: np.ndarray, signal: np.ndarray, sampwidth_bytes: int
 ):
   """Calculates the PSNR between two signals.
@@ -19,8 +20,7 @@ def calc_average_channel_psnr_pcm(
       24-bit).
 
   Returns:
-    The average PSNR in dB across all channels, or -1 if all channels are
-    identical.
+    The per channel PSNR in dB.
   """
   assert (
       sampwidth_bytes > 1
@@ -40,25 +40,78 @@ def calc_average_channel_psnr_pcm(
   for i in range(num_channels):
     mse_value = mse[i] if num_channels > 1 else mse
     if mse_value == 0:
+      psnr_list.append(np.inf)
       logging.debug("ch#%d PSNR: inf", i)
     else:
       psnr_value = 10 * math.log10(max_value**2 / mse_value)
       psnr_list.append(psnr_value)
       logging.debug("ch#%d PSNR: %f dB", i, psnr_value)
 
-  return -1 if len(psnr_list) == 0 else sum(psnr_list) / len(psnr_list)
+  return psnr_list
 
 
-def calc_average_channel_psnr_wav(ref_filepath: str, target_filepath: str):
-  """Calculates the PSNR between two WAV files.
+def calc_per_channel_lsd_pcm(ref_signal: np.ndarray,
+                             signal: np.ndarray,
+                             sampling_rate: int):
+  """Calculates the log spectral distance using Mel bins between two signals.
+
+  Args:
+    ref_signal: The reference signal as a numpy array.
+    signal: The signal to compare as a numpy array.
+    sampling rate: The sampling rate of the signals in Hz.
+
+  Returns:
+    The per channel log spectral distance in dB.
+  """
+  eps = 1e-4
+
+  # Convert to float
+  ref_signal = ref_signal / np.iinfo(ref_signal.dtype).max
+  signal = signal / np.iinfo(signal.dtype).max
+
+  lsd_list = list()
+
+  # To support mono channel
+  num_channels = 1 if ref_signal.shape[1:] == () else ref_signal.shape[1]
+  for i in range(num_channels):
+    ref_channel = ref_signal[:, i] if num_channels > 1 else ref_signal
+    signal_channel = signal[:, i] if num_channels > 1 else signal
+
+    lsd_frames = list()
+
+    # Compute mel spectrogram
+    mel_ref = librosa.feature.melspectrogram(y=ref_channel, sr=sampling_rate)
+    mel_signal = librosa.feature.melspectrogram(y=signal_channel,
+                                                sr=sampling_rate)
+
+    log_mel_ref = 10 * np.log10(mel_ref + eps)
+    log_mel_signal = 10 * np.log10(mel_signal + eps)
+
+    diff_squared = (log_mel_ref - log_mel_signal) ** 2
+
+    # Average across mel bins, which is the 0th dimension
+    lsd_per_frame = np.sqrt(np.mean(diff_squared, axis=0))
+
+    # shape: (1, num_frames) -> (num_frames,)
+    lsd_per_frame = np.squeeze(lsd_per_frame)
+
+    lsd_value = np.mean(lsd_per_frame)
+    lsd_list.append(lsd_value)
+    logging.debug('ch#d LSD: %f dB', i, lsd_value)
+
+  return lsd_list
+
+
+def calc_score_wav(ref_filepath: str, target_filepath: str, metric: str):
+  """Calculates the score between two WAV files.
 
   Args:
     ref_filepath: Path to the reference WAV file.
     target_filepath: Path to the target WAV file to compare.
+    metric: one of 'PSNR' or 'SNR'.
 
   Returns:
-    The average PSNR in dB across all channels. Or -1 if all channels are
-    identical.
+    The score in dB, averaged over all channels.
 
   Raises:
     Exception: If the wav files have different samplerate, channels, bit-depth
@@ -99,6 +152,14 @@ def calc_average_channel_psnr_wav(ref_filepath: str, target_filepath: str):
   _, ref_data = wavfile.read(ref_filepath)
   _, target_data = wavfile.read(target_filepath)
 
-  return calc_average_channel_psnr_pcm(
-      ref_data, target_data, ref_wav.getsampwidth()
-  )
+  if metric == 'PSNR':
+    scores_list = calc_per_channel_psnr_pcm(
+          ref_data, target_data, ref_wav.getsampwidth()
+    )
+  elif metric == 'LSD':
+    scores_list = calc_per_channel_lsd_pcm(ref_data, target_data,
+                                           ref_wav.getframerate())
+  else:
+    return None
+
+  return np.mean(scores_list)

tests/run_decode_and_psnr_test.py

@@ -39,12 +39,53 @@
         layout_index=1,
         reason='Extension layouts cannot be decoded.',
     ),
+    TestExclusions(
+        file_name_prefix='test_000119',
+        mix_presentation_id=42,
+        layout_index=0,
+        reason='Signal too short',
+    ),
+    TestExclusions(
+        file_name_prefix='test_000120',
+        mix_presentation_id=42,
+        layout_index=0,
+        reason='Signal too short',
+    ),
+    TestExclusions(
+        file_name_prefix='test_000122',
+        mix_presentation_id=42,
+        layout_index=0,
+        reason='Signal too short',
+    ),
+    TestExclusions(
+        file_name_prefix='test_000129',
+        mix_presentation_id=42,
+        layout_index=0,
+        reason='Signal too short',
+    ),
+    TestExclusions(
+        file_name_prefix='test_000130',
+        mix_presentation_id=42,
+        layout_index=0,
+        reason='Signal too short',
+    ),
 ]
 
 # Opus/AAC are lossy codecs, we allow a more lenient threshold for them.
-LOSSY_PSNR_THRESHOLD = 30
-LOSSLESS_PSNR_THRESHOLD = 80
-
+# Some test files use stress signals, e.g. sawtooth. Allow a more lenient
+# threshold for them.
+THRESHOLDS = {
+    # PSNR: larger is better
+    'PSNR': {
+        'lossless': {'base': 80},
+        'lossy': {'base': 30}
+    },
+    # LSD: smaller is better
+    'LSD': {
+        'lossless': {'base': 1.0},
+        'lossy': {'base': 2.5, 'stress': 10.0}
+    }
+}
 
 class ResultStatus(enum.Enum):
   SUCCESS = 1
@@ -60,16 +101,16 @@ class Result:
   mix_presentation_id: int
   sub_mix_index: int
   layout_index: int
-  psnr_score: Optional[float] = None
+  score: Optional[float] = None
   reason: Optional[str] = None
   iamfdec_command: Optional[str] = None
 
   def log(self, status: ResultStatus):
     logging.debug(
         '%s: %s >= %s for %s',
         status.name,
-        self.psnr_score,
-        self.is_lossy,
+        self.score,
+        'lossy codec' if self.is_lossy else 'lossless codec',
         self.test_prefix,
     )
     logging.debug('')
@@ -81,11 +122,14 @@ class TestSummary:
       default_factory=lambda: defaultdict(list)
   )
 
-  def print_test_summary(self, csv_summary_file=None):
+  def print_test_summary(self,
+                         csv_summary_file: str,
+                         metric: str):
     """Prints test summary to console and optionally a CSV file.
 
     Args:
       csv_summary_file: Path to CSV file to log test results.
+      metric: Name of the metric used.
     """
     logging.info('\n-----------------SUMMARY-----------------')
     for status in ResultStatus:
@@ -102,7 +146,7 @@ def print_test_summary(self, csv_summary_file=None):
             'Submix Index',
             'Layout Index',
             'Status',
-            'PSNR',
+            metric,
             'Is Lossy',
             'Reason',
             'Command',
@@ -115,7 +159,7 @@ def print_test_summary(self, csv_summary_file=None):
                 item.sub_mix_index,
                 item.layout_index,
                 status.name,
-                item.psnr_score if item.psnr_score is not None else '',
+                item.score if item.score is not None else '',
                 'lossy' if item.is_lossy else 'lossless',
                 item.reason if item.reason is not None else '',
                 item.iamfdec_command
@@ -148,15 +192,35 @@ def run_decoder(args, metadata):
   return True, cmd_str
 
 
-def run_psnr_test(args, metadata):
-  """Gets PSNR score, returns None if calculation fails.
+def get_threshold(metric: str, is_lossy: bool, is_stress_signal: bool):
+  """Gets threshold value.
+
+  Args:
+    metric: Metric name.
+    is_lossy: Whether the codec used is lossy or lossless.
+    is_stress_signal: Whether the signal used is a stress signal type.
+
+  Returns:
+    The threshold value.
+  """
+  codec_type = 'lossy' if is_lossy else 'lossless'
+  threshold_options = THRESHOLDS.get(metric, {}).get(codec_type,{})
+
+  if is_stress_signal and 'stress' in threshold_options:
+    return threshold_options['stress']
+
+  return threshold_options['base']
+
+
+def compute_metrics(args, metadata):
+  """Gets output score, returns None if calculation fails.
 
   Args:
     args: Command line arguments.
     metadata: Metadata for the test vector.
 
   Returns:
-    A tuple of (ResultStatus, reason, psnr_score).
+    A tuple of (ResultStatus, reason, score).
   """
   ref_file = os.path.join(
       args.test_file_directory, metadata.golden_wav_file_name
@@ -168,25 +232,29 @@ def run_psnr_test(args, metadata):
   assert os.path.exists(test_file), f'Test file {test_file} does not exist.'
   logging.debug('ref_file: %s', ref_file)
   logging.debug('test_file: %s', test_file)
+
   try:
-    raw_psnr_score = dsp_utils.calc_average_channel_psnr_wav(
-        ref_file, test_file
-    )
+    score = dsp_utils.calc_score_wav(ref_file, test_file, args.metric)
   except ValueError as e:
-    print(f'Failed to calculate PSNR: {e}')
-    return ResultStatus.CRASH, 'PSNR calculation failed', None
+    print(f'Failed to calculate {args.metric}: {e}')
+    return ResultStatus.CRASH, f'{args.metric} calculation failed', None
 
-  psnr_score = 100 if raw_psnr_score == -1 else raw_psnr_score
-  # Check if this PSNR is a pass or a fail, it depends on whether the test
+  # Check if this score is a pass or a fail, it depends on whether the test
   # represents a lossy or lossless codec.
-  logging.debug('psnr score: %s', psnr_score)
-  threshold = (
-      LOSSY_PSNR_THRESHOLD if metadata.is_lossy else LOSSLESS_PSNR_THRESHOLD
-  )
-  if psnr_score >= threshold:
-    return ResultStatus.SUCCESS, None, psnr_score
+  logging.debug('%s score: %s', args.metric, score)
+  threshold = get_threshold(args.metric, metadata.is_lossy,
+                            metadata.is_stress_signal)
+
+  if args.metric == 'LSD':
+    if score <= threshold:
+      return ResultStatus.SUCCESS, None, score
+    else:
+      return ResultStatus.FAILURE, 'Score above threshold.', score
   else:
-    return ResultStatus.FAILURE, 'PSNR score below threshold.', psnr_score
+    if score >= threshold:
+      return ResultStatus.SUCCESS, None, score
+    else:
+      return ResultStatus.FAILURE, 'Score below threshold.', score
 
 
 def _is_excluded(metadata, exclusions, args):
@@ -241,17 +309,17 @@ def run_tests(args, text_proto_files) -> TestSummary:
       cleanup_after_decode = (
           generated_file_is_new and not args.preserve_output_files
       )
-      status, reason, psnr_score, cmd_str = None, None, None, None
+      status, reason, score, cmd_str = None, None, None, None
       if skip_reason := _is_excluded(metadata, exclusions, args):
         # Test was intentionally excluded.
         reason = skip_reason
         status = ResultStatus.SKIPPED
       else:
         decoder_success, cmd_str = run_decoder(args, metadata)
         if decoder_success:
-          # Run the PSNR test, this could crash, or be better than or worse than
-          # the threshold PSNR.
-          status, reason, psnr_score = run_psnr_test(args, metadata)
+          # Compute the score, this could crash, or be better than or worse than
+          # the threshold score.
+          status, reason, score = compute_metrics(args, metadata)
         else:
           # Decoder crashed.
           reason = 'iamfdec crash'
@@ -264,7 +332,7 @@ def run_tests(args, text_proto_files) -> TestSummary:
           metadata.mix_presentation_id,
           metadata.sub_mix_index,
           metadata.layout_index,
-          psnr_score,
+          score,
           reason,
           iamfdec_command=cmd_str,
       )
@@ -308,6 +376,13 @@ def main():
       action='store_true',
       help='Enable testing binaural layouts.',
   )
+  parser.add_argument(
+      '-m',
+      '--metric',
+      default='PSNR',
+      choices=['PSNR', 'LSD'],
+      help='The metric to use: PSNR or LSD.',
+  )
   args = parser.parse_args()
 
   logging.basicConfig(
@@ -329,7 +404,7 @@ def main():
   text_proto_files.sort()
 
   test_summary = run_tests(args, text_proto_files)
-  test_summary.print_test_summary(args.csv_summary_file)
+  test_summary.print_test_summary(args.csv_summary_file, args.metric)
 
 
 if __name__ == '__main__':