"""Speaker diarization using TEN-VAD + ECAPA-TDNN + spectral clustering. Spectral clustering implementation adapted from FunASR/3D-Speaker: https://github.com/alibaba-damo-academy/FunASR MIT License (https://opensource.org/licenses/MIT) """ import warnings import numpy as np import scipy import sklearn.metrics.pairwise import torch from sklearn.cluster._kmeans import k_means from sklearn.preprocessing import normalize def _get_device() -> torch.device: """Get best available device for inference.""" if torch.cuda.is_available(): return torch.device("cuda") if torch.backends.mps.is_available(): return torch.device("mps") return torch.device("cpu") class SpectralCluster: """Spectral clustering using unnormalized Laplacian of affinity matrix. Adapted from FunASR/3D-Speaker and SpeechBrain implementations. Uses eigenvalue gap to automatically determine number of speakers. """ def __init__(self, min_num_spks: int = 1, max_num_spks: int = 15, pval: float = 0.06): self.min_num_spks = min_num_spks self.max_num_spks = max_num_spks self.pval = pval def __call__(self, embeddings: np.ndarray, oracle_num: int | None = None) -> np.ndarray: """Run spectral clustering on embeddings. Args: embeddings: Speaker embeddings of shape [N, D] oracle_num: Optional known number of speakers Returns: Cluster labels of shape [N] """ # Similarity matrix computation sim_mat = self.get_sim_mat(embeddings) # Refining similarity matrix with pval prunned_sim_mat = self.p_pruning(sim_mat) # Symmetrization sym_prund_sim_mat = 0.5 * (prunned_sim_mat + prunned_sim_mat.T) # Laplacian calculation laplacian = self.get_laplacian(sym_prund_sim_mat) # Get Spectral Embeddings emb, num_of_spk = self.get_spec_embs(laplacian, oracle_num) # Perform clustering return self.cluster_embs(emb, num_of_spk) def get_sim_mat(self, embeddings: np.ndarray) -> np.ndarray: """Compute cosine similarity matrix.""" return sklearn.metrics.pairwise.cosine_similarity(embeddings, embeddings) def p_pruning(self, affinity: np.ndarray) -> np.ndarray: """Prune low similarity values in affinity matrix (keep top pval fraction).""" n = affinity.shape[0] pval = max(self.pval, 6.0 / n) k_keep = max(1, int(pval * n)) # Vectorized: find top-k indices per row and zero out the rest top_k_idx = np.argpartition(affinity, -k_keep, axis=1)[:, -k_keep:] mask = np.zeros_like(affinity, dtype=bool) np.put_along_axis(mask, top_k_idx, True, axis=1) affinity[~mask] = 0 return affinity def get_laplacian(self, sim_mat: np.ndarray) -> np.ndarray: """Compute unnormalized Laplacian matrix.""" from scipy.sparse.csgraph import laplacian np.fill_diagonal(sim_mat, 0) return laplacian(sim_mat, normed=False) def get_spec_embs( self, laplacian: np.ndarray, k_oracle: int | None = None ) -> tuple[np.ndarray, int]: """Extract spectral embeddings from Laplacian.""" lambdas, eig_vecs = scipy.linalg.eigh(laplacian) if k_oracle is not None: num_of_spk = k_oracle else: lambda_gap_list = self.get_eigen_gaps( lambdas[self.min_num_spks - 1 : self.max_num_spks + 1] ) num_of_spk = np.argmax(lambda_gap_list) + self.min_num_spks emb = eig_vecs[:, :num_of_spk] return emb, num_of_spk def cluster_embs(self, emb: np.ndarray, k: int) -> np.ndarray: """Cluster spectral embeddings using k-means.""" _, labels, _ = k_means(emb, k, n_init=10) return labels def get_eigen_gaps(self, eig_vals: np.ndarray) -> np.ndarray: """Compute gaps between consecutive eigenvalues.""" return np.diff(eig_vals) class SpeakerClusterer: """Speaker clustering backend using spectral clustering with speaker merging. Features: - Spectral clustering with eigenvalue gap for auto speaker count detection - P-pruning for affinity matrix refinement - Post-clustering speaker merging by cosine similarity """ def __init__( self, min_num_spks: int = 2, max_num_spks: int = 10, merge_thr: float = 0.90, # Moderate merging ): self.min_num_spks = min_num_spks self.max_num_spks = max_num_spks self.merge_thr = merge_thr self._spectral_cluster: SpectralCluster | None = None def _get_spectral_cluster(self) -> SpectralCluster: """Lazy-load spectral clusterer.""" if self._spectral_cluster is None: self._spectral_cluster = SpectralCluster( min_num_spks=self.min_num_spks, max_num_spks=self.max_num_spks, ) return self._spectral_cluster def __call__(self, embeddings: np.ndarray, num_speakers: int | None = None) -> np.ndarray: """Cluster speaker embeddings and return labels. Args: embeddings: Speaker embeddings of shape [N, D] num_speakers: Optional oracle number of speakers Returns: Cluster labels of shape [N] """ if len(embeddings.shape) != 2: raise ValueError(f"Expected 2D array, got shape {embeddings.shape}") # Handle edge cases if embeddings.shape[0] == 0: return np.array([], dtype=int) if embeddings.shape[0] == 1: return np.array([0], dtype=int) if embeddings.shape[0] < 6: return np.zeros(embeddings.shape[0], dtype=int) # Normalize embeddings and replace NaN/inf embeddings = np.nan_to_num(embeddings, nan=0.0, posinf=0.0, neginf=0.0) embeddings = normalize(embeddings) # Run spectral clustering (suppress numerical warnings) spectral = self._get_spectral_cluster() # Update min/max for oracle case if num_speakers is not None: spectral.min_num_spks = num_speakers spectral.max_num_spks = num_speakers with warnings.catch_warnings(): warnings.filterwarnings("ignore", category=RuntimeWarning) labels = spectral(embeddings, oracle_num=num_speakers) # Reset min/max if num_speakers is not None: spectral.min_num_spks = self.min_num_spks spectral.max_num_spks = self.max_num_spks # Merge similar speakers if no oracle if num_speakers is None: labels = self._merge_by_cos(labels, embeddings, self.merge_thr) # Re-index labels sequentially _, labels = np.unique(labels, return_inverse=True) return labels def _merge_by_cos(self, labels: np.ndarray, embs: np.ndarray, cos_thr: float) -> np.ndarray: """Merge similar speakers by cosine similarity of centroids.""" from scipy.cluster.hierarchy import fcluster, linkage from scipy.spatial.distance import pdist unique_labels = np.unique(labels) if len(unique_labels) <= 1: return labels # Compute normalized speaker centroids centroids = np.array([embs[labels == lbl].mean(0) for lbl in unique_labels]) centroids = normalize(centroids) # Hierarchical clustering with cosine distance distances = pdist(centroids, metric="cosine") linkage_matrix = linkage(distances, method="average") merged_labels = fcluster(linkage_matrix, t=1.0 - cos_thr, criterion="distance") - 1 # Map original labels to merged labels label_map = dict(zip(unique_labels, merged_labels)) return np.array([label_map[lbl] for lbl in labels]) class LocalSpeakerDiarizer: """Local speaker diarization using TEN-VAD + ECAPA-TDNN + spectral clustering. Pipeline: 1. TEN-VAD detects speech segments 2. Sliding window (1.0s, 75% overlap) for uniform embedding extraction 3. ECAPA-TDNN extracts speaker embeddings per window 4. Spectral clustering with eigenvalue gap for auto speaker detection 5. Frame-level consensus voting for segment reconstruction 6. Post-processing merges short segments to reduce flicker Tunable Parameters (class attributes): - WINDOW_SIZE: Embedding extraction window size in seconds - STEP_SIZE: Sliding window step size (overlap = WINDOW_SIZE - STEP_SIZE) - VAD_THRESHOLD: Speech detection threshold (lower = more sensitive) - VAD_MIN_DURATION: Minimum speech segment duration - VAD_MAX_GAP: Maximum gap to bridge between segments - VAD_PAD_ONSET/OFFSET: Padding added to speech segments - VOTING_RATE: Frame resolution for consensus voting - MIN_SEGMENT_DURATION: Minimum final segment duration - SAME_SPEAKER_GAP: Maximum gap to merge same-speaker segments - TAIL_COVERAGE_RATIO: Minimum tail coverage to add extra window """ _ten_vad_model = None _ecapa_model = None _device = None # ==================== TUNABLE PARAMETERS ==================== # Sliding window for embedding extraction WINDOW_SIZE = 0.75 # seconds - shorter window for finer resolution STEP_SIZE = 0.15 # seconds (80% overlap for more votes) TAIL_COVERAGE_RATIO = 0.1 # Add extra window if tail > this ratio of window # VAD hysteresis parameters VAD_THRESHOLD = 0.25 # Balanced threshold VAD_MIN_DURATION = 0.05 # Minimum speech segment duration (seconds) VAD_MAX_GAP = 0.50 # Bridge gaps shorter than this (seconds) VAD_PAD_ONSET = 0.05 # Padding at segment start (seconds) VAD_PAD_OFFSET = 0.05 # Padding at segment end (seconds) # Frame-level voting VOTING_RATE = 0.01 # 10ms resolution for consensus voting # Post-processing MIN_SEGMENT_DURATION = 0.15 # Minimum final segment duration (seconds) SHORT_SEGMENT_GAP = 0.1 # Gap threshold for merging short segments SAME_SPEAKER_GAP = 0.5 # Gap threshold for merging same-speaker segments # =========================================================== @classmethod def _get_ten_vad_model(cls): """Lazy-load TEN-VAD model (singleton).""" if cls._ten_vad_model is None: from ten_vad import TenVad cls._ten_vad_model = TenVad(hop_size=256, threshold=cls.VAD_THRESHOLD) return cls._ten_vad_model @classmethod def _get_device(cls) -> torch.device: """Get the best available device.""" if cls._device is None: cls._device = _get_device() return cls._device @classmethod def _get_ecapa_model(cls): """Lazy-load ECAPA-TDNN speaker embedding model (singleton).""" if cls._ecapa_model is None: # Suppress torchaudio deprecation warning from SpeechBrain with warnings.catch_warnings(): warnings.filterwarnings("ignore", message="torchaudio._backend") from speechbrain.inference.speaker import EncoderClassifier device = cls._get_device() cls._ecapa_model = EncoderClassifier.from_hparams( source="speechbrain/spkrec-ecapa-voxceleb", run_opts={"device": str(device)}, ) return cls._ecapa_model @classmethod def diarize( cls, audio: np.ndarray | str, sample_rate: int = 16000, num_speakers: int | None = None, min_speakers: int = 2, max_speakers: int = 10, **_kwargs, ) -> list[dict]: """Run speaker diarization on audio. Args: audio: Audio waveform as numpy array or path to audio file sample_rate: Audio sample rate (default 16000) num_speakers: Exact number of speakers (if known) min_speakers: Minimum number of speakers max_speakers: Maximum number of speakers Returns: List of dicts with 'speaker', 'start', 'end' keys """ # Handle file path input if isinstance(audio, str): import librosa audio, sample_rate = librosa.load(audio, sr=16000) # Ensure correct sample rate if sample_rate != 16000: import librosa audio = librosa.resample(audio, orig_sr=sample_rate, target_sr=16000) sample_rate = 16000 audio = audio.astype(np.float32) total_duration = len(audio) / sample_rate # Step 1: VAD (returns segments and raw frame-level decisions) segments, vad_frames = cls._get_speech_segments(audio, sample_rate) if not segments: return [] # Step 2: Extract embeddings embeddings, window_segments = cls._extract_embeddings(audio, segments, sample_rate) if len(embeddings) == 0: return [] # Step 3: Cluster clusterer = SpeakerClusterer(min_num_spks=min_speakers, max_num_spks=max_speakers) labels = clusterer(embeddings, num_speakers) # Step 4: Post-process with consensus voting (VAD-aware) return cls._postprocess_segments(window_segments, labels, total_duration, vad_frames) @classmethod def _get_speech_segments( cls, audio_array: np.ndarray, sample_rate: int = 16000 ) -> tuple[list[dict], list[bool]]: """Get speech segments using TEN-VAD. Returns: Tuple of (segments list, vad_frames list of per-frame speech decisions) """ vad_model = cls._get_ten_vad_model() # Convert to int16 as required by TEN-VAD # Clip to prevent integer overflow if audio_array.dtype != np.int16: audio_int16 = (np.clip(audio_array, -1.0, 1.0) * 32767).astype(np.int16) else: audio_int16 = audio_array # Process frame by frame hop_size = 256 frame_duration = hop_size / sample_rate speech_frames: list[bool] = [] for i in range(0, len(audio_int16) - hop_size, hop_size): frame = audio_int16[i : i + hop_size] _, is_speech = vad_model.process(frame) speech_frames.append(is_speech) # Convert frame-level decisions to segments segments = [] in_speech = False start_idx = 0 for i, is_speech in enumerate(speech_frames): if is_speech and not in_speech: start_idx = i in_speech = True elif not is_speech and in_speech: start_time = start_idx * frame_duration end_time = i * frame_duration segments.append( { "start": start_time, "end": end_time, "start_sample": int(start_time * sample_rate), "end_sample": int(end_time * sample_rate), } ) in_speech = False # Handle trailing speech if in_speech: start_time = start_idx * frame_duration end_time = len(speech_frames) * frame_duration segments.append( { "start": start_time, "end": end_time, "start_sample": int(start_time * sample_rate), "end_sample": int(end_time * sample_rate), } ) return cls._apply_vad_hysteresis(segments, sample_rate), speech_frames @classmethod def _apply_vad_hysteresis(cls, segments: list[dict], sample_rate: int = 16000) -> list[dict]: """Apply hysteresis-like post-processing to VAD segments.""" if not segments: return segments segments = sorted(segments, key=lambda x: x["start"]) # Fill short gaps merged = [segments[0].copy()] for seg in segments[1:]: gap = seg["start"] - merged[-1]["end"] if gap <= cls.VAD_MAX_GAP: merged[-1]["end"] = seg["end"] merged[-1]["end_sample"] = seg["end_sample"] else: merged.append(seg.copy()) # Remove short segments filtered = [seg for seg in merged if (seg["end"] - seg["start"]) >= cls.VAD_MIN_DURATION] # Dilate segments (add padding) for seg in filtered: seg["start"] = max(0.0, seg["start"] - cls.VAD_PAD_ONSET) seg["end"] = seg["end"] + cls.VAD_PAD_OFFSET seg["start_sample"] = int(seg["start"] * sample_rate) seg["end_sample"] = int(seg["end"] * sample_rate) return filtered @classmethod def _extract_embeddings( cls, audio_array: np.ndarray, segments: list[dict], sample_rate: int ) -> tuple[np.ndarray, list[dict]]: """Extract speaker embeddings using sliding windows.""" speaker_model = cls._get_ecapa_model() window_samples = int(cls.WINDOW_SIZE * sample_rate) step_samples = int(cls.STEP_SIZE * sample_rate) embeddings = [] window_segments = [] with torch.no_grad(): for seg in segments: seg_start = seg["start_sample"] seg_end = seg["end_sample"] seg_len = seg_end - seg_start # Generate window positions if seg_len <= window_samples: starts = [seg_start] ends = [seg_end] else: starts = list(range(seg_start, seg_end - window_samples + 1, step_samples)) ends = [s + window_samples for s in starts] # Cover tail if > TAIL_COVERAGE_RATIO of window remains if ends and ends[-1] < seg_end: remainder = seg_end - ends[-1] if remainder > (window_samples * cls.TAIL_COVERAGE_RATIO): starts.append(seg_end - window_samples) ends.append(seg_end) for c_start, c_end in zip(starts, ends): chunk = audio_array[c_start:c_end] # Pad short chunks with reflection if len(chunk) < window_samples: pad_width = window_samples - len(chunk) chunk = np.pad(chunk, (0, pad_width), mode="reflect") # Extract embedding using SpeechBrain's encode_batch chunk_tensor = torch.from_numpy(chunk).float().unsqueeze(0) embedding = ( speaker_model.encode_batch(chunk_tensor).squeeze(0).squeeze(0).cpu().numpy() ) # Validate embedding if np.isfinite(embedding).all() and np.linalg.norm(embedding) > 1e-8: embeddings.append(embedding) window_segments.append( { "start": c_start / sample_rate, "end": c_end / sample_rate, } ) # Normalize all embeddings at once if embeddings: return normalize(np.array(embeddings)), window_segments return np.array([]), [] @classmethod def _resample_vad(cls, vad_frames: list[bool], num_frames: int) -> np.ndarray: """Resample VAD frame decisions to match voting grid resolution. VAD operates at 256 samples / 16000 Hz = 16ms per frame. Voting operates at VOTING_RATE (default 10ms) per frame. This maps VAD decisions to the finer voting grid. """ if not vad_frames: return np.zeros(num_frames, dtype=bool) vad_rate = 256 / 16000 # 16ms per VAD frame vad_arr = np.array(vad_frames) # Vectorized: compute VAD frame indices for each voting frame voting_times = np.arange(num_frames) * cls.VOTING_RATE vad_indices = np.clip((voting_times / vad_rate).astype(int), 0, len(vad_arr) - 1) return vad_arr[vad_indices] @classmethod def _postprocess_segments( cls, window_segments: list[dict], labels: np.ndarray, total_duration: float, vad_frames: list[bool], ) -> list[dict]: """Post-process using frame-level consensus voting with VAD-aware silence.""" if not window_segments or len(labels) == 0: return [] # Correct labels to be contiguous unique_labels = np.unique(labels) label_map = {old: new for new, old in enumerate(unique_labels)} clean_labels = np.array([label_map[lbl] for lbl in labels]) num_speakers = len(unique_labels) if num_speakers == 0: return [] # Create voting grid num_frames = int(np.ceil(total_duration / cls.VOTING_RATE)) + 1 votes = np.zeros((num_frames, num_speakers), dtype=np.float32) # Accumulate votes for win, label in zip(window_segments, clean_labels): start_frame = int(win["start"] / cls.VOTING_RATE) end_frame = int(win["end"] / cls.VOTING_RATE) end_frame = min(end_frame, num_frames) if start_frame < end_frame: votes[start_frame:end_frame, label] += 1.0 # Determine winner per frame frame_speakers = np.argmax(votes, axis=1) max_votes = np.max(votes, axis=1) # Resample VAD to voting grid resolution for silence-aware voting vad_resampled = cls._resample_vad(vad_frames, num_frames) # Convert frames to segments final_segments = [] current_speaker = -1 seg_start = 0.0 for f in range(num_frames): speaker = int(frame_speakers[f]) score = max_votes[f] # Force silence if VAD says no speech OR no votes if score == 0 or not vad_resampled[f]: speaker = -1 if speaker != current_speaker: if current_speaker != -1: final_segments.append( { "speaker": f"SPEAKER_{current_speaker}", "start": seg_start, "end": f * cls.VOTING_RATE, } ) current_speaker = speaker seg_start = f * cls.VOTING_RATE # Close last segment if current_speaker != -1: final_segments.append( { "speaker": f"SPEAKER_{current_speaker}", "start": seg_start, "end": num_frames * cls.VOTING_RATE, } ) return cls._merge_short_segments(final_segments) @classmethod def _merge_short_segments(cls, segments: list[dict]) -> list[dict]: """Merge short segments to reduce flicker.""" if not segments: return [] clean: list[dict] = [] for seg in segments: dur = seg["end"] - seg["start"] if dur < cls.MIN_SEGMENT_DURATION: if ( clean and clean[-1]["speaker"] == seg["speaker"] and seg["start"] - clean[-1]["end"] < cls.SHORT_SEGMENT_GAP ): clean[-1]["end"] = seg["end"] continue if ( clean and clean[-1]["speaker"] == seg["speaker"] and seg["start"] - clean[-1]["end"] < cls.SAME_SPEAKER_GAP ): clean[-1]["end"] = seg["end"] else: clean.append(seg) return clean @classmethod def assign_speakers_to_words( cls, words: list[dict], speaker_segments: list[dict], ) -> list[dict]: """Assign speaker labels to words based on timestamp overlap. Args: words: List of word dicts with 'word', 'start', 'end' keys speaker_segments: List of speaker dicts with 'speaker', 'start', 'end' keys Returns: Words list with 'speaker' key added to each word """ for word in words: word_mid = (word["start"] + word["end"]) / 2 # Find the speaker segment that contains this word's midpoint best_speaker = None for seg in speaker_segments: if seg["start"] <= word_mid <= seg["end"]: best_speaker = seg["speaker"] break # If no exact match, find closest segment if best_speaker is None and speaker_segments: min_dist = float("inf") for seg in speaker_segments: seg_mid = (seg["start"] + seg["end"]) / 2 dist = abs(word_mid - seg_mid) if dist < min_dist: min_dist = dist best_speaker = seg["speaker"] word["speaker"] = best_speaker return words class SpeakerDiarizer: """Speaker diarization using TEN-VAD + ECAPA-TDNN + spectral clustering. Example: >>> segments = SpeakerDiarizer.diarize(audio_array) >>> for seg in segments: ... print(f"{seg['speaker']}: {seg['start']:.2f} - {seg['end']:.2f}") """ @classmethod def diarize( cls, audio: np.ndarray | str, sample_rate: int = 16000, num_speakers: int | None = None, min_speakers: int | None = None, max_speakers: int | None = None, **_kwargs, ) -> list[dict]: """Run speaker diarization on audio. Args: audio: Audio waveform as numpy array or path to audio file sample_rate: Audio sample rate (default 16000) num_speakers: Exact number of speakers (if known) min_speakers: Minimum number of speakers max_speakers: Maximum number of speakers Returns: List of dicts with 'speaker', 'start', 'end' keys """ return LocalSpeakerDiarizer.diarize( audio, sample_rate=sample_rate, num_speakers=num_speakers, min_speakers=min_speakers or 2, max_speakers=max_speakers or 10, ) @classmethod def assign_speakers_to_words( cls, words: list[dict], speaker_segments: list[dict], ) -> list[dict]: """Assign speaker labels to words based on timestamp overlap.""" return LocalSpeakerDiarizer.assign_speakers_to_words(words, speaker_segments)