Implement basic diarization

app.py

@@ -1,8 +1,46 @@
+import os
+from dotenv import load_dotenv
 import gradio as gr
+import numpy as np
+import torch
+from pyannote.audio import Pipeline
+from pydub import AudioSegment
+from mimetypes import MimeTypes
 
+load_dotenv()
 
-def greet(name):
-    return "Hello " + name + "!!"
+hg_token = os.getenv("HG_ACCESS_TOKEN") 
 
-iface = gr.Interface(fn=greet, inputs="text", outputs="text")
-iface.launch()
+if hg_token != None:
+    pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization", use_auth_token=hg_token)
+else:
+    print('''No hugging face access token set. 
+You need to set it via an .env or environment variable HG_ACCESS_TOKEN''')
+    exit(1)
+
+
+def diarization(audio_file: tuple[int, np.array]) -> np.array:
+    """
+    Receives a tuple with the sample rate and audio data and returns the
+    a numpy array containing the audio segments, track names and speakers for 
+    each segment.
+    """
+    waveform = torch.tensor(audio_file[1].astype(np.float32, order='C')).reshape(1,-1)
+    audio_data = {
+        "waveform": waveform,
+        "sample_rate": audio_file[0]
+    }
+
+    diarization = pipeline(audio_data)
+    
+    return np.array(list(diarization.itertracks(yield_label=True)))
+
+
+
+demo = gr.Interface(
+    fn=diarization,
+    inputs=gr.Audio(type="numpy"),
+    outputs="text",
+)
+
+demo.launch()

requirements.txt

@@ -1,5 +1,6 @@
 gradio==3.21.0
 openai-whisper==20230314
+pydub==0.25.1
 python-dotenv==1.0.0
 pyannote.audio==2.1.1
 torch==1.13.1