csukuangfj/paraformer-onnxruntime-python-example

1#!/usr/bin/env python32 3# Copyright (c)  2023  Xiaomi Corporation4# Author: Fangjun Kuang5 6import kaldi_native_fbank as knf7import librosa8import numpy as np9import onnxruntime10 11 12def load_cmvn():13    neg_mean = None14    inv_std = None15 16    with open("am.mvn") as f:17        for line in f:18            if not line.startswith("<LearnRateCoef>"):19                continue20            t = line.split()[3:-1]21            t = list(map(lambda x: float(x), t))22 23            if neg_mean is None:24                neg_mean = np.array(t, dtype=np.float32)25            else:26                inv_std = np.array(t, dtype=np.float32)27 28    return neg_mean, inv_std29 30 31def compute_feat():32    sample_rate = 1600033    samples, _ = librosa.load("2.wav", sr=sample_rate)34    opts = knf.FbankOptions()35    opts.frame_opts.dither = 036    opts.frame_opts.snip_edges = False37    opts.frame_opts.samp_freq = sample_rate38    opts.mel_opts.num_bins = 8039 40    online_fbank = knf.OnlineFbank(opts)41    online_fbank.accept_waveform(sample_rate, (samples * 32768).tolist())42    online_fbank.input_finished()43 44    features = np.stack(45        [online_fbank.get_frame(i) for i in range(online_fbank.num_frames_ready)]46    )47    assert features.data.contiguous is True48    assert features.dtype == np.float32, features.dtype49 50    window_size = 7  # lfr_m51    window_shift = 6  # lfr_n52 53    T = (features.shape[0] - window_size) // window_shift + 154    features = np.lib.stride_tricks.as_strided(55        features,56        shape=(T, features.shape[1] * window_size),57        strides=((window_shift * features.shape[1]) * 4, 4),58    )59    neg_mean, inv_std = load_cmvn()60    features = (features + neg_mean) * inv_std61    return features62 63 64# tokens.txt in paraformer has only one column65# while it has two columns ins sherpa-onnx.66# This function can handle tokens.txt from both paraformer and sherpa-onnx67def load_tokens():68    ans = dict()69    i = 070    with open("tokens.txt", encoding="utf-8") as f:71        for line in f:72            ans[i] = line.strip().split()[0]73            i += 174    return ans75 76 77def main():78    features = compute_feat()79    features = np.expand_dims(features, axis=0)80    features_length = np.array([features.shape[1]], dtype=np.int32)81 82    session_opts = onnxruntime.SessionOptions()83    session_opts.log_severity_level = 3  # error level84    sess = onnxruntime.InferenceSession("model.onnx", session_opts)85 86    inputs = {87        "speech": features,88        "speech_lengths": features_length,89    }90    output_names = ["logits"]91 92    try:93        outputs = sess.run(output_names, input_feed=inputs)94    except ONNXRuntimeError:95        print("Input wav is silence or noise")96        return97 98    log_probs = outputs[0].squeeze(0)99    y = log_probs.argmax(axis=-1)100 101    tokens = load_tokens()102    text = "".join([tokens[i] for i in y if i not in (0, 2)])103    print(text)104 105 106if __name__ == "__main__":107    main()108