Total Harmonic Distorsion plus noise (THD+n) score in APM-QA.

modules/audio_processing/test/py_quality_assessment/quality_assessment/eval_scores.py

 # Copyright (c) 2017 The WebRTC project authors. All Rights Reserved.
 #
 # Use of this source code is governed by a BSD-style license
 # that can be found in the LICENSE file in the root of the source
 # tree. An additional intellectual property rights grant can be found
 # in the file PATENTS.  All contributing project authors may
 # be found in the AUTHORS file in the root of the source tree.
 
 """Evaluation score abstract class and implementations.
 """
 
 from __future__ import division
 import logging
 import os
 import re
 import subprocess
+import sys
+
+try:
+  import numpy as np
+except ImportError:
+  logging.critical('Cannot import the third-party Python package numpy')
+  sys.exit(1)
 
 from . import data_access
 from . import exceptions
 from . import signal_processing
 
 
 class EvaluationScore(object):
 
   NAME = None
   REGISTERED_CLASSES = {}
 
   def __init__(self, score_filename_prefix):
     self._score_filename_prefix = score_filename_prefix
+    self._input_signal_metadata = None
     self._reference_signal = None
     self._reference_signal_filepath = None
     self._tested_signal = None
     self._tested_signal_filepath = None
     self._output_filepath = None
     self._score = None
 
   @classmethod
   def RegisterClass(cls, class_to_register):
     """Registers an EvaluationScore implementation.
 
     Decorator to automatically register the classes that extend EvaluationScore.
     Example usage:
 
     @EvaluationScore.RegisterClass
     class AudioLevelScore(EvaluationScore):
       pass
     """
     cls.REGISTERED_CLASSES[class_to_register.NAME] = class_to_register
     return class_to_register
 
   @property
   def output_filepath(self):
     return self._output_filepath
 
   @property
   def score(self):
     return self._score
 
+  def SetInputSignalMetadata(self, metadata):
+    """Sets input signal metadata.
+
+    Args:
+      metadata: dict instance.
+    """
+    self._input_signal_metadata = metadata
+
   def SetReferenceSignalFilepath(self, filepath):
-    """ Sets the path to the audio track used as reference signal.
+    """Sets the path to the audio track used as reference signal.
 
     Args:
       filepath: path to the reference audio track.
     """
     self._reference_signal_filepath = filepath
 
   def SetTestedSignalFilepath(self, filepath):
-    """ Sets the path to the audio track used as test signal.
+    """Sets the path to the audio track used as test signal.
 
     Args:
       filepath: path to the test audio track.
     """
     self._tested_signal_filepath = filepath
 
   def Run(self, output_path):
     """Extracts the score for the set test data pair.
 
     Args:
@@ skipping 156 matching lines @@
         data.append(re.split(r'\t+', line))
 
     # Two rows expected (header and values).
     assert len(data) == 2, 'Cannot parse POLQA output'
     number_of_fields = len(data[0])
     assert number_of_fields == len(data[1])
 
     # Build and return a dictionary with field names (header) as keys and the
     # corresponding field values as values.
     return {data[0][index]: data[1][index] for index in range(number_of_fields)}
+
+
+@EvaluationScore.RegisterClass
+class TotalHarmonicDistorsionScore(EvaluationScore):
+  """Total harmonic distorsion plus noise score.
+
+  Total harmonic distorsion plus noise score.
+  See "https://en.wikipedia.org/wiki/Total_harmonic_distortion#THD.2BN".
+
+  Unit: -.
+  Ideal: 0.
+  Worst case: +inf
+  """
+
+  NAME = 'thd'
+
+  def __init__(self, score_filename_prefix):
+    EvaluationScore.__init__(self, score_filename_prefix)
+    self._input_frequency = None
+
+  def _Run(self, output_path):
+    # TODO(aleloi): Integrate changes made locally.
+    self._CheckInputSignal()
+
+    self._LoadTestedSignal()
+    if self._tested_signal.channels != 1:
+      raise exceptions.EvaluationScoreException(
+          'unsupported number of channels')
+    samples = signal_processing.SignalProcessingUtils.AudioSegmentToRawData(
+        self._tested_signal)
+
+    # Init.
+    num_samples = len(samples)
+    duration = len(self._tested_signal) / 1000.0
+    scaling = 2.0 / num_samples
+    max_freq = self._tested_signal.frame_rate / 2
+    f0_freq = float(self._input_frequency)
+    t = np.linspace(0, duration, num_samples)
+
+    # Analyze harmonics.
+    b_terms = []
+    n = 1
+    while f0_freq * n < max_freq:
+      x_n = np.sum(samples * np.sin(2.0 * np.pi * n * f0_freq * t)) * scaling
+      y_n = np.sum(samples * np.cos(2.0 * np.pi * n * f0_freq * t)) * scaling
+      b_terms.append(np.sqrt(x_n**2 + y_n**2))
+      n += 1
+
+    output_without_fundamental = samples - b_terms[0] * np.sin(
+        2.0 * np.pi * f0_freq * t)
+    distortion_and_noise = np.sqrt(np.sum(
+        output_without_fundamental**2) * np.pi * scaling)
+
+    # TODO(alessiob): Fix or remove if not needed.
+    # thd = np.sqrt(np.sum(b_terms[1:]**2)) / b_terms[0]
+
+    # TODO(alessiob): Check the range of |thd_plus_noise| and update the class
+    # docstring above if accordingly.
+    thd_plus_noise = distortion_and_noise / b_terms[0]
+
+    self._score = thd_plus_noise
+    self._SaveScore()
+
+  def _CheckInputSignal(self):
+    # Check input signal and get properties.
+    try:
+      if self._input_signal_metadata['signal'] != 'pure_tone':
+        raise exceptions.EvaluationScoreException(
+            'The THD score requires a pure tone as input signal')
+      self._input_frequency = self._input_signal_metadata['frequency']
+      if self._input_signal_metadata['test_data_gen_name'] != 'identity' or (
+          self._input_signal_metadata['test_data_gen_config'] != 'default'):
+        raise exceptions.EvaluationScoreException(
+            'The THD score cannot be used with any test data generator other '
+            'than "identity"')
+    except TypeError:
+      raise exceptions.EvaluationScoreException(
+          'The THD score requires an input signal with associated metadata')
+    except KeyError:
+      raise exceptions.EvaluationScoreException(
+          'Invalid input signal metadata to compute the THD score')