File size: 3,601 Bytes
19eb87a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 | #!/usr/bin/env python3
"""Dynamic ASR Pocket (300M, CPU): offline transcription with onnxruntime and greedy CTC decoding.
pip install onnxruntime numpy soundfile scipy
python transcribe_pocket.py clip.wav [more.wav ...] # int8 model (default)
python transcribe_pocket.py clip.wav --model model_fp32.onnx # fp32 model
Input: any wav; it is mixed to mono and resampled to 16 kHz. Clips longer than ~20 s are cut at quiet points.
Model files are looked up next to this script: model_int8.onnx (or model_fp32.onnx) and vocab.json.
Code licence: Apache-2.0. Model weights: CC-BY-SA-4.0.
"""
import argparse, json, os, time
import numpy as np
HERE = os.path.dirname(os.path.abspath(__file__))
SR = 16000
class PocketASR:
def __init__(self, model=None, vocab=None, threads=4):
import onnxruntime as ort
so = ort.SessionOptions()
so.intra_op_num_threads = threads
so.inter_op_num_threads = 1
so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
self.sess = ort.InferenceSession(model or os.path.join(HERE, "model_int8.onnx"), so,
providers=["CPUExecutionProvider"])
v = json.load(open(vocab or os.path.join(HERE, "vocab.json"), encoding="utf-8"))
self.pieces = v["pieces"] # index 0 = CTC blank; 1-3 = special tokens (never emitted)
def logprobs(self, wav):
# input "wav": float32 [1, T] at 16 kHz -> output "logprobs": [1, frames, vocab]
return self.sess.run(None, {"wav": np.asarray(wav, dtype=np.float32)[None]})[0][0]
def greedy(self, lp):
out, prev = [], -1
for p in lp.argmax(-1):
if p != prev and p > 3:
out.append(self.pieces[p])
prev = p
return " ".join("".join(out).split())
def transcribe(self, wav):
if len(wav) < SR // 10:
return ""
return self.greedy(self.logprobs(wav))
def load_wav(path):
import soundfile as sf
y, sr = sf.read(path, dtype="float32", always_2d=True)
y = y.mean(1)
if sr != SR:
from math import gcd
from scipy.signal import resample_poly
g = gcd(sr, SR)
y = resample_poly(y, SR // g, sr // g).astype(np.float32)
return y
def transcribe_long(asr, y, win=20.0):
"""Audio over ~20 s: cut at the quietest 100 ms in the last 4 s of each window, transcribe pieces, join."""
if len(y) <= win * SR:
return asr.transcribe(y)
out, s = [], 0
while s < len(y):
e = min(len(y), s + int(win * SR))
if e < len(y):
lo = s + int((win - 4) * SR)
fr = y[lo:e][: (e - lo) // 1600 * 1600].reshape(-1, 1600)
e = lo + int(np.argmin((fr ** 2).mean(1))) * 1600 + 800
out.append(asr.transcribe(y[s:e]))
s = e
return " ".join(t for t in out if t)
if __name__ == "__main__":
ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
ap.add_argument("wav", nargs="+")
ap.add_argument("--model", default=None, help="path to model_int8.onnx (default) or model_fp32.onnx")
ap.add_argument("--vocab", default=None)
ap.add_argument("--threads", type=int, default=4)
a = ap.parse_args()
asr = PocketASR(a.model, a.vocab, a.threads)
for p in a.wav:
y = load_wav(p)
t = time.perf_counter()
txt = transcribe_long(asr, y)
dt = time.perf_counter() - t
print(f"{p}\t{txt}\t[{len(y) / SR:.1f} s audio, {dt * 1000:.0f} ms, RTF {dt / max(len(y) / SR, 1e-3):.3f}]")
|