Spaces:
Running on Zero
Running on Zero
| """ | |
| Local speaker diarization ("who spoke when"). | |
| Three steps: | |
| 1. sherpa-onnx (pyannote segmentation-3.0, ONNX) cuts the recording into | |
| single-speaker turns. | |
| 2. A neural speaker-embedding model turns each turn into a voice fingerprint. | |
| 3. The fingerprints are clustered into speakers. Only turns long enough to | |
| give a trustworthy fingerprint decide how many speakers there are; short | |
| interjections ("okay", "so") are then handed to the closest voice. | |
| Everything runs offline on CPU through ONNX Runtime; no Hugging Face token | |
| and no PyTorch. | |
| """ | |
| from __future__ import annotations | |
| import logging | |
| from typing import Callable, List, Optional, Sequence, Tuple | |
| import numpy as np | |
| from config.settings import CPU_THREADS, MAX_AUTO_SPEAKERS, SAME_VOICE_SIMILARITY | |
| from modules.audio_processor import SAMPLE_RATE | |
| from modules.model_downloader import ( | |
| EMBEDDING_MODEL, | |
| SEGMENTATION_MODEL, | |
| ensure_diarization_models, | |
| ) | |
| from modules.models import DiarizationSegment | |
| logger = logging.getLogger(__name__) | |
| # A turn shorter than this gives an unreliable fingerprint: it may not vote on | |
| # the number of speakers, it is only assigned to one afterwards. | |
| MIN_RELIABLE_SECONDS = 1.5 | |
| # Below this the embedding model has too little audio to produce anything. | |
| MIN_EMBED_SECONDS = 0.4 | |
| # With auto-detection, a cluster needs this much reliable speech to count as a | |
| # person: 2% of the recording's reliable speech, but never under 4 s or over 20 s. | |
| # A real participant who says less than that is merged into the closest voice; | |
| # choosing the speaker count at upload bypasses this. | |
| MIN_SPEAKER_SECONDS = 4.0 | |
| MIN_SPEAKER_SHARE = 0.02 | |
| MAX_FLOOR_SECONDS = 20.0 | |
| class DiarizationService: | |
| """Speaker diarization on 16 kHz mono float32 samples.""" | |
| def diarize( | |
| self, | |
| audio: np.ndarray, | |
| num_speakers: Optional[int] = None, | |
| on_progress: Optional[Callable[[float], None]] = None, | |
| ) -> List[DiarizationSegment]: | |
| """ | |
| Return speaker turns ordered by start time. | |
| *num_speakers* fixes the speaker count when known (most accurate). | |
| When omitted the count is estimated from how well the voices separate. | |
| """ | |
| import sherpa_onnx | |
| ensure_diarization_models() | |
| embedding_config = sherpa_onnx.SpeakerEmbeddingExtractorConfig( | |
| model=str(EMBEDDING_MODEL), | |
| num_threads=CPU_THREADS, | |
| ) | |
| # -- 1. turns. sherpa-onnx also clusters, with a fixed distance | |
| # threshold that over-splits badly on real recordings (16 "speakers" | |
| # for a two-person interview); only its turn boundaries are kept. | |
| config = sherpa_onnx.OfflineSpeakerDiarizationConfig( | |
| segmentation=sherpa_onnx.OfflineSpeakerSegmentationModelConfig( | |
| pyannote=sherpa_onnx.OfflineSpeakerSegmentationPyannoteModelConfig( | |
| model=str(SEGMENTATION_MODEL) | |
| ), | |
| num_threads=CPU_THREADS, | |
| ), | |
| embedding=embedding_config, | |
| clustering=sherpa_onnx.FastClusteringConfig(num_clusters=-1, threshold=0.5), | |
| # Ignore blips shorter than 0.3 s; bridge pauses shorter than 0.5 s. | |
| min_duration_on=0.3, | |
| min_duration_off=0.5, | |
| ) | |
| if not config.validate(): | |
| raise RuntimeError("Diarization models are missing or invalid.") | |
| def _callback(done: int, total: int) -> int: | |
| if on_progress and total: | |
| on_progress(0.7 * done / total) | |
| return 0 | |
| result = sherpa_onnx.OfflineSpeakerDiarization(config).process( | |
| audio, callback=_callback | |
| ).sort_by_start_time() | |
| spans = [(r.start, r.end) for r in result] | |
| if not spans: | |
| return [] | |
| # -- 2. one fingerprint per turn | |
| extractor = sherpa_onnx.SpeakerEmbeddingExtractor(embedding_config) | |
| embeddings: list[Optional[np.ndarray]] = [] | |
| for i, (start, end) in enumerate(spans): | |
| embeddings.append( | |
| _embed(extractor, audio[int(start * SAMPLE_RATE): int(end * SAMPLE_RATE)]) | |
| ) | |
| if on_progress: | |
| on_progress(0.7 + 0.3 * (i + 1) / len(spans)) | |
| # -- 3. speakers | |
| labels = cluster_speakers( | |
| [end - start for start, end in spans], embeddings, num_speakers | |
| ) | |
| turns = label_turns([(s, e, label) for (s, e), label in zip(spans, labels)]) | |
| talk_time: dict[str, float] = {} | |
| for t in turns: | |
| talk_time[t.speaker] = talk_time.get(t.speaker, 0.0) + t.end - t.start | |
| logger.info( | |
| "Diarization complete: %d turns, %d speakers (%s)", | |
| len(turns), len(talk_time), | |
| ", ".join(f"{name} {seconds:.0f}s" for name, seconds in talk_time.items()), | |
| ) | |
| return turns | |
| def _embed(extractor, samples: np.ndarray) -> Optional[np.ndarray]: | |
| """Unit-length voice fingerprint for *samples*, or None if there is too little audio.""" | |
| if len(samples) < MIN_EMBED_SECONDS * SAMPLE_RATE: | |
| return None | |
| stream = extractor.create_stream() | |
| stream.accept_waveform(sample_rate=SAMPLE_RATE, waveform=samples) | |
| stream.input_finished() | |
| if not extractor.is_ready(stream): | |
| return None | |
| embedding = np.asarray(extractor.compute(stream), dtype=np.float32) | |
| norm = np.linalg.norm(embedding) | |
| return embedding / norm if norm > 0 else None | |
| def cluster_speakers( | |
| durations: Sequence[float], | |
| embeddings: Sequence[Optional[np.ndarray]], | |
| num_speakers: Optional[int] = None, | |
| ) -> List[int]: | |
| """ | |
| Return a speaker index for every turn. | |
| Reliable turns (long enough for a trustworthy fingerprint) are clustered; | |
| every other turn joins the closest cluster, or, with no fingerprint at | |
| all, the speaker of the turn before it. | |
| """ | |
| from sklearn.cluster import AgglomerativeClustering | |
| n = len(durations) | |
| reliable = [ | |
| i for i in range(n) | |
| if embeddings[i] is not None and durations[i] >= MIN_RELIABLE_SECONDS | |
| ] | |
| if len(reliable) < 2: # very short recording: use whatever fingerprints exist | |
| reliable = [i for i in range(n) if embeddings[i] is not None] | |
| if len(reliable) < 2: | |
| return [0] * n | |
| fixed = bool(num_speakers and num_speakers > 0) | |
| while True: | |
| X = np.stack([embeddings[i] for i in reliable]) | |
| k = min(num_speakers if fixed else estimate_speaker_count(X), len(reliable)) | |
| if k <= 1: | |
| return [0] * n | |
| reliable_labels = AgglomerativeClustering( | |
| n_clusters=k, metric="cosine", linkage="average" | |
| ).fit_predict(X) | |
| if fixed: | |
| break | |
| # A cluster holding almost no speech is an odd-sounding turn (laughter, | |
| # applause, a change of microphone), not a person. Clustering isolates | |
| # such turns eagerly, so set them aside and count again without them. | |
| seconds = np.array([durations[i] for i in reliable]) | |
| floor = min(MAX_FLOOR_SECONDS, max(MIN_SPEAKER_SECONDS, MIN_SPEAKER_SHARE * seconds.sum())) | |
| minor = [c for c in range(k) if seconds[reliable_labels == c].sum() < floor] | |
| if not minor: | |
| break | |
| reliable = [i for i, label in zip(reliable, reliable_labels) if label not in minor] | |
| if len(reliable) < 2: | |
| return [0] * n | |
| if not fixed: | |
| reliable_labels = merge_same_voices(X, reliable_labels) | |
| k = int(reliable_labels.max()) + 1 | |
| if k <= 1: | |
| return [0] * n | |
| centroids = _centroids(X, reliable_labels) | |
| labels: list[int] = [] | |
| for i in range(n): | |
| if embeddings[i] is None: | |
| labels.append(labels[-1] if labels else -1) | |
| else: | |
| labels.append(int(np.argmax(centroids @ embeddings[i]))) | |
| for i, label in zip(reliable, reliable_labels): | |
| labels[i] = int(label) | |
| # Leading turns without a fingerprint take the first known speaker. | |
| first_known = next(label for label in labels if label != -1) | |
| return [first_known if label == -1 else label for label in labels] | |
| def _centroids(X: np.ndarray, labels: np.ndarray) -> np.ndarray: | |
| """Average voice of each cluster, unit length, indexed by label.""" | |
| centroids = np.stack([X[labels == c].mean(axis=0) for c in range(int(labels.max()) + 1)]) | |
| return centroids / np.linalg.norm(centroids, axis=1, keepdims=True) | |
| def merge_same_voices(X: np.ndarray, labels: np.ndarray) -> np.ndarray: | |
| """ | |
| Merge clusters whose average voices are the same person, and renumber. | |
| The silhouette score always finds some way to split a recording, even one | |
| with a single speaker. What tells the cases apart is how alike the | |
| resulting voices are: on test recordings, different people scored a | |
| similarity of 0.24 to 0.39 and one person split in two scored 0.92. | |
| """ | |
| labels = np.unique(labels, return_inverse=True)[1] | |
| while labels.max() > 0: | |
| similarity = _centroids(X, labels) @ _centroids(X, labels).T | |
| np.fill_diagonal(similarity, -1.0) | |
| a, b = np.unravel_index(np.argmax(similarity), similarity.shape) | |
| logger.info("Most alike pair of voices has similarity %.2f", similarity[a, b]) | |
| if similarity[a, b] < SAME_VOICE_SIMILARITY: | |
| break | |
| labels[labels == b] = a | |
| labels = np.unique(labels, return_inverse=True)[1] | |
| return labels | |
| def estimate_speaker_count(X: np.ndarray) -> int: | |
| """ | |
| Pick the number of clusters whose members sit closest to their own cluster | |
| and furthest from the others (best silhouette score). This never answers | |
| "one"; merge_same_voices() makes that call afterwards. | |
| """ | |
| from sklearn.cluster import AgglomerativeClustering | |
| from sklearn.metrics import silhouette_score | |
| scores: dict[int, float] = {} | |
| for k in range(2, min(MAX_AUTO_SPEAKERS, len(X) - 1) + 1): | |
| labels = AgglomerativeClustering( | |
| n_clusters=k, metric="cosine", linkage="average" | |
| ).fit_predict(X) | |
| scores[k] = float(silhouette_score(X, labels, metric="cosine")) | |
| logger.info( | |
| "Speaker count scores from %d reliable turns: %s", | |
| len(X), ", ".join(f"{k}={s:.2f}" for k, s in scores.items()), | |
| ) | |
| return max(scores, key=scores.get) if scores else 1 | |
| def label_turns(raw_turns: List[Tuple[float, float, int]]) -> List[DiarizationSegment]: | |
| """Name clusters 'Speaker 1', 'Speaker 2', ... in order of first appearance.""" | |
| numbers: dict[int, int] = {} | |
| turns: list[DiarizationSegment] = [] | |
| for start, end, cluster in raw_turns: | |
| number = numbers.setdefault(cluster, len(numbers) + 1) | |
| turns.append(DiarizationSegment(start=start, end=end, speaker=f"Speaker {number}")) | |
| return turns | |