| |
| """LavaSR v2 AX650 demo: enhance/upsample a wav on the board. |
| |
| Usage: |
| python3 demo.py input.wav -o output.wav |
| python3 demo.py input.wav -o output.wav --no-denoise |
| """ |
| import argparse |
| import os |
| import sys |
| sys.dont_write_bytecode = True |
| import time |
|
|
| import numpy as np |
|
|
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) |
| from lavasr_sdk import LavaSR, read_wav, write_wav |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("input") |
| ap.add_argument("-o", "--output", default=None) |
| ap.add_argument("--models", default=os.path.join(os.path.dirname(__file__), "..", "models")) |
| ap.add_argument("--no-denoise", action="store_true", help="skip the UL-UNAS denoiser") |
| ap.add_argument("--cutoff", type=int, default=None, |
| help="LR crossover Hz (default: min(input_sr,16000)/2)") |
| args = ap.parse_args() |
|
|
| model_dir = os.path.abspath(args.models) |
| model = LavaSR( |
| os.path.join(model_dir, "den_encoder.axmodel"), |
| os.path.join(model_dir, "den_decoder.axmodel"), |
| os.path.join(model_dir, "enhancer.axmodel"), |
| model_config=os.path.join(model_dir, "model_config.json"), |
| cutoff=args.cutoff, |
| ) |
| x, sr = read_wav(args.input) |
| t0 = time.time() |
| y, osr = model.enhance(x, sr, denoise=not args.no_denoise) |
| dt = time.time() - t0 |
| out = args.output or os.path.splitext(args.input)[0] + "_enhanced.wav" |
| write_wav(out, y, osr) |
| dur = len(x) / float(sr) |
| print(f"LavaSR v2: {sr} Hz -> {osr} Hz, {dur:.2f}s audio") |
| print(f" denoise : {not args.no_denoise}") |
| print(f" time : {dt:.3f}s RTF {dt / max(dur, 1e-9):.4f}") |
| print(f" output : {out}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
