#!/usr/bin/env python3
|
"""Batch ASR example for FunASR (improved)
|
|
Features added:
|
- argparse for command-line configuration
|
- optional recursive folder scan
|
- progress reporting with simple counts
|
- safer model loading and per-file error handling
|
- configurable output file and supported extensions
|
- creates output folder if missing
|
- configurable VAD model
|
"""
|
|
import argparse
|
import sys
|
from pathlib import Path
|
from typing import List
|
|
try:
|
from funasr import AutoModel
|
from funasr.utils.postprocess_utils import rich_transcription_postprocess
|
except Exception as e:
|
print("Error importing funasr. Make sure FunASR is installed (pip install -U funasr) or you're running from the repo root.", file=sys.stderr)
|
raise
|
|
def find_audio_files(folder: Path, exts: List[str], recursive: bool) -> List[Path]:
|
if recursive:
|
return [p for p in folder.rglob("*") if p.is_file() and p.suffix.lower() in exts]
|
else:
|
return [p for p in folder.iterdir() if p.is_file() and p.suffix.lower() in exts]
|
|
def main():
|
parser = argparse.ArgumentParser(description="Batch ASR using FunASR AutoModel")
|
parser.add_argument("--input-folder", "-i", type=Path, default=Path("examples/audio_samples"),
|
help="Folder with audio files (default: examples/audio_samples)")
|
parser.add_argument("--output-file", "-o", type=Path, default=Path("examples/batch_transcriptions.txt"),
|
help="Output text file (default: examples/batch_transcriptions.txt)")
|
parser.add_argument("--model", "-m", default="paraformer-zh",
|
help="Model name (default: paraformer-zh). Examples: paraformer-zh, paraformer-en, SenseVoiceSmall")
|
parser.add_argument("--device", "-d", default="cpu", help="Device for inference (default: cpu)")
|
parser.add_argument("--recursive", "-r", action="store_true", help="Recursively search input folder for audio files")
|
parser.add_argument("--extensions", "-e", nargs="+", default=[".wav", ".mp3"], help="Accepted audio extensions (default: .wav .mp3)")
|
parser.add_argument("--vad-model", default="fsmn-vad", help="VAD model to use. Set to 'none' to disable. (default: fsmn-vad)")
|
args = parser.parse_args()
|
|
if not args.input_folder.exists():
|
print(f"Input folder does not exist: {args.input_folder}", file=sys.stderr)
|
sys.exit(2)
|
|
exts = [ext.lower() if ext.startswith('.') else f'.{ext.lower()}' for ext in args.extensions]
|
files = sorted(find_audio_files(args.input_folder, exts, args.recursive))
|
if len(files) == 0:
|
print(f"No audio files found in {args.input_folder} with extensions {exts}")
|
sys.exit(0)
|
|
vad_arg = args.vad_model if args.vad_model.lower() != 'none' else None
|
print(f"Loading model '{args.model}' on device '{args.device}'... (this may take a while)")
|
try:
|
model = AutoModel(model=args.model, vad_model=vad_arg, device=args.device)
|
except Exception as e:
|
print("Failed to load model:", e, file=sys.stderr)
|
raise
|
|
results = {}
|
total = len(files)
|
for idx, fpath in enumerate(files, start=1):
|
rel_key = str(fpath.relative_to(args.input_folder))
|
try:
|
print(f"[{idx}/{total}] Processing: {fpath}")
|
# `cache` was removed per review because it's not used by generate
|
res = model.generate(input=str(fpath), language="auto")
|
text = rich_transcription_postprocess(res[0]["text"]) if res and isinstance(res, list) and "text" in res[0] else ""
|
results[rel_key] = text
|
print(f" -> {text}")
|
except Exception as e:
|
print(f" Error processing {fpath}: {e}", file=sys.stderr)
|
results[rel_key] = f"<ERROR: {e}>"
|
|
# Ensure output directory exists
|
args.output_file.parent.mkdir(parents=True, exist_ok=True)
|
with args.output_file.open("w", encoding="utf-8") as outf:
|
for fname, transcription in results.items():
|
outf.write(f"{fname}: {transcription}\n")
|
|
print(f"\nAll transcriptions saved to {args.output_file}")
|
|
|
if __name__ == "__main__":
|
main()
|
