update app.py

app.py

@@ -1,22 +1,5 @@
-from google.colab import files
-uploaded = files.upload()
-path = next(iter(uploaded))
-
-num_speakers = 2 #@param {type:"integer"}
-
-language = 'English' #@param ['any', 'English']
-
-model_size = 'large' #@param ['tiny', 'base', 'small', 'medium', 'large']
-
-
-model_name = model_size
-if language == 'English' and model_size != 'tiny':
-  model_name += '.en'
-
-!pip install -q git+https://github.com/openai/whisper.git > /dev/null
-!pip install -q git+https://github.com/pyannote/pyannote-audio > /dev/null
-
 import whisper
+import gradio as gr
 import datetime
 
 import subprocess
@@ -24,9 +7,6 @@ import subprocess
 import torch
 import pyannote.audio
 from pyannote.audio.pipelines.speaker_verification import PretrainedSpeakerEmbedding
-embedding_model = PretrainedSpeakerEmbedding(
-    "speechbrain/spkrec-ecapa-voxceleb",
-    device=torch.device("cuda"))
 
 from pyannote.audio import Audio
 from pyannote.core import Segment
@@ -37,98 +17,93 @@ import contextlib
 from sklearn.cluster import AgglomerativeClustering
 import numpy as np
 
-
-if path[-3:] != 'wav':
-  subprocess.call(['ffmpeg', '-i', path, 'audio.wav', '-y'])
-  path = 'audio.wav'
-
-model = whisper.load_model(model_size)
-
-result = model.transcribe(path)
-segments = result["segments"]
-
-with contextlib.closing(wave.open(path,'r')) as f:
-  frames = f.getnframes()
-  rate = f.getframerate()
-  duration = frames / float(rate)
+model = whisper.load_model("large-v2")
+embedding_model = PretrainedSpeakerEmbedding( 
+    "speechbrain/spkrec-ecapa-voxceleb",
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+)
+
+def transcribe(audio, num_speakers):
+  path, error = convert_to_wav(audio)
+  if error is not None:
+    return error
+
+  duration = get_duration(path)
+  if duration > 4 * 60 * 60:
+    return "Audio duration too long"
+
+  result = model.transcribe(path)
+  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(path, segments, duration)
+    add_speaker_labels(segments, embeddings, num_speakers)
+  output = get_output(segments)
+  return output
+
+def convert_to_wav(path):
+  if path[-3:] != 'wav':
+    new_path = '.'.join(path.split('.')[:-1]) + '.wav'
+    try:
+      subprocess.call(['ffmpeg', '-i', path, new_path, '-y'])
+    except:
+      return path, 'Error: Could not convert file to .wav'
+    path = new_path
+  return path, None
+
+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)
 
 audio = Audio()
 
-def segment_embedding(segment):
+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])
 
-embeddings = np.zeros(shape=(len(segments), 192))
-for i, segment in enumerate(segments):
-  embeddings[i] = segment_embedding(segment)
-
-embeddings = np.nan_to_num(embeddings)
-
-clustering = AgglomerativeClustering(num_speakers).fit(embeddings)
-labels = clustering.labels_
-for i in range(len(segments)):
-  segments[i]["speaker"] = 'SPEAKER ' + str(labels[i] + 1)
-  # speaker = 'Held'
-  # speaker = 'Heldisha'
-  # if segments[i]["speaker"]== 'SPEAKER 1':
-  #   segments[i]["speaker"] = 'Held'
-  # elif segments[i]["speaker"]== 'SPEAKER 2':
-  #   segments[i]["speaker"] = 'Heldisha'
-  # if segments[i]["speaker"]== 'SPEAKER 1':
-    # segments[i]["speaker"] = segments.index('n')
-# k = list(segments)
-# print(k[5])
+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))
 
-f = open("transcript.txt", "w")
-
-for (i, segment) in enumerate(segments):
-  if i == 0 or segments[i - 1]["speaker"] != segment["speaker"]:
-    f.write("\n" + segment["speaker"] + ' ' + str(time(segment["start"])) + '\n')
-  f.write(segment["text"][1:] + ' ')
-f.close()
-
-# with open('transcript.txt', 'r') as file:
-#     text = file.read()
-#     words = text.split()
-#     i = words.index('name')
-#     if (words[i-1] == 'My') or (words[i-1] == 'my') and (words[i+1] == 'is'):
-#         name1 = words[i+2]
-#         print(name1)
-  
-#     with open('transcript.txt', 'r') as file:
-#         text = file.read()
-#         new_text = text.replace('SPEAKER 1', name1)
-#     with open('transcript.txt', 'w') as file:
-#         file.write(new_text)
-
-
-
-# with open('transcript.txt', 'r') as file:
-    
-#     text = file.read()
-#     words = text.split()
-#     i = words.index('name')
-#     if (words[i+3] == 'What') or (1<2) and (words[i+1] == 'is') or 1<2:
-#         name2 = words[i+22]
-#         print(name2)
-#     with open('transcript.txt', 'r') as file:
-#         text = file.read()
-#         new_text = text.replace('SPEAKER 2', name2)
-#     with open('transcript.txt', 'w') as file:
-#         file.write(new_text)
-    
-    
-    
-
-
-
-    
-    
-    
+def get_output(segments):
+  output = ''
+  for (i, segment) in enumerate(segments):
+    if i == 0 or segments[i - 1]["speaker"] != segment["speaker"]:
+      if i != 0:
+        output += '\n\n'
+      output += segment["speaker"] + ' ' + str(time(segment["start"])) + '\n\n'
+    output += segment["text"][1:] + ' '
+  return output
+
+gr.Interface(
+    title = 'Whisper with Speaker Recognition', 
+    fn=transcribe, 
+    inputs=[
+        gr.inputs.Audio(source="upload", type="filepath"),
+        gr.inputs.Number(default=2, label="Number of Speakers")
+
+    ],
+    outputs=[
+        gr.outputs.Textbox(label='Transcript')
+    ]
+  ).launch()