diff --git a/main.py b/main.py index 9395b2e..21027bb 100755 --- a/main.py +++ b/main.py @@ -1,8 +1,21 @@ #!/usr/bin/env python import argparse +import dataclasses -from pyannote.audio import Pipeline +import numpy as np +from sklearn.cluster import AgglomerativeClustering +from sklearn.preprocessing import normalize +from speechbrain.inference.speaker import EncoderClassifier +import torch +import torchaudio +import whisper + +@dataclasses.dataclass +class Segment: + start: np.float32 + end: np.float32 + text: str def srt_timestamp(milis): h = milis // (1000 * 60 * 60) @@ -19,18 +32,75 @@ def main(): parser.add_argument("-o", "--output", default="output.srt") args = parser.parse_args() - pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization-community-1", - token="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") - output = pipeline(args.input, num_speakers=2) - for turn, speaker in output.speaker_diarization: - print(f"{speaker} speaks between t={turn.start:.3f}s and t={turn.end:.3f}s") + model = whisper.load_model("turbo") + result = model.transcribe(args.input, word_timestamps=True, language="en") + segments = [] + for seg in result["segments"]: + words = seg["words"] + start = words[0]["start"] + end = words[-1]["end"] + segments.append(Segment(start, end, seg["text"].strip())) + + print("Transcription done. Starting speaker classification") + + classifier = EncoderClassifier.from_hparams( + source="speechbrain/spkrec-ecapa-voxceleb", + ) + + signal, sample_rate = torchaudio.load(args.input) + + print(sample_rate) + # resample to 16khz + if sample_rate != 16000: + resampler = torchaudio.transforms.Resample(sample_rate, 16000) + signal = resampler(signal) + sample_rate = 16000 + + # convert to mono if needed + if signal.shape[0] > 1: + signal = torch.mean(signal, dim=0, keepdim=True) + + embeddings = [] + for seg in segments: + start_sample = int(seg.start * sample_rate) + end_sample = int(seg.end * sample_rate) + + chunk = signal[:, start_sample:end_sample] + + # SpeechBrain expects [batch, time] + emb = classifier.encode_batch(chunk) + emb = emb.squeeze().detach().cpu().numpy() + + embeddings.append(emb) + embeddings = np.vstack(embeddings) + embeddings = normalize(embeddings) + + n_speakers = 2 + clustering = AgglomerativeClustering( + n_clusters=n_speakers, + metric="cosine", + linkage="average" + ) + labels = clustering.fit_predict(embeddings) + + for seg, label in zip(segments, labels): + seg.text = f"Speaker {label + 1}: {seg.text}" with open(args.output, "w") as srt_file: - for i, (turn, speaker) in enumerate(output.speaker_diarization, 1): - srt_file.write(f"{i}\n" - f"{srt_timestamp(int(turn.start * 1000))} --> " - f"{srt_timestamp(int(turn.end * 1000))}\n" - f"{speaker}\n\n") + for i, seg in enumerate(segments, 1): + start = srt_timestamp(round(seg.start * 1000)) + end = srt_timestamp(round(seg.end * 1000)) + srt_file.write(f"{i}\n{start} --> {end}\n{seg.text}\n\n") + print(start, end, seg.text) + + from sklearn.decomposition import PCA + import matplotlib.pyplot as plt + + pca = PCA(n_components=2) + points = pca.fit_transform(embeddings) + + plt.scatter(points[:,0], points[:,1]) + plt.show() if __name__ == "__main__": main() diff --git a/requirements.txt b/requirements.txt index 8de5ef5..8af18af 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,2 +1,6 @@ -pyannote-audio==4.0.4 +numpy==2.4.6 openai-whisper==20250625 +scikit-learn==1.8.0 +speechbrain==1.1.0 +torch==2.12.0 +torchaudio==2.11.0