amine@317: from unittest import TestCase
amine@157: import math
amine@157: from array import array
amine@317: from genty import genty, genty_dataset
amine@317: from auditok.util import AudioEnergyValidator
amine@157: from auditok.io import DATA_FORMAT
amine@157: 
amine@157: 
amine@157: def _sample_generator(*data_buffers):
amine@157:     """
amine@157:     Takes a list of many mono audio data buffers and makes a sample generator
amine@157:     of interleaved audio samples, one sample from each channel. The resulting
amine@157:     generator can be used to build a multichannel audio buffer.
amine@157:     >>> gen = _sample_generator("abcd", "ABCD")
amine@157:     >>> list(gen)
amine@157:     ["a", "A", "b", "B", "c", "C", "d", "D"]
amine@157:     """
amine@157:     frame_gen = zip(*data_buffers)
amine@157:     return (sample for frame in frame_gen for sample in frame)
amine@157: 
amine@157: 
amine@157: def _generate_pure_tone(
amine@157:     frequency, duration_sec=1, sampling_rate=16000, sample_width=2, volume=1e4
amine@157: ):
amine@157:     """
amine@157:     Generates a pure tone with the given frequency.
amine@157:     """
amine@157:     assert frequency <= sampling_rate / 2
amine@157:     max_value = (2 ** (sample_width * 8) // 2) - 1
amine@157:     if volume > max_value:
amine@157:         volume = max_value
amine@157:     fmt = DATA_FORMAT[sample_width]
amine@157:     total_samples = int(sampling_rate * duration_sec)
amine@157:     step = frequency / sampling_rate
amine@157:     two_pi_step = 2 * math.pi * step
amine@157:     data = array(
amine@157:         fmt,
amine@157:         (
amine@157:             int(math.sin(two_pi_step * i) * volume)
amine@157:             for i in range(total_samples)
amine@157:         ),
amine@157:     )
amine@157:     return data
amine@157: 
amine@157: 
amine@157: PURE_TONE_DICT = {
amine@157:     freq: _generate_pure_tone(freq, 1, 16000, 2) for freq in (400, 800, 1600)
amine@157: }
amine@157: PURE_TONE_DICT.update(
amine@157:     {
amine@157:         freq: _generate_pure_tone(freq, 0.1, 16000, 2)
amine@157:         for freq in (600, 1150, 2400, 7220)
amine@157:     }
amine@157: )
amine@317: 
amine@317: 
amine@317: @genty
amine@317: class TestAudioEnergyValidator(TestCase):
amine@317:     @genty_dataset(
amine@317:         mono_valid_uc_None=([350, 400], 1, None, True),
amine@317:         mono_valid_uc_any=([350, 400], 1, "any", True),
amine@317:         mono_valid_uc_0=([350, 400], 1, 0, True),
amine@317:         mono_valid_uc_mix=([350, 400], 1, "mix", True),
amine@317:         # previous cases are all the same since we have mono audio
amine@317:         mono_invalid_uc_None=([300, 300], 1, None, False),
amine@317:         stereo_valid_uc_None=([300, 400, 350, 300], 2, None, True),
amine@317:         stereo_valid_uc_any=([300, 400, 350, 300], 2, "any", True),
amine@317:         stereo_valid_uc_mix=([300, 400, 350, 300], 2, "mix", True),
amine@317:         stereo_valid_uc_avg=([300, 400, 350, 300], 2, "avg", True),
amine@317:         stereo_valid_uc_average=([300, 400, 300, 300], 2, "average", True),
amine@317:         stereo_valid_uc_mix_with_null_channel=(
amine@317:             [634, 0, 634, 0],
amine@317:             2,
amine@317:             "mix",
amine@317:             True,
amine@317:         ),
amine@317:         stereo_valid_uc_0=([320, 100, 320, 100], 2, 0, True),
amine@317:         stereo_valid_uc_1=([100, 320, 100, 320], 2, 1, True),
amine@317:         stereo_invalid_uc_None=([280, 100, 280, 100], 2, None, False),
amine@317:         stereo_invalid_uc_any=([280, 100, 280, 100], 2, "any", False),
amine@317:         stereo_invalid_uc_mix=([400, 200, 400, 200], 2, "mix", False),
amine@317:         stereo_invalid_uc_0=([300, 400, 300, 400], 2, 0, False),
amine@317:         stereo_invalid_uc_1=([400, 300, 400, 300], 2, 1, False),
amine@317:         zeros=([0, 0, 0, 0], 2, None, False),
amine@317:     )
amine@317:     def test_audio_energy_validator(
amine@317:         self, data, channels, use_channel, expected
amine@317:     ):
amine@317: 
amine@317:         data = array("h", data)
amine@317:         sample_width = 2
amine@317:         energy_threshold = 50
amine@317:         validator = AudioEnergyValidator(
amine@317:             energy_threshold, sample_width, channels, use_channel
amine@317:         )
amine@317: 
amine@317:         if expected:
amine@317:             self.assertTrue(validator.is_valid(data))
amine@317:         else:
amine@317:             self.assertFalse(validator.is_valid(data))