Add ComputeKaldiPitch

build_tools/setup_helpers/extension.py

@@ -64,6 +64,8 @@ def _get_ela(debug):
 def _get_srcs():
     srcs = [_CSRC_DIR / 'pybind.cpp']
     srcs += list(_CSRC_DIR.glob('sox/**/*.cpp'))
+    srcs += list(_CSRC_DIR.glob('kaldi/**/*.cpp'))
+    srcs += list(_CSRC_DIR.glob('kaldi/**/*.cc'))
     if _BUILD_TRANSDUCER:
         srcs += [_CSRC_DIR / 'transducer.cpp']
     return [str(path) for path in srcs]
@@ -72,6 +74,7 @@ def _get_srcs():
 def _get_include_dirs():
     dirs = [
         str(_ROOT_DIR),
+        str(_CSRC_DIR / 'kaldi'),
     ]
     if _BUILD_SOX or _BUILD_TRANSDUCER:
         dirs.append(str(_TP_INSTALL_DIR / 'include'))

test/torchaudio_unittest/assets/kaldi_test_pitch_args.json

@@ -0,0 +1,5 @@
+{"sample_frequency": 8000}
+{"sample_frequency": 8000, "frames_per_chunk": 200}
+{"sample_frequency": 8000, "frames_per_chunk": 200, "simulate_first_pass_online": true}
+{"sample_frequency": 16000}
+{"sample_frequency": 44100}

test/torchaudio_unittest/kaldi_compatibility_cpu_test.py

@@ -1,7 +1,7 @@
 import torch
 
 from torchaudio_unittest import common_utils
-from .kaldi_compatibility_impl import Kaldi
+from .kaldi_compatibility_impl import Kaldi, KaldiCPUOnly
 
 
 class TestKaldiFloat32(Kaldi, common_utils.PytorchTestCase):
@@ -12,3 +12,8 @@ class TestKaldiFloat32(Kaldi, common_utils.PytorchTestCase):
 class TestKaldiFloat64(Kaldi, common_utils.PytorchTestCase):
     dtype = torch.float64
     device = torch.device('cpu')
+
+
+class TestKaldiCPUOnly(KaldiCPUOnly, common_utils.PytorchTestCase):
+    dtype = torch.float32
+    device = torch.device('cpu')

test/torchaudio_unittest/kaldi_compatibility_impl.py

@@ -9,10 +9,13 @@
 
 from torchaudio_unittest.common_utils import (
     TestBaseMixin,
+    TempDirMixin,
     load_params,
     skipIfNoExec,
     get_asset_path,
-    load_wav
+    load_wav,
+    save_wav,
+    get_sinusoid,
 )
 
 
@@ -48,11 +51,13 @@ def _run_kaldi(command, input_type, input_value):
     return torch.from_numpy(result.copy())  # copy supresses some torch warning
 
 
-class Kaldi(TestBaseMixin):
+class KaldiTestBase(TempDirMixin, TestBaseMixin):
     def assert_equal(self, output, *, expected, rtol=None, atol=None):
         expected = expected.to(dtype=self.dtype, device=self.device)
         self.assertEqual(output, expected, rtol=rtol, atol=atol)
 
+
+class Kaldi(KaldiTestBase):
     @skipIfNoExec('apply-cmvn-sliding')
     def test_sliding_window_cmn(self):
         """sliding_window_cmn should be numerically compatible with apply-cmvn-sliding"""
@@ -101,3 +106,21 @@ def test_mfcc(self, kwargs):
         command = ['compute-mfcc-feats'] + _convert_args(**kwargs) + ['scp:-', 'ark:-']
         kaldi_result = _run_kaldi(command, 'scp', wave_file)
         self.assert_equal(result, expected=kaldi_result, rtol=1e-4, atol=1e-8)
+
+
+class KaldiCPUOnly(KaldiTestBase):
+    @parameterized.expand(load_params('kaldi_test_pitch_args.json'))
+    @skipIfNoExec('compute-kaldi-pitch-feats')
+    def test_pitch_feats(self, kwargs):
+        """compute_kaldi_pitch produces numerically compatible result with compute-kaldi-pitch-feats"""
+        sample_rate = kwargs['sample_frequency']
+        waveform = get_sinusoid(dtype='float32', sample_rate=sample_rate)
+        result = F.compute_kaldi_pitch(waveform[0], **kwargs)
+
+        waveform = get_sinusoid(dtype='int16', sample_rate=sample_rate)
+        wave_file = self.get_temp_path('test.wav')
+        save_wav(wave_file, waveform, sample_rate)
+
+        command = ['compute-kaldi-pitch-feats'] + _convert_args(**kwargs) + ['scp:-', 'ark:-']
+        kaldi_result = _run_kaldi(command, 'scp', wave_file)
+        self.assert_equal(result, expected=kaldi_result)