1# Copyright (c) 2017 The WebRTC project authors. All Rights Reserved.
2#
3# Use of this source code is governed by a BSD-style license
4# that can be found in the LICENSE file in the root of the source
5# tree. An additional intellectual property rights grant can be found
6# in the file PATENTS.  All contributing project authors may
7# be found in the AUTHORS file in the root of the source tree.
8
9"""Unit tests for the signal_processing module.
10"""
11
12import unittest
13
14import numpy as np
15import pydub
16
17from . import exceptions
18from . import signal_processing
19
20
21class TestSignalProcessing(unittest.TestCase):
22  """Unit tests for the signal_processing module.
23  """
24
25  def testMixSignals(self):
26    # Generate a template signal with which white noise can be generated.
27    silence = pydub.AudioSegment.silent(duration=1000, frame_rate=48000)
28
29    # Generate two distinct AudioSegment instances with 1 second of white noise.
30    signal = signal_processing.SignalProcessingUtils.GenerateWhiteNoise(
31        silence)
32    noise = signal_processing.SignalProcessingUtils.GenerateWhiteNoise(
33        silence)
34
35    # Extract samples.
36    signal_samples = signal.get_array_of_samples()
37    noise_samples = noise.get_array_of_samples()
38
39    # Test target SNR -Inf (noise expected).
40    mix_neg_inf = signal_processing.SignalProcessingUtils.MixSignals(
41        signal, noise, -np.Inf)
42    self.assertTrue(len(noise), len(mix_neg_inf))  # Check duration.
43    mix_neg_inf_samples = mix_neg_inf.get_array_of_samples()
44    self.assertTrue(  # Check samples.
45        all([x == y for x, y in zip(noise_samples, mix_neg_inf_samples)]))
46
47    # Test target SNR 0.0 (different data expected).
48    mix_0 = signal_processing.SignalProcessingUtils.MixSignals(
49        signal, noise, 0.0)
50    self.assertTrue(len(signal), len(mix_0))  # Check duration.
51    self.assertTrue(len(noise), len(mix_0))
52    mix_0_samples = mix_0.get_array_of_samples()
53    self.assertTrue(
54        any([x != y for x, y in zip(signal_samples, mix_0_samples)]))
55    self.assertTrue(
56        any([x != y for x, y in zip(noise_samples, mix_0_samples)]))
57
58    # Test target SNR +Inf (signal expected).
59    mix_pos_inf = signal_processing.SignalProcessingUtils.MixSignals(
60        signal, noise, np.Inf)
61    self.assertTrue(len(signal), len(mix_pos_inf))  # Check duration.
62    mix_pos_inf_samples = mix_pos_inf.get_array_of_samples()
63    self.assertTrue(  # Check samples.
64        all([x == y for x, y in zip(signal_samples, mix_pos_inf_samples)]))
65
66  def testMixSignalsMinInfPower(self):
67    silence = pydub.AudioSegment.silent(duration=1000, frame_rate=48000)
68    signal = signal_processing.SignalProcessingUtils.GenerateWhiteNoise(
69        silence)
70
71    with self.assertRaises(exceptions.SignalProcessingException):
72      _ = signal_processing.SignalProcessingUtils.MixSignals(
73          signal, silence, 0.0)
74
75    with self.assertRaises(exceptions.SignalProcessingException):
76      _ = signal_processing.SignalProcessingUtils.MixSignals(
77          silence, signal, 0.0)
78
79  def testMixSignalNoiseDifferentLengths(self):
80    # Test signals.
81    shorter = signal_processing.SignalProcessingUtils.GenerateWhiteNoise(
82        pydub.AudioSegment.silent(duration=1000, frame_rate=8000))
83    longer = signal_processing.SignalProcessingUtils.GenerateWhiteNoise(
84        pydub.AudioSegment.silent(duration=2000, frame_rate=8000))
85
86    # When the signal is shorter than the noise, the mix length always equals
87    # that of the signal regardless of whether padding is applied.
88    # No noise padding, length of signal less than that of noise.
89    mix = signal_processing.SignalProcessingUtils.MixSignals(
90      signal=shorter,
91      noise=longer,
92      pad_noise=signal_processing.SignalProcessingUtils.MixPadding.NO_PADDING)
93    self.assertEqual(len(shorter), len(mix))
94    # With noise padding, length of signal less than that of noise.
95    mix = signal_processing.SignalProcessingUtils.MixSignals(
96      signal=shorter,
97      noise=longer,
98      pad_noise=signal_processing.SignalProcessingUtils.MixPadding.ZERO_PADDING)
99    self.assertEqual(len(shorter), len(mix))
100
101    # When the signal is longer than the noise, the mix length depends on
102    # whether padding is applied.
103    # No noise padding, length of signal greater than that of noise.
104    mix = signal_processing.SignalProcessingUtils.MixSignals(
105      signal=longer,
106      noise=shorter,
107      pad_noise=signal_processing.SignalProcessingUtils.MixPadding.NO_PADDING)
108    self.assertEqual(len(shorter), len(mix))
109    # With noise padding, length of signal greater than that of noise.
110    mix = signal_processing.SignalProcessingUtils.MixSignals(
111      signal=longer,
112      noise=shorter,
113      pad_noise=signal_processing.SignalProcessingUtils.MixPadding.ZERO_PADDING)
114    self.assertEqual(len(longer), len(mix))
115
116  def testMixSignalNoisePaddingTypes(self):
117    # Test signals.
118    shorter = signal_processing.SignalProcessingUtils.GenerateWhiteNoise(
119        pydub.AudioSegment.silent(duration=1000, frame_rate=8000))
120    longer = signal_processing.SignalProcessingUtils.GeneratePureTone(
121        pydub.AudioSegment.silent(duration=2000, frame_rate=8000), 440.0)
122
123    # Zero padding: expect pure tone only in 1-2s.
124    mix_zero_pad = signal_processing.SignalProcessingUtils.MixSignals(
125      signal=longer,
126      noise=shorter,
127      target_snr=-6,
128      pad_noise=signal_processing.SignalProcessingUtils.MixPadding.ZERO_PADDING)
129
130    # Loop: expect pure tone plus noise in 1-2s.
131    mix_loop = signal_processing.SignalProcessingUtils.MixSignals(
132      signal=longer,
133      noise=shorter,
134      target_snr=-6,
135      pad_noise=signal_processing.SignalProcessingUtils.MixPadding.LOOP)
136
137    def Energy(signal):
138      samples = signal_processing.SignalProcessingUtils.AudioSegmentToRawData(
139          signal).astype(np.float32)
140      return np.sum(samples * samples)
141
142    e_mix_zero_pad = Energy(mix_zero_pad[-1000:])
143    e_mix_loop = Energy(mix_loop[-1000:])
144    self.assertLess(0, e_mix_zero_pad)
145    self.assertLess(e_mix_zero_pad, e_mix_loop)
146
147  def testMixSignalSnr(self):
148    # Test signals.
149    tone_low = signal_processing.SignalProcessingUtils.GeneratePureTone(
150        pydub.AudioSegment.silent(duration=64, frame_rate=8000), 250.0)
151    tone_high = signal_processing.SignalProcessingUtils.GeneratePureTone(
152        pydub.AudioSegment.silent(duration=64, frame_rate=8000), 3000.0)
153
154    def ToneAmplitudes(mix):
155      """Returns the amplitude of the coefficients #16 and #192, which
156         correspond to the tones at 250 and 3k Hz respectively."""
157      mix_fft = np.absolute(signal_processing.SignalProcessingUtils.Fft(mix))
158      return mix_fft[16], mix_fft[192]
159
160    mix = signal_processing.SignalProcessingUtils.MixSignals(
161      signal=tone_low,
162      noise=tone_high,
163      target_snr=-6)
164    ampl_low, ampl_high = ToneAmplitudes(mix)
165    self.assertLess(ampl_low, ampl_high)
166
167    mix = signal_processing.SignalProcessingUtils.MixSignals(
168      signal=tone_high,
169      noise=tone_low,
170      target_snr=-6)
171    ampl_low, ampl_high = ToneAmplitudes(mix)
172    self.assertLess(ampl_high, ampl_low)
173
174    mix = signal_processing.SignalProcessingUtils.MixSignals(
175      signal=tone_low,
176      noise=tone_high,
177      target_snr=6)
178    ampl_low, ampl_high = ToneAmplitudes(mix)
179    self.assertLess(ampl_high, ampl_low)
180
181    mix = signal_processing.SignalProcessingUtils.MixSignals(
182      signal=tone_high,
183      noise=tone_low,
184      target_snr=6)
185    ampl_low, ampl_high = ToneAmplitudes(mix)
186    self.assertLess(ampl_low, ampl_high)
187