Update app.py

app.py

@@ -24,6 +24,7 @@ import wave
 import contextlib
 from sklearn.cluster import AgglomerativeClustering
 import numpy as np
+import json
 
 __FILES = set()
       
@@ -157,6 +158,32 @@ 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_onset"] += 0.1
+        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()
+        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=[]
@@ -209,6 +236,8 @@ def Transcribe_V2(num_speakers, speaker_names, audio="temp_audio.wav"):
         return "Audio duration too long"
 
     result = model.transcribe(audio)
+    json.dumps(result)
+    json.dumps(diarization(audio))
 
     segments = result["segments"]