comparison tests/test_signal.py @ 315:5f1859160fd7

Refactor signal processing code
    - Use audioop for channel averaging and energy computation

+import unittest
 from unittest import TestCase
 from array import array as array_
 from genty import genty, genty_dataset
 import numpy as np
 from auditok import signal as signal_

 class TestSignal(TestCase):
     def setUp(self):
         self.data = b"012345679ABC"
         self.numpy_fmt = {"b": np.int8, "h": np.int16, "i": np.int32}
 
+
     @genty_dataset(
-        int8=("b", [48, 49, 50, 51, 52, 53, 54, 55, 57, 65, 66, 67]),
-        int16=("h", [12592, 13106, 13620, 14134, 16697, 17218]),
-        int32=("i", [858927408, 926299444, 1128415545]),
+        int8_mono=(1, [48, 49, 50, 51, 52, 53, 54, 55, 57, 65, 66, 67]),
+        int16_mono=(2, [12592, 13106, 13620, 14134, 16697, 17218]),
+        int32_mono=(4, [858927408, 926299444, 1128415545]),
+        int8_stereo=(1,  [[48, 50, 52, 54, 57, 66], [49, 51, 53, 55, 65, 67]]),
+        int16_stereo=(2, [[12592, 13620, 16697], [13106, 14134, 17218]]),
+        int32_3channel=(4, [[858927408], [926299444], [1128415545]]),
     )
-    def test_to_array(self, fmt, expected):
-        resutl = signal_.to_array(self.data, fmt)
-        expected = array_(fmt, expected)
+    def test_to_array(self, sample_width, expected):
+        if isinstance(expected[0], list):
+            channels = len(expected)
+            expected = [array_(signal_.FORMAT[sample_width], xi) for xi in expected]
+        else:
+            channels = 1
+            expected = array_(signal_.FORMAT[sample_width], expected)
+        resutl = signal_.to_array(self.data, sample_width, channels)
+        resutl_numpy = signal_numpy.to_array(self.data, sample_width, channels)
         self.assertEqual(resutl, expected)
-        resutl_numpy = signal_numpy.to_array(self.data, self.numpy_fmt[fmt])
-        self.assertTrue(all(resutl_numpy == expected))
+        self.assertTrue((resutl_numpy == np.asarray(expected)).all())
         self.assertEqual(resutl_numpy.dtype, np.float64)
 
     @genty_dataset(
         int8_1channel_select_0=(
             "b",

     def test_extract_single_channel(self, fmt, channels, selected, expected):
         resutl = signal_.extract_single_channel(
             self.data, fmt, channels, selected
         )
         expected = array_(fmt, expected)
+        expected_numpy_fmt = self.numpy_fmt[fmt]
         self.assertEqual(resutl, expected)
         resutl_numpy = signal_numpy.extract_single_channel(
             self.data, self.numpy_fmt[fmt], channels, selected
         )
         self.assertTrue(all(resutl_numpy == expected))
-        self.assertEqual(resutl_numpy.dtype, np.float64)
+        self.assertEqual(resutl_numpy.dtype, expected_numpy_fmt)
 
     @genty_dataset(
         int8_2channel=("b", 2, [48, 50, 52, 54, 61, 66]),
         int8_4channel=("b", 4, [50, 54, 64]),
         int16_1channel=("h", 1, [12592, 13106, 13620, 14134, 16697, 17218]),

         int32_3channel=("i", 3, [971214132]),
     )
     def test_average_channels(self, fmt, channels, expected):
         resutl = signal_.average_channels(self.data, fmt, channels)
         expected = array_(fmt, expected)
+        expected_numpy_fmt = self.numpy_fmt[fmt]
         self.assertEqual(resutl, expected)
         resutl_numpy = signal_numpy.average_channels(
             self.data, self.numpy_fmt[fmt], channels
         )
         self.assertTrue(all(resutl_numpy == expected))
-        self.assertEqual(resutl_numpy.dtype, np.float64)
+        self.assertEqual(resutl_numpy.dtype, expected_numpy_fmt)
 
     @genty_dataset(
         int8_1channel=(
             "b",
             1,

         int32_3channel=("i", 3, [[858927408], [926299444], [1128415545]]),
     )
     def test_separate_channels(self, fmt, channels, expected):
         resutl = signal_.separate_channels(self.data, fmt, channels)
         expected = [array_(fmt, exp) for exp in expected]
+        expected_numpy_fmt = self.numpy_fmt[fmt]
         self.assertEqual(resutl, expected)
 
         resutl_numpy = signal_numpy.separate_channels(
             self.data, self.numpy_fmt[fmt], channels
         )
         self.assertTrue((resutl_numpy == expected).all())
-        self.assertEqual(resutl_numpy.dtype, np.float64)
+        self.assertEqual(resutl_numpy.dtype, expected_numpy_fmt)
 
     @genty_dataset(
-        simple=([300, 320, 400, 600], 52.506639194632434),
-        zero=([0], -200),
-        zeros=([0, 0, 0], -200),
+        simple=([300, 320, 400, 600], 2, 52.50624901923348),
+        zero=([0], 2, -200),
+        zeros=([0, 0, 0], 2, -200),
     )
-    def test_calculate_energy_single_channel(self, x, expected):
-        energy = signal_.calculate_energy_single_channel(x)
+    def test_calculate_energy_single_channel(self, x, sample_width, expected):
+        x = array_(signal_.FORMAT[sample_width], x)
+        energy = signal_.calculate_energy_single_channel(x, sample_width)
         self.assertEqual(energy, expected)
-        energy = signal_numpy.calculate_energy_single_channel(x)
+        energy = signal_numpy.calculate_energy_single_channel(x, sample_width)
         self.assertEqual(energy, expected)
+
 
     @genty_dataset(
         min_=(
             [[300, 320, 400, 600], [150, 160, 200, 300]],
+            2,
             min,
-            46.48603928135281,
+            46.485649105953854,
         ),
         max_=(
             [[300, 320, 400, 600], [150, 160, 200, 300]],
+            2,
             max,
-            52.506639194632434,
+            52.50624901923348,
         ),
     )
-    def test_calculate_energy_multichannel(self, x, aggregation_fn, expected):
-        energy = signal_.calculate_energy_multichannel(x, aggregation_fn)
+    def test_calculate_energy_multichannel(self, x, sample_width, aggregation_fn, expected):
+        x = [array_(signal_.FORMAT[sample_width], xi) for xi in x]
+        energy = signal_.calculate_energy_multichannel(x, sample_width, aggregation_fn)
         self.assertEqual(energy, expected)
 
-        energy = signal_numpy.calculate_energy_multichannel(x, aggregation_fn)
+        energy = signal_numpy.calculate_energy_multichannel(x, sample_width, aggregation_fn)
         self.assertEqual(energy, expected)
+
+if __name__ == "__main__":
+    unittest.main()