|
|
|
|
| import torch
|
| import numpy as np
|
| import librosa
|
| import asyncio
|
| import tempfile
|
| import os
|
| import time
|
| import traceback
|
| from typing import AsyncGenerator, Dict, Any, Optional, Union, Iterable
|
|
|
| import speech_recognition as sr
|
| from MMM import MMM
|
|
|
|
|
| class Speaker_ID:
|
| def __init__(
|
| self,
|
| mmm_manager,
|
| base_model_id: str = "unknown",
|
| device: Union[str, torch.device, None] = None,
|
| seq_len: int = 1200,
|
| sr: int = 1200,
|
| ):
|
| self.mmm = mmm_manager
|
| self.base_model_id = base_model_id
|
| self.seq_len = int(seq_len)
|
| self.sr = int(sr)
|
|
|
| if device is None:
|
| self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
| else:
|
| self.device = torch.device(device)
|
|
|
| if not hasattr(self.mmm, "models"):
|
| raise ValueError("Provided mmm_manager does not look like an MMM manager (missing .models).")
|
|
|
| if self.base_model_id not in self.mmm.models:
|
| available = list(self.mmm.models.keys())
|
| raise KeyError(f"Base model id '{self.base_model_id}' not found. Available keys: {available}")
|
|
|
| self.base_model = self.mmm.models[self.base_model_id].to(self.device)
|
| self.base_model.eval()
|
|
|
| def _audio_to_tensor(self, wav_path: str) -> torch.Tensor:
|
| y, _ = librosa.load(str(wav_path), sr=self.sr, mono=True)
|
| y = y.astype(np.float32)
|
| if y.size == 0:
|
| raise RuntimeError(f"Empty audio file: {wav_path}")
|
| maxv = float(np.max(np.abs(y)))
|
| if maxv > 0:
|
| y = y / maxv
|
| if y.shape[0] < self.seq_len:
|
| y = np.pad(y, (0, self.seq_len - y.shape[0]))
|
| else:
|
| y = y[: self.seq_len]
|
| return torch.from_numpy(y).unsqueeze(-1)
|
|
|
| def _ensure_tensor(self, features: Union[np.ndarray, torch.Tensor]) -> torch.Tensor:
|
| if isinstance(features, np.ndarray):
|
| t = torch.from_numpy(features)
|
| elif torch.is_tensor(features):
|
| t = features.clone()
|
| else:
|
| raise TypeError("audio_features must be numpy array or torch tensor or audio file path")
|
|
|
| if t.dim() == 1:
|
| t = t.unsqueeze(-1)
|
| if t.dim() == 2:
|
| return t.float()
|
| raise ValueError(f"Unexpected features tensor shape: {t.shape}")
|
|
|
| def generate_embedding(self, audio_input: Union[str, np.ndarray, torch.Tensor]) -> np.ndarray:
|
| if isinstance(audio_input, str):
|
| x = self._audio_to_tensor(audio_input)
|
| else:
|
| x = self._ensure_tensor(audio_input)
|
| x = x.to(self.device)
|
| if x.dim() == 2:
|
| x = x.unsqueeze(1)
|
|
|
| with torch.no_grad():
|
| out = self.base_model(x)
|
|
|
| if isinstance(out, dict):
|
| if "mu" in out:
|
| mu = out["mu"]
|
| emb_bz = mu.mean(dim=0)
|
| emb = emb_bz.squeeze(0).cpu().numpy()
|
| return emb
|
| if "z" in out:
|
| z = out["z"].mean(dim=0).squeeze(0).cpu().numpy()
|
| return z
|
| if "reconstruction" in out:
|
| recon = out["reconstruction"].mean(dim=0).squeeze(0).cpu().numpy()
|
| return recon
|
|
|
| if torch.is_tensor(out):
|
| arr = out.mean(dim=0).squeeze(0).cpu().numpy()
|
| return arr
|
|
|
| raise KeyError("Base model forward did not return 'mu', 'z', 'reconstruction' or a tensor to use as embedding.")
|
|
|
| def enroll_speaker(
|
| self,
|
| speaker_id: str,
|
| audio_input: Union[str, np.ndarray, torch.Tensor],
|
| model_type: str = "mmm",
|
| n_components: int = 4,
|
| epochs: int = 50,
|
| lr: float = 1e-3,
|
| seq_len_for_mmm: int = None,
|
| **fit_kwargs,
|
| ) -> str:
|
| model_type = model_type.lower()
|
| if model_type not in ("gmm", "hmm", "mmm"):
|
| raise ValueError("model_type must be 'gmm', 'hmm', or 'mmm'")
|
|
|
| emb = self.generate_embedding(audio_input)
|
| if model_type == "gmm":
|
| X = np.asarray(emb, dtype=np.float32)[None, :]
|
| self.mmm.fit_and_add(
|
| data=X,
|
| model_type="gmm",
|
| model_id=speaker_id,
|
| n_components=n_components,
|
| lr=lr,
|
| epochs=epochs,
|
| **fit_kwargs,
|
| )
|
| else:
|
| T = int(seq_len_for_mmm or self.seq_len)
|
| z = torch.tensor(emb, dtype=torch.float32, device=self.device)
|
| seq = z.unsqueeze(0).repeat(T, 1)
|
| seq = seq.unsqueeze(1)
|
| self.mmm.fit_and_add(
|
| data=seq,
|
| model_type="mmm" if model_type == "mmm" else "hmm",
|
| model_id=speaker_id,
|
| input_dim=emb.shape[-1],
|
| output_dim=emb.shape[-1],
|
| hidden_dim=emb.shape[-1] * 2,
|
| z_dim=min(256, emb.shape[-1]),
|
| rnn_hidden=emb.shape[-1],
|
| num_states=fit_kwargs.get("num_states", 8),
|
| n_mix=fit_kwargs.get("n_mix", 2),
|
| trans_d_model=fit_kwargs.get("trans_d_model", 64),
|
| trans_nhead=fit_kwargs.get("trans_nhead", 4),
|
| trans_layers=fit_kwargs.get("trans_layers", 2),
|
| lr=lr,
|
| epochs=epochs,
|
| **fit_kwargs,
|
| )
|
|
|
| return speaker_id
|
|
|
| def identify(
|
| self,
|
| audio_input: Union[str, np.ndarray, torch.Tensor],
|
| unknown_label_confidence_margin: float = 0.0,
|
| ):
|
| emb = self.generate_embedding(audio_input)
|
| emb_np = np.asarray(emb, dtype=np.float32)
|
| X_try = emb_np[None, :]
|
|
|
| scores: Dict[str, float] = {}
|
| for model_id in list(self.mmm.models.keys()):
|
| try:
|
| sc = self.mmm.score(model_id, X_try)
|
| if isinstance(sc, dict):
|
| vals = []
|
| for v in sc.values():
|
| try:
|
| vals.append(float(np.asarray(v).mean()))
|
| except Exception:
|
| pass
|
| score_val = float(np.mean(vals)) if vals else float("nan")
|
| else:
|
| try:
|
| score_val = float(np.asarray(sc).mean())
|
| except Exception:
|
| score_val = float(sc)
|
| scores[model_id] = score_val
|
| except Exception:
|
| try:
|
| T = self.seq_len
|
| seq = np.tile(emb_np[None, :], (T, 1, 1))
|
| sc = self.mmm.score(model_id, seq)
|
| try:
|
| scores[model_id] = float(np.asarray(sc).mean())
|
| except Exception:
|
| scores[model_id] = float(sc)
|
| except Exception:
|
| continue
|
|
|
| if not scores:
|
| return self.base_model_id, float("nan"), {}
|
|
|
| best_model, best_score = max(scores.items(), key=lambda kv: kv[1])
|
|
|
| if best_model != self.base_model_id and unknown_label_confidence_margin > 0.0:
|
| unknown_score = scores.get(self.base_model_id, float("-inf"))
|
| if best_score <= unknown_score + unknown_label_confidence_margin:
|
| return self.base_model_id, unknown_score, scores
|
|
|
| return best_model, best_score, scores
|
|
|
|
|
|
|
|
|
| async def ASI(
|
| phrase_time_limit: Optional[float] = 3.0,
|
| queue_maxsize: int = 8,
|
| mmm_pt_path: str = "models/MMM/mmm.pt",
|
| ) -> AsyncGenerator[Dict[str, Any], None]:
|
| mgr = MMM.load(mmm_pt_path)
|
| speaker_system = Speaker_ID(mmm_manager=mgr, base_model_id="unknown", seq_len=1200, sr=1200)
|
|
|
| loop = asyncio.get_running_loop()
|
| audio_q: asyncio.Queue = asyncio.Queue(maxsize=queue_maxsize)
|
|
|
| recognizer = sr.Recognizer()
|
| try:
|
| mic = sr.Microphone()
|
| except Exception as e:
|
| raise RuntimeError("Could not open microphone. Check drivers / permissions.") from e
|
|
|
| def _bg_callback(recognizer_obj: sr.Recognizer, audio: sr.AudioData) -> None:
|
| try:
|
| wav_bytes = audio.get_wav_data()
|
| try:
|
| loop.call_soon_threadsafe(audio_q.put_nowait, wav_bytes)
|
| except Exception:
|
| pass
|
| except Exception:
|
| traceback.print_exc()
|
|
|
| stop_listening = recognizer.listen_in_background(mic, _bg_callback, phrase_time_limit=phrase_time_limit)
|
|
|
| try:
|
| while True:
|
| try:
|
| wav_bytes = await audio_q.get()
|
| except asyncio.CancelledError:
|
| break
|
|
|
| if wav_bytes is None:
|
| continue
|
|
|
| def _write_temp_wav(b: bytes) -> str:
|
| tf = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
|
| try:
|
| tf.write(b)
|
| tf.flush()
|
| return tf.name
|
| finally:
|
| tf.close()
|
|
|
| tmp_path = await loop.run_in_executor(None, _write_temp_wav, wav_bytes)
|
|
|
| try:
|
| result = await loop.run_in_executor(None, speaker_system.identify, tmp_path)
|
| best_speaker, best_score, scores = result
|
| yield {
|
| "speaker": best_speaker,
|
| "score": best_score,
|
| "scores": scores,
|
| "path": tmp_path,
|
| "timestamp": time.time(),
|
| }
|
| except Exception as e:
|
| yield {
|
| "error": str(e),
|
| "traceback": traceback.format_exc(),
|
| "path": tmp_path,
|
| "timestamp": time.time(),
|
| }
|
| finally:
|
| try:
|
| os.remove(tmp_path)
|
| except Exception:
|
| pass
|
|
|
| finally:
|
| try:
|
| stop_listening(wait_for_stop=False)
|
| except Exception:
|
| pass
|
|
|
|
|
| async def _main_cli():
|
| async for res in ASI(phrase_time_limit=3.0):
|
| if "error" in res:
|
| print("ID error:", res["error"])
|
| else:
|
| ts = time.ctime(res["timestamp"])
|
| print(f"[{ts}] Predicted: {res['speaker']} (score={res['score']})")
|
| print("All scores:", res["scores"])
|
|
|
|
|
| if __name__ == "__main__":
|
| asyncio.run(_main_cli())
|
|
|
