whisper

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()