kaldifeat/kaldifeat/python/tests/test_spectrogram.py
Fangjun Kuang 72aa5eab2b
Add support for pickling. (#17)
* Add pickle support to FrameExtractionOptions.

* Add pickle support to MelBanksOptions.

* Add pickle support.

* Update README to mention other projects.

* Fix style issues.
2021-11-04 11:55:22 +08:00

72 lines
2.1 KiB
Python
Executable File

#!/usr/bin/env python3
# Copyright 2021 Xiaomi Corporation (authors: Fangjun Kuang)
import pickle
from pathlib import Path
from utils import get_devices, read_ark_txt, read_wave
import kaldifeat
cur_dir = Path(__file__).resolve().parent
def test_spectrogram_default():
print("=====test_spectrogram_default=====")
for device in get_devices():
print("device", device)
opts = kaldifeat.SpectrogramOptions()
opts.device = device
opts.frame_opts.dither = 0
spectrogram = kaldifeat.Spectrogram(opts)
filename = cur_dir / "test_data/test.wav"
wave = read_wave(filename).to(opts.device)
features = spectrogram(wave)
gt = read_ark_txt(cur_dir / "test_data/test-spectrogram.txt")
# assert torch.allclose(features.cpu(), gt, atol=1.1)
print(features[1, 145:148], gt[1, 145:148]) # they are different
def test_spectrogram_no_snip_edges():
print("=====test_spectrogram_no_snip_edges=====")
for device in get_devices():
print("device", device)
opts = kaldifeat.SpectrogramOptions()
opts.device = device
opts.frame_opts.dither = 0
opts.frame_opts.snip_edges = False
spectrogram = kaldifeat.Spectrogram(opts)
filename = cur_dir / "test_data/test.wav"
wave = read_wave(filename).to(opts.device)
features = spectrogram(wave)
gt = read_ark_txt(
cur_dir / "test_data/test-spectrogram-no-snip-edges.txt"
)
# assert torch.allclose(features.cpu(), gt, atol=1.5)
print(features[1, 145:148], gt[1, 145:148]) # they are different
def test_pickle():
for device in get_devices():
opts = kaldifeat.SpectrogramOptions()
opts.device = device
opts.frame_opts.dither = 0
opts.frame_opts.snip_edges = False
spec = kaldifeat.Spectrogram(opts)
data = pickle.dumps(spec)
spec2 = pickle.loads(data)
assert str(spec.opts) == str(spec2.opts)
if __name__ == "__main__":
test_spectrogram_default()
test_spectrogram_no_snip_edges()
test_pickle()