157 lines
3.9 KiB
Python
Executable File
157 lines
3.9 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
# Copyright 2024 Xiaomi Corp. (authors: Fangjun Kuang)
|
|
|
|
from typing import Tuple
|
|
|
|
import kaldi_native_fbank as knf
|
|
import numpy as np
|
|
import onnxruntime as ort
|
|
import soundfile as sf
|
|
|
|
"""
|
|
NodeArg(name='feats', type='tensor(float)', shape=[1, 'T', 40])
|
|
-----
|
|
NodeArg(name='logits', type='tensor(float)', shape=['Addlogits_dim_0', 1, 7535])
|
|
"""
|
|
|
|
|
|
class OnnxModel:
|
|
def __init__(
|
|
self,
|
|
filename: str,
|
|
):
|
|
session_opts = ort.SessionOptions()
|
|
session_opts.inter_op_num_threads = 1
|
|
session_opts.intra_op_num_threads = 1
|
|
|
|
self.session_opts = session_opts
|
|
|
|
self.model = ort.InferenceSession(
|
|
filename,
|
|
sess_options=self.session_opts,
|
|
providers=["CPUExecutionProvider"],
|
|
)
|
|
|
|
self.show()
|
|
|
|
def show(self):
|
|
for i in self.model.get_inputs():
|
|
print(i)
|
|
|
|
print("-----")
|
|
|
|
for i in self.model.get_outputs():
|
|
print(i)
|
|
|
|
def __call__(self, x):
|
|
"""
|
|
Args:
|
|
x: a float32 tensor of shape (N, T, C)
|
|
"""
|
|
logits = self.model.run(
|
|
[
|
|
self.model.get_outputs()[0].name,
|
|
],
|
|
{
|
|
self.model.get_inputs()[0].name: x,
|
|
},
|
|
)[0]
|
|
|
|
return logits
|
|
|
|
|
|
def load_audio(filename: str) -> Tuple[np.ndarray, int]:
|
|
data, sample_rate = sf.read(
|
|
filename,
|
|
always_2d=True,
|
|
dtype="float32",
|
|
)
|
|
data = data[:, 0] # use only the first channel
|
|
samples = np.ascontiguousarray(data)
|
|
return samples, sample_rate
|
|
|
|
|
|
def get_features(test_wav_filename):
|
|
samples, sample_rate = load_audio(test_wav_filename)
|
|
|
|
if sample_rate != 16000:
|
|
import librosa
|
|
|
|
samples = librosa.resample(samples, orig_sr=sample_rate, target_sr=16000)
|
|
sample_rate = 16000
|
|
|
|
samples *= 32768
|
|
|
|
opts = knf.MfccOptions()
|
|
# See https://github.com/Tele-AI/TeleSpeech-ASR/blob/master/mfcc_hires.conf
|
|
opts.frame_opts.dither = 0
|
|
|
|
opts.num_ceps = 40
|
|
opts.use_energy = False
|
|
|
|
opts.mel_opts.num_bins = 40
|
|
opts.mel_opts.low_freq = 40
|
|
opts.mel_opts.high_freq = -200
|
|
|
|
mfcc = knf.OnlineMfcc(opts)
|
|
mfcc.accept_waveform(16000, samples)
|
|
frames = []
|
|
for i in range(mfcc.num_frames_ready):
|
|
frames.append(mfcc.get_frame(i))
|
|
|
|
frames = np.stack(frames, axis=0)
|
|
return frames
|
|
|
|
|
|
def cmvn(features):
|
|
# See https://github.com/Tele-AI/TeleSpeech-ASR/blob/master/wenet_representation/conf/train_d2v2_ark_conformer.yaml#L70
|
|
# https://github.com/Tele-AI/TeleSpeech-ASR/blob/master/wenet_representation/wenet/dataset/dataset.py#L184
|
|
# https://github.com/Tele-AI/TeleSpeech-ASR/blob/master/wenet_representation/wenet/dataset/processor.py#L278
|
|
mean = features.mean(axis=0, keepdims=True)
|
|
std = features.std(axis=0, keepdims=True)
|
|
return (features - mean) / (std + 1e-5)
|
|
|
|
|
|
def main():
|
|
# Please download the test data from
|
|
# https://hf-mirror.com/csukuangfj/sherpa-onnx-paraformer-zh-small-2024-03-09/tree/main/test_wavs
|
|
test_wav_filename = "./3-sichuan.wav"
|
|
test_wav_filename = "./4-tianjin.wav"
|
|
test_wav_filename = "./5-henan.wav"
|
|
|
|
features = get_features(test_wav_filename)
|
|
|
|
features = cmvn(features)
|
|
|
|
features = np.expand_dims(features, axis=0) # (T, C) -> (N, T, C)
|
|
|
|
model_filename = "./model.int8.onnx"
|
|
model = OnnxModel(model_filename)
|
|
logits = model(features)
|
|
logits = logits.squeeze(axis=1) # remove batch axis
|
|
ids = logits.argmax(axis=-1)
|
|
|
|
id2token = dict()
|
|
with open("./tokens.txt", encoding="utf-8") as f:
|
|
for line in f:
|
|
t, idx = line.split()
|
|
id2token[int(idx)] = t
|
|
|
|
tokens = []
|
|
|
|
blank = 0
|
|
prev = -1
|
|
|
|
for k in ids:
|
|
if k != blank and k != prev:
|
|
tokens.append(k)
|
|
prev = k
|
|
|
|
tokens = [id2token[i] for i in tokens]
|
|
text = "".join(tokens)
|
|
print(text)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|