135 lines
4.3 KiB
Python
135 lines
4.3 KiB
Python
from pathlib import Path
|
|
import os
|
|
import sys
|
|
import threading
|
|
import logging
|
|
import wave
|
|
import json
|
|
import subprocess
|
|
import shutil
|
|
|
|
try:
|
|
from vosk import Model, KaldiRecognizer
|
|
except Exception:
|
|
Model = None
|
|
KaldiRecognizer = None
|
|
|
|
|
|
class FFmpegConverter:
|
|
"""Utility to convert audio files using ffmpeg."""
|
|
|
|
@staticmethod
|
|
def convert_mp3_to_wav(src: Path, dest: Path) -> None:
|
|
cmd = [
|
|
"ffmpeg",
|
|
"-y",
|
|
"-i",
|
|
str(src),
|
|
"-ac",
|
|
"1",
|
|
"-ar",
|
|
"16000",
|
|
"-vn",
|
|
"-f",
|
|
"wav",
|
|
str(dest),
|
|
]
|
|
subprocess.run(cmd, check=True)
|
|
|
|
|
|
class VoskModelManager:
|
|
"""Loads a Vosk model and provides transcription methods.
|
|
|
|
Loading captures model stderr into Python logs to show progress.
|
|
"""
|
|
|
|
def __init__(self, model_path: Path):
|
|
self.logger = logging.getLogger("vosk.model_manager")
|
|
if Model is None or KaldiRecognizer is None:
|
|
raise RuntimeError("vosk is not installed in runtime")
|
|
self.model = self._load_model_with_progress(model_path)
|
|
|
|
def _load_model_with_progress(self, path: Path):
|
|
logger = logging.getLogger("vosk.loader")
|
|
r_fd, w_fd = os.pipe()
|
|
saved_stderr_fd = os.dup(sys.stderr.fileno())
|
|
os.dup2(w_fd, sys.stderr.fileno())
|
|
|
|
def _reader(fd):
|
|
with os.fdopen(fd, "rb") as fh:
|
|
for raw in iter(fh.readline, b""):
|
|
try:
|
|
logger.info(raw.decode(errors="ignore").rstrip())
|
|
except Exception:
|
|
pass
|
|
|
|
reader_thread = threading.Thread(target=_reader, args=(r_fd,), daemon=True)
|
|
reader_thread.start()
|
|
|
|
try:
|
|
model = Model(str(path))
|
|
finally:
|
|
try:
|
|
os.dup2(saved_stderr_fd, sys.stderr.fileno())
|
|
except Exception:
|
|
pass
|
|
try:
|
|
os.close(saved_stderr_fd)
|
|
except Exception:
|
|
pass
|
|
reader_thread.join(timeout=2)
|
|
|
|
return model
|
|
|
|
def transcribe_wav(self, wav_path: Path) -> str:
|
|
if not wav_path.exists():
|
|
raise FileNotFoundError(wav_path)
|
|
try:
|
|
with wave.open(str(wav_path), "rb") as wf:
|
|
if wf.getnchannels() != 1:
|
|
raise ValueError("Audio must be mono (1 channel)")
|
|
recognizer = KaldiRecognizer(self.model, wf.getframerate())
|
|
try:
|
|
recognizer.SetWords(True)
|
|
except Exception:
|
|
# older vosk bindings may not expose SetWords; ignore if absent
|
|
pass
|
|
while True:
|
|
data = wf.readframes(4000)
|
|
if not data:
|
|
break
|
|
recognizer.AcceptWaveform(data)
|
|
result = json.loads(recognizer.FinalResult())
|
|
return result.get("text", "")
|
|
except wave.Error as e:
|
|
raise ValueError(f"Invalid WAV file: {e}")
|
|
|
|
def transcribe_wav_with_words(self, wav_path: Path):
|
|
"""Transcribe WAV and return dict with 'text' and 'words' (word-level timestamps).
|
|
|
|
Returns: { 'text': str, 'words': [{'word': str, 'start': float, 'end': float}, ...] }
|
|
"""
|
|
if not wav_path.exists():
|
|
raise FileNotFoundError(wav_path)
|
|
try:
|
|
with wave.open(str(wav_path), "rb") as wf:
|
|
if wf.getnchannels() != 1:
|
|
raise ValueError("Audio must be mono (1 channel)")
|
|
recognizer = KaldiRecognizer(self.model, wf.getframerate())
|
|
try:
|
|
recognizer.SetWords(True)
|
|
except Exception:
|
|
pass
|
|
while True:
|
|
data = wf.readframes(4000)
|
|
if not data:
|
|
break
|
|
recognizer.AcceptWaveform(data)
|
|
res = json.loads(recognizer.FinalResult())
|
|
words = res.get("result", [])
|
|
# words are dictionaries with word, start, end
|
|
text = res.get("text", "")
|
|
return {"text": text, "words": words}
|
|
except wave.Error as e:
|
|
raise ValueError(f"Invalid WAV file: {e}")
|