Update app.py

app.py

@@ -13,6 +13,22 @@ import json
 pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization", use_auth_token="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")
 from pydub.effects import speedup
 import moviepy.editor as mp
+import datetime
+import torch
+import pyannote.audio
+from pyannote.audio.pipelines.speaker_verification import PretrainedSpeakerEmbedding
+from pyannote.audio import Audio
+from pyannote.core import Segment
+import wave
+import contextlib
+from sklearn.cluster import AgglomerativeClustering
+import numpy as np
+
+model = whisper.load_model("medium")
+embedding_model = PretrainedSpeakerEmbedding(
+    "speechbrain/spkrec-ecapa-voxceleb",
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+)
 
 
 __FILES = set()
@@ -131,6 +147,86 @@ def Transcribe(NumberOfSpeakers, SpeakerNames="", audio="temp_audio.wav"):
     RemoveAllFiles()
     return (t_text, ({ "data": [{"speaker": speaker, "text": text} for speaker, text in conversation]}))
 
+
+def Transcribe_V2(num_speakers, speaker_names, audio="temp_audio.wav"):
+    SPEAKER_DICT = {}
+    SPEAKERS = [speaker.strip() for speaker in speaker_names.split(',')]
+    def GetSpeaker(sp):
+        speaker = sp
+        if sp not in list(SPEAKER_DICT.keys()):
+            if len(SPEAKERS):
+                t = SPEAKERS.pop(0)
+                SPEAKER_DICT[sp] = t
+                speaker = SPEAKER_DICT[sp]
+        else:
+            speaker = SPEAKER_DICT[sp]
+        return speaker
+    
+    # audio = Audio()
+    def get_output(segments):
+        # print(segments)
+        conversation=[]
+        for (i, segment) in enumerate(segments):
+            # print(f"{i}, {segment["speaker"]}, {segments[i - 1]["speaker"]}, {}")
+            if not len(conversation):
+                conversation.append([GetSpeaker(segment["speaker"]), segment["text"].lstrip()])
+            elif conversation[-1][0] == GetSpeaker(segment["speaker"]):
+                conversation[-1][1] +=  segment["text"].lstrip()
+            else:
+                conversation.append([GetSpeaker(segment["speaker"]), segment["text"].lstrip()])
+            # if i == 0 or segments[i - 1]["speaker"] != segment["speaker"]:
+            #     if i != 0:
+            #         conversation.append([GetSpeaker(segment["speaker"]), segment["text"][1:]]) # segment["speaker"] + ' ' + str(time(segment["start"])) + '\n\n'
+            # conversation[-1][1] += segment["text"][1:]
+        # return output
+        return ("".join([f"{speaker} --> {text}\n" for speaker, text in conversation])), conversation
+
+    def get_duration(path):
+        with contextlib.closing(wave.open(path,'r')) as f:
+            frames = f.getnframes()
+            rate = f.getframerate()
+        return frames / float(rate)
+
+    def make_embeddings(path, segments, duration):
+        embeddings = np.zeros(shape=(len(segments), 192))
+        for i, segment in enumerate(segments):
+            embeddings[i] = segment_embedding(path, segment, duration)
+        return np.nan_to_num(embeddings)
+
+    def segment_embedding(path, segment, duration):
+        start = segment["start"]
+        # Whisper overshoots the end timestamp in the last segment
+        end = min(duration, segment["end"])
+        clip = Segment(start, end)
+        waveform, sample_rate = Audio().crop(path, clip)
+        return embedding_model(waveform[None])
+
+    def add_speaker_labels(segments, embeddings, num_speakers):
+        clustering = AgglomerativeClustering(num_speakers).fit(embeddings)
+        labels = clustering.labels_
+        for i in range(len(segments)):
+            segments[i]["speaker"] = 'SPEAKER ' + str(labels[i] + 1)
+
+    def time(secs):
+        return datetime.timedelta(seconds=round(secs))
+
+    duration = get_duration(audio)
+    if duration > 4 * 60 * 60:
+        return "Audio duration too long"
+
+    result = model.transcribe(audio)
+
+    segments = result["segments"]
+
+    num_speakers = min(max(round(num_speakers), 1), len(segments))
+    if len(segments) == 1:
+        segments[0]['speaker'] = 'SPEAKER 1'
+    else:
+        embeddings = make_embeddings(audio, segments, duration)
+        add_speaker_labels(segments, embeddings, num_speakers)
+    return get_output(segments)
+    # return output
+
 def AudioTranscribe(NumberOfSpeakers=None, SpeakerNames="", audio="", retries=5):
     if retries:
         # subprocess.call(['ffmpeg', '-i', audio,'temp_audio.wav'])
@@ -141,7 +237,7 @@ def AudioTranscribe(NumberOfSpeakers=None, SpeakerNames="", audio="", retries=5)
             return AudioTranscribe(NumberOfSpeakers, SpeakerNames, audio, retries-1)
         if not (os.path.isfile("temp_audio.wav")):
             return AudioTranscribe(NumberOfSpeakers, SpeakerNames, audio, retries-1)
-        return Transcribe(NumberOfSpeakers, SpeakerNames)
+        return Transcribe_V2(NumberOfSpeakers, SpeakerNames)
     else:
         raise gr.Error("There is some issue ith Audio Transcriber. Please try again later!")
 
@@ -157,10 +253,10 @@ def VideoTranscribe(NumberOfSpeakers=None, SpeakerNames="", video="", retries=5)
             return VideoTranscribe(NumberOfSpeakers, SpeakerNames, video, retries-1)
         if not (os.path.isfile("temp_audio.wav")):
             return VideoTranscribe(NumberOfSpeakers, SpeakerNames, video, retries-1)
-        return Transcribe(NumberOfSpeakers, SpeakerNames)
+        return Transcribe_V2(NumberOfSpeakers, SpeakerNames)
     else:
         raise gr.Error("There is some issue ith Video Transcriber. Please try again later!")
-    return Transcribe(NumberOfSpeakers, SpeakerNames)
+    return Transcribe_V2(NumberOfSpeakers, SpeakerNames)
 
 def YoutubeTranscribe(NumberOfSpeakers=None, SpeakerNames="", URL="", retries = 5):
     if retries:
@@ -184,7 +280,7 @@ def YoutubeTranscribe(NumberOfSpeakers=None, SpeakerNames="", URL="", retries =
             stream = ffmpeg.input('temp_audio.m4a')
             stream = ffmpeg.output(stream, 'temp_audio.wav')
             RemoveFile("temp_audio.m4a")
-            return Transcribe(NumberOfSpeakers, SpeakerNames)
+            return Transcribe_V2(NumberOfSpeakers, SpeakerNames)
     else:
         raise gr.Error(f"Unable to get video from {URL}")