Update app.py

app.py

@@ -1,17 +1,56 @@
-# load the pipeline from Hugginface Hub
 import gradio as gr
 from pyannote.audio import Pipeline
-pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization")
+from transformers import pipeline
 
-def speakerDiarization(audio):
-  # apply the pipeline to an audio file
-  diarization = pipeline(audio)
-  # dump the diarization output to disk using RTTM format
-  return diarization
+asr = pipeline(
+    "automatic-speech-recognition",
+    model="facebook/wav2vec2-large-960h-lv60-self",
+    feature_extractor="facebook/wav2vec2-large-960h-lv60-self",
+    
+)
+speaker_segmentation = Pipeline.from_pretrained("pyannote/speaker-segmentation")
 
-gr.Interface(speakerDiarization,
-             inputs = gr.inputs.Audio(source="microphone", type="filepath", optional=True, label="Please record your voice"),
-             outputs = gr.outputs.Image(label="Output Text"),
-             title="ASR using Wav2Vec 2.0",
-             description = "This application displays transcribed text for given audio input",
-             examples = [["TestAudio1.wav"]], theme="grass").launch()
+def segmentation(audio):
+    speaker_output = speaker_segmentation(audio)
+    text_output = asr(audio,return_timestamps="word")
+    
+    full_text = text_output['text'].lower()
+    chunks = text_output['chunks']
+
+    diarized_output = ""
+    i = 0
+    for turn, _, speaker in speaker_output.itertracks(yield_label=True):
+        diarized = ""
+        while i < len(chunks) and chunks[i]['timestamp'][1] <= turn.end:
+            diarized += chunks[i]['text'].lower() + ' '
+            i += 1
+        
+        if diarized != "":
+            diarized_output += "{}: ''{}'' from {:.3f}-{:.3f}\n".format(speaker,diarized,turn.start,turn.end)
+        
+    return diarized_output, full_text
+
+title = "Speech Recognition with Speaker Segmentation"
+description = "Speaker Diarization is the act of attributing individual speakers to their corresponding parts in an audio recording. This space aims to distinguish the speakers with speaker segmentation and their speech with speech-to-text from a given input audio file. Pre-trained models used are Pyannote[1] for the Speaker Segmentation and Wav2Vec2[2] for the Automatic Speech Recognition."
+article = "<p style='text-align: center'><a href='https://github.com/pyannote/pyannote-audio' target='_blank'>[1] Pyannote - Speaker Segmentation model (GitHub repo)</a></p>"
+article += "<p style='text-align: center'><a href='https://github.com/pytorch/fairseq/tree/main/examples/wav2vec#wav2vec-20' target='_blank'>[2] Facebook Wav2Vec2 (GitHub repo)</a></p>"
+article += "<p style='text-align: center'>Audio File Sources: <a href='https://www.youtube.com/watch?v=DYu_bGbZiiQ&t=132s' target='_blank'>1</a> <a href='https://www.youtube.com/watch?v=DDjWTWHHkpk&t=29s' target='_blank'>2</a> <a href='https://www.youtube.com/watch?v=G2xWg2ckKHI&t=24s' target='_blank'>3</a> <a href='https://www.youtube.com/watch?v=sCcv9uqSBU0&t=32s' target='_blank'>4</a> <a href='https://www.youtube.com/watch?v=K1hlp0DCE_8&t=71s' target='_blank'>5</a></p>"
+
+inputs = gr.inputs.Audio(source="upload", type="filepath", label="Upload your audio file here:")
+outputs = [gr.outputs.Textbox(type="auto", label="Diarized Output"),
+            gr.outputs.Textbox(type="auto",label="Full ASR Text for comparison")]
+examples = [["meeting_audio.wav"],
+            ["noisy_london_interview.wav"],
+            ["clean_london_interview.wav"],
+            ["podcast_audio.wav"],
+            ["air_traffic_control_audio.wav"],]
+
+app = gr.Interface(fn=segmentation,
+                inputs=inputs,
+                outputs=outputs,
+                examples=examples,
+                title=title,
+                description=description,
+                article=article,
+                allow_flagging=False)
+app.launch()