""" Speaker Diarization Module ========================== Implements VAD + Speaker Embedding + Clustering pipeline for speaker diarization. """ from __future__ import annotations import logging from dataclasses import dataclass, field from pathlib import Path from typing import Any, Dict, List, Optional, Tuple import numpy as np import torch from sklearn.cluster import AgglomerativeClustering, KMeans, SpectralClustering from sklearn.metrics import silhouette_score from sklearn.preprocessing import StandardScaler from src.utils import setup_logger @dataclass class DiarizationConfig: """Configuration for speaker diarization""" # VAD settings vad_threshold: float = 0.5 min_speech_duration: float = 0.3 min_silence_duration: float = 0.3 # Segmentation settings segment_window: float = 1.5 segment_hop: float = 0.75 # Clustering settings clustering_method: str = "agglomerative" clustering_threshold: float = 0.7 min_cluster_size: int = 2 max_speakers: Optional[int] = None # Post-processing merge_gap_threshold: float = 0.5 min_segment_duration: float = 0.3 # Model settings embedding_model_id: str = "speechbrain/spkrec-ecapa-voxceleb" use_speechbrain: bool = True # prefer SpeechBrain embeddings allow_fallback: bool = False # if False, raise an error when SpeechBrain cannot be loaded # Collapse heuristics collapse_threshold: float = 0.15 # When negative, do not automatically collapse clusters to a single speaker based on silhouette. silhouette_collapse_threshold: float = -1.0 # Iterative merging (centroid-based) iterative_merge_threshold: float = 0.15 iterative_merge_silhouette_threshold: float = 0.0 iterative_merge_max_iters: int = 10 # Performance tuning embedding_batch_size: int = 32 embedding_cache: bool = True # write/load embedding arrays to cache_dir use_fast_embedding: bool = False # use MFCC deterministic embeddings for speed # Optional: target speaker count - if set, clusters will be greedily merged to meet target target_num_speakers: Optional[int] = None target_force_threshold: float = ( 1.0 # 1.0 => allow merges regardless of distance; lower = more conservative ) # Device device: str = "cuda" if torch.cuda.is_available() else "cpu" @dataclass class SpeakerSegment: """Represents a speaker segment with timing and metadata""" speaker_id: str start: float end: float confidence: float = 1.0 is_overlap: bool = False embedding: Optional[np.ndarray] = None metadata: Dict[str, Any] = field(default_factory=dict) @property def duration(self) -> float: """Get segment duration in seconds""" return self.end - self.start def to_dict(self) -> Dict[str, Any]: """Convert to dictionary""" return { "speaker_id": self.speaker_id, "start": self.start, "end": self.end, "confidence": self.confidence, "is_overlap": self.is_overlap, "duration": self.duration, } class SpeakerDiarizer: """ Speaker Diarization using SpeechBrain ECAPA-TDNN embeddings. Pipeline: 1. Voice Activity Detection (VAD) 2. Audio segmentation into windows 3. Speaker embedding extraction (ECAPA-TDNN) 4. Clustering to assign speaker labels 5. Post-processing (merging, smoothing) Attributes: config: DiarizationConfig object Example: >>> diarizer = SpeakerDiarizer() >>> segments = diarizer.process(waveform, sample_rate=16000, num_speakers=4) >>> for seg in segments: ... print(f"{seg.speaker_id}: {seg.start:.2f}s - {seg.end:.2f}s") """ def __init__(self, config: Optional[DiarizationConfig] = None, models_dir: str = "./models"): """ Initialize SpeakerDiarizer. Args: config: DiarizationConfig object models_dir: Directory to cache downloaded models """ self.config = config or DiarizationConfig() self.models_dir = Path(models_dir) self.models_dir.mkdir(parents=True, exist_ok=True) self.device = self.config.device # Setup logger self.logger = setup_logger("SpeakerDiarizer") # Model placeholders (lazy loading) self._embedding_model = None self._vad_model = None self._embedding_model_is_speechbrain = False def _load_embedding_model(self): """Lazy load speaker embedding model This function will attempt to patch missing torchaudio APIs (e.g., list_audio_backends) so that SpeechBrain imports cleanly on environments with older torchaudio builds. """ if self._embedding_model is None: # Shim torchaudio compatibility if needed (some torchaudio versions lack list_audio_backends) try: import importlib if importlib.util.find_spec("torchaudio"): import torchaudio if not hasattr(torchaudio, "list_audio_backends"): def _list_audio_backends(): # best-effort guess of available backends; not exhaustive backends = [] try: # prefer sox_io and soundfile as common options backends.append("sox_io") except Exception: pass try: backends.append("soundfile") except Exception: pass if not backends: backends = ["sox_io"] return backends torchaudio.list_audio_backends = _list_audio_backends if not hasattr(torchaudio, "get_audio_backend"): torchaudio.get_audio_backend = lambda: torchaudio.list_audio_backends()[0] except Exception: # best-effort only, don't prevent embedding loading attempt pass try: from speechbrain.inference.speaker import EncoderClassifier self.logger.info(f"Loading embedding model: {self.config.embedding_model_id}") import os # Prefer to disable HF symlinks up-front on Windows to prevent permission errors os.environ.setdefault("HF_HUB_DISABLE_SYMLINKS", "1") # Try a robust direct download into a local models directory to avoid symlinks entirely dest_dir = str(self.models_dir / self.config.embedding_model_id.replace("/", "_")) try: from huggingface_hub import snapshot_download self.logger.info( f"Attempting to snapshot_download model to local dir {dest_dir} (no symlinks)" ) os.makedirs(dest_dir, exist_ok=True) snapshot_download( repo_id=self.config.embedding_model_id, local_dir=dest_dir, local_dir_use_symlinks=False, ) # Try to load from the locally downloaded snapshot try: self._embedding_model = EncoderClassifier.from_hparams( source=dest_dir, savedir=dest_dir, run_opts={"device": self.device}, ) self.logger.info("Embedding model loaded successfully from local snapshot") # mark that we used speechbrain self._embedding_model_is_speechbrain = True return except Exception as e_local: self.logger.warning(f"Local snapshot load failed: {e_local}") except Exception: # snapshot_download not available or failed; continue with other strategies pass try: # First try: load directly from hf cache (no savedir) - this typically avoids writing symlinks self._embedding_model = EncoderClassifier.from_hparams( source=self.config.embedding_model_id, run_opts={"device": self.device}, ) self.logger.info("Embedding model loaded successfully (from HF cache)") self._embedding_model_is_speechbrain = True return except Exception as e: err_msg = str(e) # Detect Windows symlink permission error and retry with savedir + disabled symlink env if ( ("A required privilege" in err_msg) or ("symlink" in err_msg.lower()) or getattr(e, "winerror", None) == 1314 ): try: os.environ["HF_HUB_DISABLE_SYMLINKS"] = "1" self.logger.warning( "Detected symlink/permission issue; retrying model load with HF_HUB_DISABLE_SYMLINKS=1 and specifying savedir" ) self._embedding_model = EncoderClassifier.from_hparams( source=self.config.embedding_model_id, savedir=str(self.models_dir / "spkrec-ecapa"), run_opts={"device": self.device}, ) self.logger.info( "Embedding model loaded successfully (after disabling symlinks)" ) self._embedding_model_is_speechbrain = True return except Exception: # Try monkeypatching SB fetch to use COPY try: import speechbrain.utils.fetching as sbfetch orig_fetch = sbfetch.fetch def _fetch_copy(*args, **kwargs): kwargs.setdefault("local_strategy", sbfetch.LocalStrategy.COPY) return orig_fetch(*args, **kwargs) sbfetch.fetch = _fetch_copy self.logger.info( "Retrying model load with SpeechBrain fetch set to COPY strategy" ) self._embedding_model = EncoderClassifier.from_hparams( source=self.config.embedding_model_id, savedir=str(self.models_dir / "spkrec-ecapa"), run_opts={"device": self.device}, ) self.logger.info( "Embedding model loaded successfully (after switching fetch strategy)" ) self._embedding_model_is_speechbrain = True return except Exception as e3: err_msg = str(e3) finally: try: sbfetch.fetch = orig_fetch except Exception: pass self.logger.error(f"Failed to load SpeechBrain embedding model: {err_msg}") # Try to salvage by copying an existing cached snapshot or downloading directly into dest_dir try: import re import shutil m = re.search(r"'([^']+)'\s*->\s*'([^']+)'", err_msg) if m: src_file = m.group(1) src_dir = os.path.dirname(src_file) self.logger.info( f"Attempting to copy cached snapshot from {src_dir} to {dest_dir}" ) shutil.copytree(src_dir, dest_dir, dirs_exist_ok=True) # Retry loading from the local copied directory try: self._embedding_model = EncoderClassifier.from_hparams( source=dest_dir, savedir=dest_dir, run_opts={"device": self.device}, ) self.logger.info( "Embedding model loaded successfully (after copying cached snapshot)" ) self._embedding_model_is_speechbrain = True return except Exception as e4: err_msg = str(e4) # As a last resort, try to download model files directly into dest_dir using huggingface_hub APIs from huggingface_hub import hf_hub_download, list_repo_files self.logger.info( f"Attempting direct HF download into {dest_dir} to avoid symlinks" ) os.makedirs(dest_dir, exist_ok=True) files = list_repo_files(self.config.embedding_model_id) for fname in files: if fname.endswith("/"): continue hf_hub_download( repo_id=self.config.embedding_model_id, filename=fname, local_dir=dest_dir, local_dir_use_symlinks=False, ) # Retry loading now that files are locally present self._embedding_model = EncoderClassifier.from_hparams( source=dest_dir, savedir=dest_dir, run_opts={"device": self.device}, ) self.logger.info( "Embedding model loaded successfully (after direct HF download)" ) self._embedding_model_is_speechbrain = True return except Exception as e5: err_msg = str(e5) self.logger.warning( "Common fixes: install a compatible torchaudio (matching your PyTorch), and install 'soundfile' or enable 'sox_io' backend." ) # If user allows fallback, provide MFCC fallback; otherwise raise an error to enforce SpeechBrain usage if getattr(self.config, "allow_fallback", False): self.logger.warning( "Falling back to MFCC-based deterministic embeddings (will be less accurate)." ) self._embedding_model = "FALLBACK" self._fallback_extractor = self._mfcc_embedding return else: raise RuntimeError( "Failed to load SpeechBrain embedding model and 'allow_fallback' is False. " "Ensure torchaudio and speechbrain are installed, or set 'allow_fallback=True' in DiarizationConfig." ) except Exception: # Import of SpeechBrain failed entirely; honor allow_fallback setting self.logger.warning( "Could not import SpeechBrain; checking 'allow_fallback' setting" ) if getattr(self.config, "allow_fallback", False): self.logger.warning( "Falling back to MFCC-based deterministic embeddings (allow_fallback=True)" ) self._embedding_model = "FALLBACK" self._fallback_extractor = self._mfcc_embedding else: raise RuntimeError( "Failed to import or initialize SpeechBrain embedding model and 'allow_fallback' is False. " "Install SpeechBrain or set 'allow_fallback=True' in DiarizationConfig to allow deterministic fallback." ) def _mfcc_embedding( self, segment_np: np.ndarray, sample_rate: int, target_dim: int = 192 ) -> np.ndarray: """Compute a deterministic embedding from audio segment using MFCCs. Falls back to simple waveform statistics if librosa is not available. Returns a fixed-size vector of length `target_dim`. """ try: import librosa mfcc = librosa.feature.mfcc(y=segment_np, sr=sample_rate, n_mfcc=40) mfcc_mean = mfcc.mean(axis=1) mfcc_std = mfcc.std(axis=1) vec = np.concatenate([mfcc_mean, mfcc_std]) except Exception: # Minimal deterministic fallback: use downsampled waveform statistics + spectral centroid approximation vec = [] vec.append(np.mean(segment_np)) vec.append(np.std(segment_np)) # simple spectral centroid proxy freqs = np.fft.rfftfreq(len(segment_np), d=1.0 / sample_rate) spec = np.abs(np.fft.rfft(segment_np)) if spec.sum() > 0: centroid = float((freqs * spec).sum() / spec.sum()) / (sample_rate / 2) else: centroid = 0.0 vec.append(centroid) vec = np.array(vec, dtype=float) # Pad or trim to target_dim if len(vec) < target_dim: padded = np.zeros(target_dim, dtype=float) padded[: len(vec)] = vec vec = padded elif len(vec) > target_dim: vec = vec[:target_dim] # normalize norm = np.linalg.norm(vec) + 1e-12 return (vec / norm).astype(np.float32) def process( self, waveform: torch.Tensor, sample_rate: int = 16000, num_speakers: Optional[int] = None, cache_dir: Optional[str] = None, audio_id: Optional[str] = None, fast_mode: bool = False, ) -> List[SpeakerSegment]: """ Main diarization pipeline. Args: waveform: Audio waveform [1, T] sample_rate: Audio sample rate num_speakers: Known number of speakers (auto-detect if None) Returns: List of SpeakerSegment with speaker assignments """ self._load_embedding_model() # Step 1: Voice Activity Detection speech_regions = self._detect_speech(waveform, sample_rate) if not speech_regions: self.logger.warning("No speech detected in audio") return [] self.logger.info(f"Detected {len(speech_regions)} speech regions") # Step 2: Create analysis windows windows = self._create_windows(speech_regions) if not windows: self.logger.warning("No valid windows created") return [] self.logger.info(f"Created {len(windows)} analysis windows") # Step 3: Extract speaker embeddings embeddings = self._extract_embeddings(waveform, windows, sample_rate) self.logger.info(f"Extracted embeddings with shape: {embeddings.shape}") # Step 4: Cluster embeddings labels = self._cluster_embeddings( embeddings, num_speakers=num_speakers or self.config.max_speakers ) num_speakers_found = len(set(labels)) self.logger.info(f"Found {num_speakers_found} speakers") # Step 5: Create segments from windows and labels raw_segments = self._create_segments(windows, labels, embeddings) # Step 6: Post-processing processed_segments = self._postprocess_segments(raw_segments) # Step 7: Detect overlapping speech processed_segments = self._detect_overlaps(processed_segments) self.logger.info(f"Final: {len(processed_segments)} segments") return processed_segments def auto_tune( self, waveform: torch.Tensor, sample_rate: int = 16000, num_speakers: Optional[int] = None ) -> dict: """Auto-tune clustering-related hyperparameters by searching simple parameter grid. This method extracts embeddings and tries different clustering thresholds and minimum cluster sizes, scoring candidates by silhouette score (and closeness to `num_speakers` if provided). The best parameter set is applied to `self.config` and returned for inspection. """ # Quick extraction path speech_regions = self._detect_speech(waveform, sample_rate) if not speech_regions: self.logger.warning("Auto-tune: no speech regions detected; aborting tuning") return {} windows = self._create_windows(speech_regions) if not windows: self.logger.warning("Auto-tune: no analysis windows created; aborting tuning") return {} embeddings = self._extract_embeddings(waveform, windows, sample_rate) if embeddings is None or len(embeddings) < 4: self.logger.warning("Auto-tune: insufficient embeddings for tuning; aborting tuning") return {} # Parameter grid (coarse) clustering_thresholds = [0.95, 0.85, 0.7, 0.5, 0.3, 0.15] min_cluster_sizes = [1, 2, 3, 4] best_score = -1e9 best_params = { "clustering_threshold": self.config.clustering_threshold, "min_cluster_size": self.config.min_cluster_size, "iterative_merge_threshold": self.config.iterative_merge_threshold, } # Save original values to restore if needed orig_threshold = self.config.clustering_threshold orig_min_size = self.config.min_cluster_size orig_iter_thresh = self.config.iterative_merge_threshold try: for thr in clustering_thresholds: for msize in min_cluster_sizes: # Temporarily set self.config.clustering_threshold = thr self.config.min_cluster_size = msize try: labels = self._cluster_embeddings(embeddings, num_speakers=None) k = len(np.unique(labels)) if k <= 1: sil = 0.0 else: try: sil = silhouette_score(embeddings, labels, metric="cosine") except Exception: sil = 0.0 # Scoring: prefer higher silhouette and closeness to desired num_speakers score = sil if num_speakers is not None: score -= 0.1 * abs(k - num_speakers) # small penalty for many clusters score -= 0.02 * k self.logger.debug( f"Auto-tune candidate: thr={thr}, min_size={msize} -> k={k}, sil={sil:.4f}, score={score:.4f}" ) if score > best_score: best_score = score best_params = { "clustering_threshold": thr, "min_cluster_size": msize, "achieved_k": k, "silhouette": sil, } except Exception as e: self.logger.debug(f"Auto-tune candidate failed: {e}") continue # Apply best params self.config.clustering_threshold = float( best_params.get("clustering_threshold", orig_threshold) ) self.config.min_cluster_size = int(best_params.get("min_cluster_size", orig_min_size)) # If a desired num_speakers was provided, set target merge accordingly if num_speakers is not None: self.config.target_num_speakers = int(num_speakers) self.logger.info(f"Auto-tune selected: {best_params}") return best_params finally: # nothing to restore; we've intentionally applied best params pass def _detect_speech(self, waveform: torch.Tensor, sample_rate: int) -> List[Tuple[float, float]]: """ Detect speech regions using energy-based VAD. Args: waveform: Audio waveform sample_rate: Sample rate Returns: List of (start, end) tuples for speech regions """ waveform_np = waveform.squeeze().cpu().numpy() # Frame parameters frame_length_ms = 25 # 25ms frames hop_length_ms = 10 # 10ms hop frame_length = int(frame_length_ms * sample_rate / 1000) hop_length = int(hop_length_ms * sample_rate / 1000) # Calculate energy per frame num_frames = max(1, 1 + (len(waveform_np) - frame_length) // hop_length) energies = np.zeros(num_frames) for i in range(num_frames): start_idx = i * hop_length end_idx = min(start_idx + frame_length, len(waveform_np)) frame = waveform_np[start_idx:end_idx] if len(frame) > 0: energies[i] = np.sqrt(np.mean(frame**2) + 1e-10) # Compute adaptive threshold if len(energies) > 0: energy_sorted = np.sort(energies) # Use 30th percentile as noise floor estimate noise_floor = energy_sorted[int(0.3 * len(energy_sorted))] threshold = noise_floor + self.config.vad_threshold * np.std(energies) else: threshold = self.config.vad_threshold # Find speech regions is_speech = energies > threshold # Apply morphological operations to smooth # (simple dilation and erosion using convolution) kernel_size = max(1, int(self.config.min_speech_duration * 1000 / hop_length_ms)) if kernel_size > 1 and len(is_speech) > kernel_size: # Simple smoothing kernel = np.ones(kernel_size) / kernel_size smoothed = np.convolve(is_speech.astype(float), kernel, mode="same") is_speech = smoothed > 0.5 # Convert to time regions regions = [] in_speech = False speech_start = 0.0 for i, speech in enumerate(is_speech): time = i * hop_length / sample_rate if speech and not in_speech: speech_start = time in_speech = True elif not speech and in_speech: duration = time - speech_start if duration >= self.config.min_speech_duration: regions.append((speech_start, time)) in_speech = False # Handle last region if in_speech: end_time = len(waveform_np) / sample_rate duration = end_time - speech_start if duration >= self.config.min_speech_duration: regions.append((speech_start, end_time)) # Merge nearby regions regions = self._merge_nearby_regions(regions, self.config.min_silence_duration) return regions def _merge_nearby_regions( self, regions: List[Tuple[float, float]], min_gap: float ) -> List[Tuple[float, float]]: """Merge regions that are close together""" if not regions: return [] merged = [regions[0]] for start, end in regions[1:]: last_start, last_end = merged[-1] if start - last_end <= min_gap: merged[-1] = (last_start, end) else: merged.append((start, end)) return merged def _create_windows( self, speech_regions: List[Tuple[float, float]] ) -> List[Tuple[float, float]]: """Create sliding windows over speech regions for embedding extraction""" windows = [] for region_start, region_end in speech_regions: t = region_start while t < region_end: window_end = min(t + self.config.segment_window, region_end) # Only include windows with sufficient duration if (window_end - t) >= self.config.min_segment_duration: # Avoid creating too many tiny windows across short recordings if (region_end - region_start) < (self.config.segment_window * 2): # for short regions, use a single window covering the region windows.append((region_start, region_end)) break windows.append((t, window_end)) t += self.config.segment_hop return windows def _extract_embeddings( self, waveform: torch.Tensor, windows: List[Tuple[float, float]], sample_rate: int, cache_dir: Optional[str] = None, audio_id: Optional[str] = None, fast_mode: bool = False, ) -> np.ndarray: """Extract speaker embeddings for each window. Optimizations implemented: - Disk cache (if enabled in config and cache_dir provided) - Batch extraction using model's batch API when available - Fast MFCC embedding path when `use_fast_embedding` is True """ # Try disk cache first if ( cache_dir and audio_id and self.config.embedding_cache and getattr(self.config, "embedding_cache", True) ): try: import os cache_path = Path(cache_dir) / f"{audio_id}_embeddings.npy" if cache_path.exists(): arr = np.load(str(cache_path)) if arr.shape[0] == len(windows): self.logger.info(f"Loaded embeddings from cache: {cache_path}") return arr except Exception: pass n = len(windows) embeddings = [None] * n # If fallback or user requested fast embedding, compute MFCC-based embeddings vectorized if ( (self._embedding_model == "FALLBACK" or self._embedding_model is None) or getattr(self.config, "use_fast_embedding", False) or fast_mode ): for i, (start, end) in enumerate(windows): start_sample = int(start * sample_rate) end_sample = int(end * sample_rate) segment = waveform[:, start_sample:end_sample] try: seg_np = segment.squeeze().cpu().numpy() emb = self._fallback_extractor(seg_np, sample_rate) except Exception: seg_np = segment.squeeze().cpu().numpy() emb = self._mfcc_embedding(seg_np, sample_rate) embeddings[i] = emb embeddings = np.stack(embeddings, axis=0) # Save to cache try: if cache_dir and audio_id and self.config.embedding_cache: Path(cache_dir).mkdir(parents=True, exist_ok=True) np.save(str(Path(cache_dir) / f"{audio_id}_embeddings.npy"), embeddings) except Exception: pass return embeddings # Otherwise use model batch encoding when available batch_size = max(1, int(getattr(self.config, "embedding_batch_size", 32))) # Prepare segment numpy arrays segs = [] seg_indices = [] for i, (start, end) in enumerate(windows): start_sample = int(start * sample_rate) end_sample = int(end * sample_rate) segment = waveform[:, start_sample:end_sample] segs.append(segment) seg_indices.append(i) # Try batch processing try: # If model supports encode_batch on a list or stacked tensor, process in chunks for i in range(0, len(segs), batch_size): batch = segs[i : i + batch_size] # Stack into a tensor batch try: batch_tensor = torch.stack( [b.squeeze(0) if b.dim() == 2 else b for b in batch], dim=0 ) except Exception: # Some models expect list of tensors; keep as list batch_tensor = batch with torch.no_grad(): try: # Move to model device if available if hasattr(self._embedding_model, "device") and isinstance( batch_tensor, torch.Tensor ): batch_tensor = batch_tensor.to(self._embedding_model.device) out = None # Try the most common batch API names if hasattr(self._embedding_model, "encode_batch"): out = self._embedding_model.encode_batch(batch_tensor) elif hasattr(self._embedding_model, "encode"): out = self._embedding_model.encode(batch_tensor) else: # fallback: try to call on each separately out = [self._embedding_model.encode_batch(x) for x in batch] # Normalize outputs into numpy array if isinstance(out, torch.Tensor): out_np = out.cpu().numpy() elif isinstance(out, list): out_np = np.stack( [ ( o.squeeze().cpu().numpy() if isinstance(o, torch.Tensor) else np.array(o) ) for o in out ], axis=0, ) else: out_np = np.array(out) # assign back to embeddings for j, idx in enumerate(range(i, i + out_np.shape[0])): embeddings[idx] = out_np[j] except Exception as e: # fallback to per-segment extraction for this batch self.logger.debug(f"Batch embedding failed, falling back per-segment: {e}") for bb_idx, seg in enumerate(batch): try: with torch.no_grad(): if hasattr(self._embedding_model, "device") and isinstance( seg, torch.Tensor ): seg = seg.to(self._embedding_model.device) emb = self._embedding_model.encode_batch(seg) emb = emb.squeeze().cpu().numpy() except Exception: emb = np.random.randn(192).astype(np.float32) embeddings[i + bb_idx] = emb embeddings = np.stack(embeddings, axis=0) # Save to cache try: if cache_dir and audio_id and self.config.embedding_cache: Path(cache_dir).mkdir(parents=True, exist_ok=True) np.save(str(Path(cache_dir) / f"{audio_id}_embeddings.npy"), embeddings) except Exception: pass return embeddings except Exception as e: self.logger.warning(f"Batch embedding extraction failed: {e}") # final fallback: single extraction loop embeddings = [] for start, end in windows: start_sample = int(start * sample_rate) end_sample = int(end * sample_rate) segment = waveform[:, start_sample:end_sample] try: with torch.no_grad(): if hasattr(self._embedding_model, "device"): segment = segment.to(self._embedding_model.device) emb = self._embedding_model.encode_batch(segment) emb = emb.squeeze().cpu().numpy() except Exception: emb = np.random.randn(192).astype(np.float32) embeddings.append(emb) embeddings = np.stack(embeddings, axis=0) return embeddings def _cluster_embeddings( self, embeddings: np.ndarray, num_speakers: Optional[int] = None, method_override: Optional[str] = None ) -> np.ndarray: """Cluster embeddings to assign speaker labels, with small-cluster merging. Args: embeddings: (N, D) array of embeddings num_speakers: Optional target number of speakers method_override: If set, use this clustering method ('agglomerative','spectral','kmeans') """ if len(embeddings) < 2: return np.zeros(len(embeddings), dtype=int) # Normalize embeddings scaler = StandardScaler() embeddings_norm = scaler.fit_transform(embeddings) # Support both nested (Config.diarization.clustering) and flat config shapes if method_override is not None: method = method_override # default thresholds - allow config overrides below threshold = getattr(self.config, "clustering_threshold", 0.7) linkage = getattr(self.config, "clustering_linkage", "average") min_size_cfg = getattr(self.config, "min_cluster_size", 2) max_speakers_cfg = getattr(self.config, "max_speakers", None) elif hasattr(self.config, "clustering"): method = self.config.clustering.method threshold = self.config.clustering.threshold linkage = self.config.clustering.linkage min_size_cfg = getattr( self.config.clustering, "min_cluster_size", getattr(self.config, "min_cluster_size", 2), ) max_speakers_cfg = getattr(self.config, "max_speakers", None) else: method = getattr(self.config, "clustering_method", "spectral") threshold = getattr(self.config, "clustering_threshold", 0.7) linkage = getattr(self.config, "clustering_linkage", "average") min_size_cfg = getattr(self.config, "min_cluster_size", 2) max_speakers_cfg = getattr(self.config, "max_speakers", None) if method == "agglomerative": if num_speakers is not None: clustering = AgglomerativeClustering( n_clusters=num_speakers, metric="cosine", linkage=linkage ) else: # If no target provided, estimate number of speakers via silhouette search est_max = min(8, max(2, len(embeddings) // 2)) est_min = 2 best_k = None best_score = -1.0 # Only try silhouette search on reasonably-sized inputs if len(embeddings) >= 8: for k in range(est_min, est_max + 1): try: tmp = AgglomerativeClustering(n_clusters=k, metric="cosine", linkage=linkage) labels_tmp = tmp.fit_predict(embeddings_norm) # silhouette requires at least 2 clusters and < n_samples clusters if len(np.unique(labels_tmp)) > 1 and len(np.unique(labels_tmp)) < len(embeddings): score = silhouette_score(embeddings_norm, labels_tmp, metric="cosine") else: score = -1.0 except Exception: score = -1.0 if score > best_score: best_score = score best_k = k # If silhouette search found a sensible k use it; else fallback to threshold style if best_k is not None and best_score > 0.01: clustering = AgglomerativeClustering(n_clusters=best_k, metric="cosine", linkage=linkage) self.logger.info(f"Agglomerative autodetected k={best_k} (silhouette={best_score:.3f})") else: clustering = AgglomerativeClustering( n_clusters=None, distance_threshold=threshold, metric="cosine", linkage=linkage, ) elif method == "spectral": n_clusters = num_speakers or min(8, len(embeddings) // 2) clustering = SpectralClustering( n_clusters=n_clusters, affinity="nearest_neighbors", n_neighbors=min(10, len(embeddings) - 1), ) elif method == "kmeans": n_clusters = num_speakers or min(8, len(embeddings) // 2) clustering = KMeans(n_clusters=n_clusters, random_state=42, n_init=10) else: raise ValueError(f"Unknown clustering method: {method}") try: labels = clustering.fit_predict(embeddings_norm) except Exception as e: self.logger.error(f"Clustering failed: {e}") labels = np.array([i % 2 for i in range(len(embeddings))]) # Debug: cluster sizes unique, counts = np.unique(labels, return_counts=True) sizes = dict(zip(unique.tolist(), counts.tolist())) self.logger.debug(f"Initial clusters: {len(unique)}, sizes: {sizes}") # Global check: if all embeddings are very similar, collapse directly to 1 speaker try: # First, perform a row-normalized (per-embedding) cosine check on raw embeddings row_norm = embeddings / (np.linalg.norm(embeddings, axis=1, keepdims=True) + 1e-12) n_sample = min(200, len(row_norm)) idx = np.linspace(0, len(row_norm) - 1, n_sample).astype(int) sub = row_norm[idx] sims = np.dot(sub, sub.T) sims = np.clip(sims, -1.0, 1.0) dists = 1.0 - sims mean_row_dist = ( float(np.mean(dists[np.triu_indices_from(dists, k=1)])) if n_sample > 1 else 1.0 ) global_row_threshold = getattr(self.config, "global_collapse_threshold", 0.03) # Be more permissive for short recordings (few windows) if len(embeddings) < 40: global_row_threshold = max(global_row_threshold, 0.08) if mean_row_dist < global_row_threshold: self.logger.info( f"Row-normalized embeddings too similar (mean dist={mean_row_dist:.6f}), collapsing to 1 speaker" ) return np.zeros(len(embeddings), dtype=int) # Next, check on scaled embeddings (existing logic) n_sample = min(200, len(embeddings_norm)) idx = np.linspace(0, len(embeddings_norm) - 1, n_sample).astype(int) sub = embeddings_norm[idx] sims = np.dot(sub, sub.T) sims = np.clip(sims, -1.0, 1.0) dists = 1.0 - sims mean_global_dist = ( float(np.mean(dists[np.triu_indices_from(dists, k=1)])) if n_sample > 1 else 1.0 ) global_collapse_threshold = getattr(self.config, "global_collapse_threshold", 0.03) if mean_global_dist < global_collapse_threshold: self.logger.info( f"Global embeddings too similar (mean dist={mean_global_dist:.4f}), collapsing to 1 speaker" ) return np.zeros(len(embeddings), dtype=int) # Additional small-variance heuristic: if feature-wise std is tiny, collapse as well mean_std = float(np.mean(np.std(embeddings_norm, axis=0))) std_threshold = getattr(self.config, "global_std_threshold", 1e-2) if mean_std < std_threshold: self.logger.info( f"Embeddings have tiny variance (mean std={mean_std:.6f}), collapsing to 1 speaker" ) return np.zeros(len(embeddings), dtype=int) except Exception: pass # If centroids are very close to each other, this is likely a single-speaker recording. # Compute mean pairwise centroid cosine distance; if below a threshold, collapse to 1 cluster. try: labels_unique = np.unique(labels) centroids = [embeddings_norm[labels == l].mean(axis=0) for l in labels_unique] if len(centroids) > 1: pair_dists = [] for i in range(len(centroids)): for j in range(i + 1, len(centroids)): a = centroids[i] / (np.linalg.norm(centroids[i]) + 1e-12) b = centroids[j] / (np.linalg.norm(centroids[j]) + 1e-12) pair_dists.append(1.0 - float(np.dot(a, b))) mean_pair_dist = float(np.mean(pair_dists)) if pair_dists else 1.0 else: mean_pair_dist = 1.0 collapse_threshold = getattr(self.config, "collapse_threshold", 0.15) if mean_pair_dist < collapse_threshold: self.logger.info( f"Centroids too similar (mean dist={mean_pair_dist:.3f}), collapsing to 1 speaker" ) labels = np.zeros_like(labels) # If SpeechBrain embeddings are used and clusters have a very low silhouette score, # it's likely that the recording is single-speaker and clustering is over-fragmenting. try: if getattr(self.config, "use_speechbrain", True) and getattr( self, "_embedding_model_is_speechbrain", False ): unique_labels = np.unique(labels) if len(unique_labels) > 1: try: score = silhouette_score(embeddings_norm, labels, metric="cosine") if score < getattr(self.config, "silhouette_collapse_threshold", 0.05): self.logger.info( f"Low silhouette score ({score:.4f}) detected with SpeechBrain embeddings; collapsing to 1 speaker" ) return np.zeros(len(embeddings), dtype=int) except Exception: pass except Exception: pass except Exception: pass # Merge clusters smaller than min_cluster_size min_size = min_size_cfg if min_size and min_size > 1: changed = True while changed: changed = False labels_unique, label_counts = np.unique(labels, return_counts=True) small_labels = [l for l, c in zip(labels_unique, label_counts) if c < min_size] if not small_labels: break # compute centroids for existing labels centroids = {l: embeddings_norm[labels == l].mean(axis=0) for l in labels_unique} for sl in small_labels: candidates = [l for l in labels_unique if l != sl] if not candidates: continue # find nearest centroid (cosine distance) def cosine_dist(a, b): a_norm = a / (np.linalg.norm(a) + 1e-12) b_norm = b / (np.linalg.norm(b) + 1e-12) return 1.0 - float(np.dot(a_norm, b_norm)) distances = [(c, cosine_dist(centroids[sl], centroids[c])) for c in candidates] nearest = min(distances, key=lambda x: x[1])[0] # reassign labels labels[labels == sl] = nearest changed = True # Final cluster sizes unique2, counts2 = np.unique(labels, return_counts=True) sizes2 = dict(zip(unique2.tolist(), counts2.tolist())) self.logger.debug(f"Clusters after merge: {len(unique2)}, sizes: {sizes2}") # Additional centroid-based merging: merge clusters whose centroids are very close try: labels_unique = np.unique(labels) centroids = {l: embeddings_norm[labels == l].mean(axis=0) for l in labels_unique} # compute pairwise centroid distances pairs = [] for i, a in enumerate(labels_unique): for j, b in enumerate(labels_unique): if j <= i: continue dist = 1.0 - float( np.dot( centroids[a] / (np.linalg.norm(centroids[a]) + 1e-12), centroids[b] / (np.linalg.norm(centroids[b]) + 1e-12), ) ) pairs.append((dist, a, b)) # merge pairs with distance < threshold pairs.sort() merged = False for dist, a, b in pairs: if dist < threshold: # merge b into a labels[labels == b] = a merged = True if merged: labels_unique2, counts2 = np.unique(labels, return_counts=True) sizes2 = dict(zip(labels_unique2.tolist(), counts2.tolist())) self.logger.debug( f"Clusters after centroid-merge: {len(labels_unique2)}, sizes: {sizes2}" ) # Iterative silhouette-guided merging: try merging closest centroid pairs while it improves or meets configured criteria try: iterative_thresh = getattr(self.config, "iterative_merge_threshold", threshold) silhouette_min = getattr(self.config, "iterative_merge_silhouette_threshold", 0.0) max_merge_iters = getattr(self.config, "iterative_merge_max_iters", 10) def compute_centroids(curr_labels): uniq = np.unique(curr_labels) return {l: embeddings_norm[curr_labels == l].mean(axis=0) for l in uniq} def pairwise_min_pair(centroids_dict): uniq = list(centroids_dict.keys()) best = (1.0, None, None) for i, a in enumerate(uniq): for j in range(i + 1, len(uniq)): b = uniq[j] a_c = centroids_dict[a] / (np.linalg.norm(centroids_dict[a]) + 1e-12) b_c = centroids_dict[b] / (np.linalg.norm(centroids_dict[b]) + 1e-12) dist = 1.0 - float(np.dot(a_c, b_c)) if dist < best[0]: best = (dist, a, b) return best curr_labels = labels.copy() prev_score = None try: if len(np.unique(curr_labels)) > 1: prev_score = silhouette_score(embeddings_norm, curr_labels, metric="cosine") except Exception: prev_score = None iters = 0 while iters < max_merge_iters: iters += 1 cent = compute_centroids(curr_labels) if len(cent) <= 1: break min_dist, a, b = pairwise_min_pair(cent) if min_dist >= iterative_thresh: break # simulate merge and evaluate silhouette next_labels = curr_labels.copy() next_labels[next_labels == b] = a try: if len(np.unique(next_labels)) > 1: next_score = silhouette_score( embeddings_norm, next_labels, metric="cosine" ) else: next_score = 1.0 except Exception: next_score = None accept = False if next_score is not None: if prev_score is None: # accept merges that meet a minimum silhouette threshold if next_score >= silhouette_min: accept = True else: # accept if silhouette improves by a small margin or stays acceptable if next_score >= prev_score or next_score >= silhouette_min: accept = True if accept: curr_labels = next_labels prev_score = next_score labels = curr_labels.copy() # continue iterating else: break if iters > 1: labels_unique2, counts2 = np.unique(labels, return_counts=True) sizes2 = dict(zip(labels_unique2.tolist(), counts2.tolist())) self.logger.debug( f"Clusters after iterative-merge (iters={iters}): {len(labels_unique2)}, sizes: {sizes2}" ) # If user requested a target speaker count, greedily merge closest centroid pairs until we meet it try: target_k = getattr(self.config, "target_num_speakers", None) force_thresh = float(getattr(self.config, "target_force_threshold", 1.0)) if target_k is not None: curr_labels = labels.copy() def compute_centroids(curr): uniq = np.unique(curr) return {l: embeddings_norm[curr == l].mean(axis=0) for l in uniq} merged_iters = 0 while len(np.unique(curr_labels)) > target_k: cent = compute_centroids(curr_labels) if len(cent) <= 1: break # find closest pair uniq = list(cent.keys()) best = (1.0, None, None) for i, a in enumerate(uniq): for j in range(i + 1, len(uniq)): b = uniq[j] a_c = cent[a] / (np.linalg.norm(cent[a]) + 1e-12) b_c = cent[b] / (np.linalg.norm(cent[b]) + 1e-12) dist = 1.0 - float(np.dot(a_c, b_c)) if dist < best[0]: best = (dist, a, b) min_dist, a, b = best # if min_dist is too large and force_thresh < 1.0, break if min_dist > force_thresh and force_thresh < 1.0: self.logger.warning( f"Stopping target-merge early: nearest cluster dist {min_dist:.3f} > force_thresh {force_thresh}" ) break # merge b into a curr_labels[curr_labels == b] = a merged_iters += 1 # safety to avoid infinite loops if merged_iters > 1000: break if merged_iters: labels = curr_labels.copy() labels_unique2, counts2 = np.unique(labels, return_counts=True) sizes2 = dict(zip(labels_unique2.tolist(), counts2.tolist())) self.logger.info( f"Clusters after target-merge (target={target_k}, iters={merged_iters}): {len(labels_unique2)}, sizes: {sizes2}" ) except Exception: pass except Exception: # don't let merging errors break the pipeline pass # Heuristic fallback: if still too fragmented, run KMeans with estimated speaker count n_clusters_found = len(np.unique(labels)) max_allowed = 20 if n_clusters_found > max_allowed: est_k = min(12, max(2, int(len(embeddings) / 80))) self.logger.warning( f"Too many clusters ({n_clusters_found}), falling back to KMeans with k={est_k}" ) try: km = KMeans(n_clusters=est_k, random_state=42, n_init=10) labels = km.fit_predict(embeddings_norm) # Re-merge small clusters after KMeans labels_unique2, counts2 = np.unique(labels, return_counts=True) sizes2 = dict(zip(labels_unique2.tolist(), counts2.tolist())) self.logger.info( f"Clusters after KMeans fallback: {len(labels_unique2)}, sizes: {sizes2}" ) except Exception as e: self.logger.error(f"KMeans fallback failed: {e}") except Exception: pass return labels def _create_segments( self, windows: List[Tuple[float, float]], labels: np.ndarray, embeddings: np.ndarray ) -> List[SpeakerSegment]: """Create SpeakerSegment objects from windows and labels""" segments = [] for (start, end), label, emb in zip(windows, labels, embeddings): segments.append( SpeakerSegment( speaker_id=f"SPEAKER_{label:02d}", start=start, end=end, confidence=1.0, embedding=emb, ) ) # If we used the fallback extractor, update segment embeddings to the deterministic MFCC embeddings if getattr(self, "_fallback_extractor", None) is not None: try: for i, seg in enumerate(segments): # reuse windows to create a deterministic embedding s, e = windows[i] # external code expects embeddings array, but ensure segment.embedding is deterministic if ( segments[i].embedding is None or isinstance(self._embedding_model, str) and self._embedding_model == "FALLBACK" ): # compute on-demand using fallback extractor seg_np = self._extract_waveform_segment(windows[i]) segments[i].embedding = self._fallback_extractor(seg_np, sample_rate) except Exception: pass return segments def _postprocess_segments(self, segments: List[SpeakerSegment]) -> List[SpeakerSegment]: """Post-process segments: merge adjacent, filter short""" if not segments: return [] # Sort by start time segments = sorted(segments, key=lambda x: x.start) # Merge adjacent segments from same speaker merged = [segments[0]] for seg in segments[1:]: last = merged[-1] gap = seg.start - last.end if seg.speaker_id == last.speaker_id and gap <= self.config.merge_gap_threshold: # Merge: extend last segment last.end = max(last.end, seg.end) last.confidence = (last.confidence + seg.confidence) / 2 else: merged.append(seg) # Smoothing: fix short isolated segments between identical speakers smoothed = merged if len(smoothed) >= 3: changed = False for i in range(1, len(smoothed) - 1): seg = smoothed[i] prev = smoothed[i - 1] nxt = smoothed[i + 1] threshold = max(1.0, self.config.min_segment_duration) if seg.duration < threshold and prev.speaker_id == nxt.speaker_id: seg.speaker_id = prev.speaker_id changed = True if changed: # merge again after smoothing merged2 = [smoothed[0]] for seg in smoothed[1:]: last = merged2[-1] gap = seg.start - last.end if seg.speaker_id == last.speaker_id and gap <= self.config.merge_gap_threshold: last.end = max(last.end, seg.end) last.confidence = (last.confidence + seg.confidence) / 2 else: merged2.append(seg) merged = merged2 # Filter short segments filtered = [seg for seg in merged if seg.duration >= self.config.min_segment_duration] return filtered def _merge_segments( self, segments: List[SpeakerSegment], max_gap: float = 0.5 ) -> List[SpeakerSegment]: """Compatibility helper: merge adjacent segments from same speaker within max_gap""" if not segments: return [] segments = sorted(segments, key=lambda x: x.start) merged_list = [segments[0]] for seg in segments[1:]: last = merged_list[-1] gap = seg.start - last.end if seg.speaker_id == last.speaker_id and gap <= max_gap: # Merge: extend last segment last.end = max(last.end, seg.end) last.confidence = (last.confidence + seg.confidence) / 2 else: merged_list.append(seg) return merged_list def _detect_overlaps(self, segments: List[SpeakerSegment]) -> List[SpeakerSegment]: """Mark segments that overlap with other speakers""" for i, seg1 in enumerate(segments): for j, seg2 in enumerate(segments): if i != j and seg1.speaker_id != seg2.speaker_id: # Check for time overlap overlap_start = max(seg1.start, seg2.start) overlap_end = min(seg1.end, seg2.end) if overlap_start < overlap_end: seg1.is_overlap = True seg2.is_overlap = True return segments def get_speaker_stats(self, segments: List[SpeakerSegment]) -> Dict[str, Dict[str, float]]: """ Get statistics for each speaker. Returns: Dict mapping speaker_id to stats (total_duration, num_segments, etc.) """ stats = {} for seg in segments: if seg.speaker_id not in stats: stats[seg.speaker_id] = { "total_duration": 0.0, "num_segments": 0, "avg_segment_duration": 0.0, "overlap_duration": 0.0, } stats[seg.speaker_id]["total_duration"] += seg.duration stats[seg.speaker_id]["num_segments"] += 1 if seg.is_overlap: stats[seg.speaker_id]["overlap_duration"] += seg.duration # Calculate averages for speaker_id in stats: num_segs = stats[speaker_id]["num_segments"] if num_segs > 0: stats[speaker_id]["avg_segment_duration"] = ( stats[speaker_id]["total_duration"] / num_segs ) return stats