Update app.py

app.py

@@ -17,14 +17,15 @@ import moviepy.editor as mp
 import datetime
 import torch
 import pyannote.audio
-from pyannote.audio.pipelines.speaker_verification import PretrainedSpeakerEmbedding
+from pyannote.audio.pipelines.speaker_verification import SpeechBrainPretrainedSpeakerEmbedding #PyannoteAudioPretrainedSpeakerEmbedding
 from pyannote.audio import Audio
 from pyannote.core import Segment
 import wave
 import contextlib
 from sklearn.cluster import AgglomerativeClustering
 import numpy as np
-
+import json
+from datetime import timedelta
 
 __FILES = set()
       
@@ -140,7 +141,8 @@ def Transcribe_V1(NumberOfSpeakers, SpeakerNames="", audio="temp_audio.wav"):
 
 def Transcribe_V2(num_speakers, speaker_names, audio="temp_audio.wav"):
     model = whisper.load_model("medium")
-    embedding_model = PretrainedSpeakerEmbedding(
+    # embedding_model = SpeechBrainPretrainedSpeakerEmbedding("speechbrain/spkrec-ecapa-voxceleb")
+    embedding_model = SpeechBrainPretrainedSpeakerEmbedding(
         "speechbrain/spkrec-ecapa-voxceleb",
         device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     )
@@ -158,23 +160,50 @@ def Transcribe_V2(num_speakers, speaker_names, audio="temp_audio.wav"):
         return speaker
     
     # audio = Audio()
+    def diarization(audio):
+        def millisec(timeStr):
+            spl = timeStr.split(":")
+            s = (int)((int(spl[0]) * 60 * 60 + int(spl[1]) * 60 + float(spl[2]) )* 1000)
+            return s
+        as_audio = AudioSegment.from_wav(audio)
+        DEMO_FILE = {'uri': 'blabal', 'audio': audio}
+        hparams = pipeline.parameters(instantiated=True)
+        hparams["segmentation"]["min_duration_off"] -= 0.25
+        pipeline.instantiate(hparams)
+        if num_speakers:
+            dz = pipeline(DEMO_FILE, num_speakers=num_speakers)  
+        else:
+            dz = pipeline(DEMO_FILE)  
+        with open(CreateFile(f"diarization_{audio}.txt"), "w") as text_file:
+            text_file.write(str(dz))
+        dz = open(CreateFile(f"diarization_{audio}.txt")).read().splitlines()
+        print(dz)
+        dzList = []
+        for l in dz:
+            start, end =  tuple(re.findall('[0-9]+:[0-9]+:[0-9]+\.[0-9]+', string=l))
+            start = millisec(start)
+            end = millisec(end)
+            lex = GetSpeaker(re.findall('(SPEAKER_[0-9][0-9])', string=l)[0])
+            dzList.append([start, end, lex])
+        return dzList
+    
     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()
+                conversation.append([str(timedelta(seconds=float(segment['start']))),str(timedelta(seconds=float(segment['end']))),GetSpeaker(segment["speaker"]), segment["text"].lstrip()])
+            elif conversation[-1][2] == GetSpeaker(segment["speaker"]):
+                conversation[-1][3] +=  segment["text"].lstrip()
             else:
-                conversation.append([GetSpeaker(segment["speaker"]), segment["text"].lstrip()])
+                conversation.append([str(timedelta(seconds=float(segment['start']))),str(timedelta(seconds=float(segment['end']))),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])), ({ "data": [{"speaker": speaker, "text": text} for speaker, text in conversation]})
+        return ("".join([f"[{start}] - {speaker} \n{text}\n" for start, end, speaker, text in conversation])), ({ "data": [{"start": start, "end":end, "speaker": speaker, "text": text} for start, end, speaker, text in conversation]})
 
     def get_duration(path):
         with contextlib.closing(wave.open(path,'r')) as f:
@@ -209,7 +238,9 @@ def Transcribe_V2(num_speakers, speaker_names, audio="temp_audio.wav"):
     if duration > 4 * 60 * 60:
         return "Audio duration too long"
 
+    print(json.dumps(diarization(audio)))
     result = model.transcribe(audio)
+    print(json.dumps(result))
 
     segments = result["segments"]
 
@@ -233,7 +264,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_V1(NumberOfSpeakers, SpeakerNames)
+        return Transcribe_V2(NumberOfSpeakers, SpeakerNames)
     else:
         raise gr.Error("There is some issue ith Audio Transcriber. Please try again later!")
 
@@ -249,10 +280,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_V1(NumberOfSpeakers, SpeakerNames)
+        return Transcribe_V2(NumberOfSpeakers, SpeakerNames)
     else:
         raise gr.Error("There is some issue ith Video Transcriber. Please try again later!")
-    return Transcribe_V1(NumberOfSpeakers, SpeakerNames)
+    return Transcribe_V2(NumberOfSpeakers, SpeakerNames)
 
 def YoutubeTranscribe(NumberOfSpeakers=None, SpeakerNames="", URL="", retries = 5):
     if retries:
@@ -276,7 +307,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_V1(NumberOfSpeakers, SpeakerNames)
+            return Transcribe_V2(NumberOfSpeakers, SpeakerNames)
     else:
         raise gr.Error(f"Unable to get video from {URL}")
  
@@ -299,27 +330,28 @@ at = gr.Interface(
 # demo = gr.TabbedInterface([ut, vt, at], ["Youtube URL", "Video", "Audio"])
 # demo.launch()
 with gr.Blocks() as yav_ui:
-    with gr.Tab("Input"):
-        with gr.Tab("Youtube", id=1):
-            yinput_nos = gr.Number(label="Number of Speakers", placeholder="2")
-            yinput_sn = gr.Textbox(label="Name of the Speakers (ordered by the time they speak and separated by comma)", placeholder="If Speaker 1 is first to speak followed by Speaker 2 then -> Speaker 1, Speaker 2")
-            yinput = gr.Textbox(label="Youtube Link", placeholder="https://www.youtube.com/watch?v=GECcjrYHH8w")
-            ybutton_transcribe = gr.Button("Transcribe")
-        with gr.Tab("Video", id=2):
-            vinput_nos = gr.Number(label="Number of Speakers", placeholder="2")
-            vinput_sn = gr.Textbox(label="Name of the Speakers (ordered by the time they speak and separated by comma)", placeholder="If Speaker 1 is first to speak followed by Speaker 2 then -> Speaker 1, Speaker 2")
-            vinput = gr.Video(label="Video")
-            vbutton_transcribe = gr.Button("Transcribe")
-        with gr.Tab("Audio", id=3):
-            ainput_nos = gr.Number(label="Number of Speakers", placeholder="2")
-            ainput_sn = gr.Textbox(label="Name of the Speakers (ordered by the time they speak and separated by comma)", placeholder="If Speaker 1 is first to speak followed by Speaker 2 then -> Speaker 1, Speaker 2")
-            ainput = gr.Audio(label="Audio")
-            abutton_transcribe = gr.Button("Transcribe")
-    with gr.Tab("Output"):
-        with gr.Tab("Text"):
-            output_textbox = gr.Textbox(label="Transcribed Text", lines=15)
-        with gr.Tab("JSON"):
-            output_json = gr.JSON(label="Transcribed JSON")
+    with gr.Row():
+        with gr.Column():
+            with gr.Tab("Youtube", id=1):
+                yinput_nos = gr.Number(label="Number of Speakers", placeholder="2")
+                yinput_sn = gr.Textbox(label="Name of the Speakers (ordered by the time they speak and separated by comma)", placeholder="If Speaker 1 is first to speak followed by Speaker 2 then -> Speaker 1, Speaker 2")
+                yinput = gr.Textbox(label="Youtube Link", placeholder="https://www.youtube.com/watch?v=GECcjrYHH8w")
+                ybutton_transcribe = gr.Button("Transcribe", show_progress=True, scroll_to_output=True)
+            with gr.Tab("Video", id=2):
+                vinput_nos = gr.Number(label="Number of Speakers", placeholder="2")
+                vinput_sn = gr.Textbox(label="Name of the Speakers (ordered by the time they speak and separated by comma)", placeholder="If Speaker 1 is first to speak followed by Speaker 2 then -> Speaker 1, Speaker 2")
+                vinput = gr.Video(label="Video")
+                vbutton_transcribe = gr.Button("Transcribe", show_progress=True, scroll_to_output=True)
+            with gr.Tab("Audio", id=3):
+                ainput_nos = gr.Number(label="Number of Speakers", placeholder="2")
+                ainput_sn = gr.Textbox(label="Name of the Speakers (ordered by the time they speak and separated by comma)", placeholder="If Speaker 1 is first to speak followed by Speaker 2 then -> Speaker 1, Speaker 2")
+                ainput = gr.Audio(label="Audio", type="filepath")
+                abutton_transcribe = gr.Button("Transcribe", show_progress=True, scroll_to_output=True)
+        with gr.Column():
+            with gr.Tab("Text"):
+                output_textbox = gr.Textbox(label="Transcribed Text", lines=15)
+            with gr.Tab("JSON"):
+                output_json = gr.JSON(label="Transcribed JSON")
     ybutton_transcribe.click(
                 fn=YoutubeTranscribe,
                 inputs=[yinput_nos,yinput_sn,yinput],