File size: 5,243 Bytes
c3eb7ea | 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 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 | import functools
import numpy as np
import torch
from .config import AST_MODEL, EMOTION_CAT_MODEL, EMOTION_DIM_MODEL, SEGMENTATION_MODEL, SR
from .logger import log
_AUDIOSET_TO_NOISE_TYPE = {
"Television": "TV",
"Music": "music",
"Musical instrument": "music",
"Vehicle": "traffic/road noise",
"Car": "traffic/road noise",
"Traffic noise, roadway noise": "traffic/road noise",
"Wind": "wind",
"Wind noise (microphone)": "wind",
"Typing": "keyboard typing",
"Computer keyboard": "keyboard typing",
"Conversation": "office chatter",
"Chatter": "office chatter",
"Speech": None,
"Crowd": "office chatter",
"Static": "static",
"White noise": "static",
"Hum": "mechanical noise",
"Mechanisms": "mechanical noise",
"Engine": "mechanical noise",
"Telephone bell ringing": "office chatter",
"Dog": "background chatter/animal noise",
"Silence": None,
}
@functools.lru_cache(maxsize=1)
def _silero_vad():
model, utils = torch.hub.load(
"snakers4/silero-vad", "silero_vad", trust_repo=True, onnx=False
)
return model, utils
def speech_segments(y: np.ndarray) -> list[tuple[float, float]] | None:
try:
model, utils = _silero_vad()
get_speech_timestamps = utils[0]
wav = torch.from_numpy(y)
ts = get_speech_timestamps(wav, model, sampling_rate=SR)
return [(t["start"] / SR, t["end"] / SR) for t in ts]
except Exception as e:
log.exception("speech_segments failed: %s", e)
return None
@functools.lru_cache(maxsize=1)
def _emotion_model():
import torch.nn as nn
from transformers import Wav2Vec2Processor
from transformers.models.wav2vec2.modeling_wav2vec2 import (
Wav2Vec2Model,
Wav2Vec2PreTrainedModel,
)
class RegressionHead(nn.Module):
def __init__(self, config):
super().__init__()
self.dense = nn.Linear(config.hidden_size, config.hidden_size)
self.dropout = nn.Dropout(config.final_dropout)
self.out_proj = nn.Linear(config.hidden_size, config.num_labels)
def forward(self, features):
x = self.dropout(features)
x = torch.tanh(self.dense(x))
x = self.dropout(x)
return self.out_proj(x)
class EmotionModel(Wav2Vec2PreTrainedModel):
def __init__(self, config):
super().__init__(config)
self.wav2vec2 = Wav2Vec2Model(config)
self.classifier = RegressionHead(config)
self.init_weights()
def forward(self, input_values):
hidden = self.wav2vec2(input_values)[0]
pooled = torch.mean(hidden, dim=1)
return pooled, self.classifier(pooled)
processor = Wav2Vec2Processor.from_pretrained(EMOTION_DIM_MODEL)
model = EmotionModel.from_pretrained(EMOTION_DIM_MODEL).eval()
return processor, model
def dimensional_emotion(y: np.ndarray) -> dict | None:
try:
processor, model = _emotion_model()
inputs = processor(y, sampling_rate=SR, return_tensors="pt")
with torch.no_grad():
_, logits = model(inputs["input_values"])
arousal, dominance, valence = logits[0].tolist()
return {"arousal": arousal, "dominance": dominance, "valence": valence}
except Exception as e:
log.exception("dimensional_emotion failed: %s", e)
return None
@functools.lru_cache(maxsize=1)
def _categorical_ser_pipeline():
from transformers import pipeline
return pipeline("audio-classification", model=EMOTION_CAT_MODEL, top_k=4)
def categorical_emotion(y: np.ndarray) -> list[dict] | None:
try:
clf = _categorical_ser_pipeline()
return clf({"array": y, "sampling_rate": SR})
except Exception as e:
log.exception("categorical_emotion failed: %s", e)
return None
@functools.lru_cache(maxsize=1)
def _ast_pipeline():
from transformers import pipeline
return pipeline("audio-classification", model=AST_MODEL, top_k=10)
def noise_tags(y: np.ndarray) -> list[dict] | None:
try:
clf = _ast_pipeline()
return clf({"array": y, "sampling_rate": SR})
except Exception as e:
log.exception("noise_tags failed: %s", e)
return None
def noise_type_from_tags(tags: list[dict], min_score: float = 0.15) -> str:
for tag in tags:
mapped = _AUDIOSET_TO_NOISE_TYPE.get(tag["label"])
if mapped and tag["score"] >= min_score:
return mapped
return ""
@functools.lru_cache(maxsize=1)
def _segmentation_inference():
from pyannote.audio import Inference, Model
model = Model.from_pretrained(SEGMENTATION_MODEL)
return Inference(model, step=2.5)
def overlap_seconds(path: str) -> float | None:
try:
inference = _segmentation_inference()
output = inference(path)
frame_speaker_count = output.data.sum(axis=-1)
overlap_frames = int((frame_speaker_count >= 2).sum())
frame_duration = output.sliding_window.duration / output.data.shape[1]
return float(overlap_frames * frame_duration)
except Exception as e:
log.exception("overlap_seconds failed: %s", e)
return None
|