Spaces:
Running
Running
File size: 2,704 Bytes
d7dc3c9 | 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 | from __future__ import annotations
from functools import lru_cache
import numpy as np
import soundfile as sf
import torch
from scipy.signal import resample_poly
from transformers import AutoFeatureExtractor, AutoModelForAudioClassification
MODEL_ID = "mispeech/ced-tiny"
MODEL_REVISION = "ace276d29dd0bb3f3517b0fa8cf300738c409019"
TARGET_SAMPLE_RATE = 16000
WINDOW_SAMPLES = 10 * TARGET_SAMPLE_RATE
@lru_cache(maxsize=1)
def _load_model():
feature_extractor = AutoFeatureExtractor.from_pretrained(
MODEL_ID,
revision=MODEL_REVISION,
trust_remote_code=True,
)
model = AutoModelForAudioClassification.from_pretrained(
MODEL_ID,
revision=MODEL_REVISION,
trust_remote_code=True,
)
model.to("cpu")
model.eval()
return feature_extractor, model
def _load_audio(path: str) -> np.ndarray:
audio, sample_rate = sf.read(path, dtype="float32", always_2d=True)
if sample_rate <= 0 or audio.size == 0:
raise ValueError("The audio file is empty or has an invalid sample rate.")
if not np.isfinite(audio).all():
raise ValueError("The audio file contains non-finite samples.")
audio = audio.mean(axis=1)
if sample_rate != TARGET_SAMPLE_RATE:
divisor = int(np.gcd(sample_rate, TARGET_SAMPLE_RATE))
audio = resample_poly(
audio,
TARGET_SAMPLE_RATE // divisor,
sample_rate // divisor,
).astype(np.float32, copy=False)
return np.ascontiguousarray(audio, dtype=np.float32)
def classify_audio(path: str, top_k: int = 10) -> list[dict[str, float | str]]:
audio = _load_audio(path)
feature_extractor, model = _load_model()
window_probabilities = []
for start in range(0, len(audio), WINDOW_SAMPLES):
window = audio[start : start + WINDOW_SAMPLES]
inputs = feature_extractor(
window,
sampling_rate=TARGET_SAMPLE_RATE,
return_tensors="pt",
)
with torch.inference_mode():
logits = model(**inputs).logits[0]
window_probabilities.append(torch.sigmoid(logits))
probabilities = torch.stack(window_probabilities).mean(dim=0)
if not torch.isfinite(probabilities).all():
raise ValueError("The model returned non-finite scores.")
count = min(max(1, int(top_k)), probabilities.numel())
scores, indices = torch.topk(probabilities, k=count)
labels = model.config.id2label
return [
{
"rank": rank,
"label": str(labels.get(int(index), int(index))),
"score": round(float(score), 6),
}
for rank, (score, index) in enumerate(zip(scores, indices), start=1)
]
|