106 lines
3.2 KiB
Python
106 lines
3.2 KiB
Python
# Copyright 2026 Umbra Tools (umbratools.com)
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class QualityChecker:
|
|
def __init__(self, config: Any = None) -> None:
|
|
self.config = config
|
|
self._whisper = None
|
|
self._mos_model = None
|
|
|
|
def _ensure_whisper(self) -> None:
|
|
if self._whisper is not None:
|
|
return
|
|
try:
|
|
import whisper
|
|
|
|
self._whisper = whisper.load_model("base")
|
|
logger.info("Whisper model loaded for WER checks")
|
|
except Exception as e:
|
|
logger.warning("Could not load Whisper: %s", e)
|
|
|
|
def _ensure_mos(self) -> None:
|
|
if self._mos_model is not None:
|
|
return
|
|
try:
|
|
import speechmos
|
|
|
|
self._mos_model = speechmos.jnt_fer_2023
|
|
logger.info("UTMOS model loaded")
|
|
except Exception as e:
|
|
logger.warning("Could not load UTMOS: %s", e)
|
|
|
|
def measure_mos(self, audio_path: Path) -> float | None:
|
|
self._ensure_mos()
|
|
if self._mos_model is None:
|
|
return None
|
|
try:
|
|
import numpy as np
|
|
import soundfile as sf
|
|
|
|
audio, sr = sf.read(str(audio_path))
|
|
if audio.ndim > 1:
|
|
audio = audio.mean(axis=1)
|
|
|
|
score = self._mos_model.run(audio, sr)
|
|
return float(score) if score else None
|
|
except Exception as e:
|
|
logger.error("MOS measurement failed: %s", e)
|
|
return None
|
|
|
|
def measure_wer(self, audio_path: Path, reference_text: str) -> float | None:
|
|
self._ensure_whisper()
|
|
if self._whisper is None:
|
|
return None
|
|
try:
|
|
result = self._whisper.transcribe(str(audio_path))
|
|
transcribed = result["text"].strip().lower()
|
|
reference = reference_text.strip().lower()
|
|
|
|
ref_words = reference.split()
|
|
trans_words = transcribed.split()
|
|
|
|
distance = self._levenshtein(ref_words, trans_words)
|
|
wer = distance / max(len(ref_words), 1)
|
|
return min(wer, 1.0)
|
|
except Exception as e:
|
|
logger.error("WER measurement failed: %s", e)
|
|
return None
|
|
|
|
def check_audio_quality(self, audio_path: Path) -> dict[str, float | None]:
|
|
mos = self.measure_mos(audio_path)
|
|
return {
|
|
"mos_score": mos,
|
|
"wer_score": None,
|
|
"speaker_similarity": None,
|
|
"artifact_score": None,
|
|
}
|
|
|
|
@staticmethod
|
|
def _levenshtein(s1: list[str], s2: list[str]) -> int:
|
|
if len(s1) < len(s2):
|
|
return QualityChecker._levenshtein(s2, s1)
|
|
|
|
if len(s2) == 0:
|
|
return len(s1)
|
|
|
|
previous_row = range(len(s2) + 1)
|
|
for i, c1 in enumerate(s1):
|
|
current_row = [i + 1]
|
|
for j, c2 in enumerate(s2):
|
|
insertions = previous_row[j + 1] + 1
|
|
deletions = current_row[j] + 1
|
|
substitutions = previous_row[j] + (c1 != c2)
|
|
current_row.append(min(insertions, deletions, substitutions))
|
|
previous_row = current_row
|
|
|
|
return previous_row[-1]
|