Update app.py

app.py

@@ -15,6 +15,17 @@ from pydub.effects import speedup
 
 
 __FILES = set()
+SPEAKER_DICT = {}
+SPEAKERS = []
+
+def GetSpeaker(sp):
+    if sp not in list(SPEAKER_DICT.keys):
+        SPEAKER_DICT[sp] = SPEAKERS.pop(0)
+    return SPEAKER_DICT[sp]
+
+def GenerateSpeakerDict(sp):
+    SPEAKERS = [speaker.strip() for speaker in sp.split(',')]
+        
 
 def CreateFile(filename):
     __FILES.add(filename)
@@ -29,7 +40,7 @@ def RemoveAllFiles():
         if (os.path.isfile(file)):
             os.remove(file)
     
-def Transcribe(audio="temp_audio.wav"):
+def Transcribe(NumberOfSpeakers, SpeakerNames="", audio="temp_audio.wav"):
     def millisec(timeStr):
         spl = timeStr.split(":")
         s = (int)((int(spl[0]) * 60 * 60 + int(spl[1]) * 60 + float(spl[2]) )* 1000)
@@ -58,7 +69,7 @@ def Transcribe(audio="temp_audio.wav"):
             start, end =  tuple(re.findall('[0-9]+:[0-9]+:[0-9]+\.[0-9]+', string=l))
             start = millisec(start)
             end = millisec(end)
-            lex = re.findall('(SPEAKER_[0-9][0-9])', string=l)[0]
+            lex = GetSpeaker(re.findall('(SPEAKER_[0-9][0-9])', string=l)[0])
             dzList.append([start, end, lex])
         sounds = spacer
         segments = []
@@ -101,6 +112,7 @@ def Transcribe(audio="temp_audio.wav"):
                   #print(f"[{dzList[i][2]}] {c[2]}")
         return conversation, ("".join([f"{speaker} --> {text}\n" for speaker, text in conversation]))
 
+    GenerateSpeakerDict(SpeakerNames)
     spacermilli, spacer = preprocess(audio)
     dz_audio, dzList, segments = diarization(audio)
     conversation, t_text = transcribe(dz_audio)
@@ -116,7 +128,7 @@ def AudioTranscribe(NumberOfSpeakers=None, SpeakerNames="", audio="", retries=5)
             return AudioTranscribe(audio, retries-1)
         if not (os.path.isfile("temp_audio.wav")):
             return AudioTranscribe(audio, retries-1)
-        return Transcribe()
+        return Transcribe(NumberOfSpeakers, SpeakerNames)
     else:
         raise gr.Error("There is some issue ith Audio Transcriber. Please try again later!")
 
@@ -130,10 +142,10 @@ def VideoTranscribe(NumberOfSpeakers=None, SpeakerNames="", video="", retries=5)
             return VideoTranscribe(video, retries-1)
         if not (os.path.isfile("temp_audio.wav")):
             return VideoTranscribe(video, retries-1)
-        return Transcribe()
+        return Transcribe(NumberOfSpeakers, SpeakerNames)
     else:
         raise gr.Error("There is some issue ith Video Transcriber. Please try again later!")
-    return Transcribe()
+    return Transcribe(NumberOfSpeakers, SpeakerNames)
 
 def YoutubeTranscribe(NumberOfSpeakers=None, SpeakerNames="", URL="", retries = 5):
     if retries:
@@ -157,7 +169,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()
+            return Transcribe(NumberOfSpeakers, SpeakerNames)
     else:
         raise gr.Error(f"Unable to get video from {URL}")