Implement diarization

app.py

@@ -1,7 +1,25 @@
+from huggingface_hub import login
+from diarization import startDiarization
+import ffmpeg
 import gradio as gr
+import os
 
-def greet(name):
-    return "Hello " + name + "!!"
+hugging_face_token = os.environ["HUGGING_FACE_TOKEN"]
+login(token=hugging_face_token)
 
-iface = gr.Interface(fn=greet, inputs="text", outputs="text")
-iface.launch()
+# Get video & convert it
+inputs_interface = gr.inputs.Video(label="Insert video")
+in_file = ffmpeg.input(inputs_interface).audio
+out_file = ffmpeg.output(in_file, "input.wav", f="wav").run()
+
+
+def prepareInput():
+    return startDiarization(out_file)
+
+
+gr.Interface(
+    prepareInput,
+    inputs=inputs_interface,
+    outputs="text",
+    title="Get Diarization"
+).launch()

diarization.py

@@ -0,0 +1,44 @@
+from pyannote.audio import Pipeline
+from pydub import AudioSegment
+import torch
+import json
+
+pipeline = Pipeline.from_pretrained('pyannote/speaker-diarization')
+device = torch.device("cpu")
+pipeline.to(device)
+
+
+def startDiarization(input_file):
+    diarization = pipeline(input_file)
+
+    sample_groups = []
+    speaker_groups = {}
+    for turn, _, speaker in diarization.itertracks(yield_label=True):
+        if (speaker not in sample_groups):
+            sample_groups.append(str(speaker))
+
+        suffix = 1
+        file_name = f"{speaker}-{suffix}"
+        while file_name in speaker_groups:
+            suffix += 1
+            file_name = f"{speaker}-{suffix}"
+        speaker_groups[file_name] = [turn.start, turn.end]
+    saveGroupsJson(sample_groups, speaker_groups)
+    audioSegmentation(input_file, speaker_groups)
+    return str(speaker_groups)
+
+
+def audioSegmentation(input_file, speaker_groups_dict):
+    audioSegment = AudioSegment.from_wav(input_file)
+    for speaker in speaker_groups_dict:
+        time = speaker_groups_dict[speaker]
+        audioSegment[time[0]*1000: time[1] *
+                     1000].export(f"{speaker}.wav", format='wav')
+        print(f"group {speaker}: {time[0]*1000}--{time[1]*1000}")
+
+
+def saveGroupsJson(sample_groups_list: list, speaker_groups_dict: dict):
+    with open("sample_groups.json", "w") as json_file_sample:
+        json.dump(sample_groups_list, json_file_sample)
+    with open("speaker_groups.json", "w") as json_file_speaker:
+        json.dump(speaker_groups_dict, json_file_speaker)

requirements.txt

@@ -0,0 +1,6 @@
+ffmpeg-python
+pyannote @ git+https://github.com/pyannote/pyannote-audio.git@develop
+pydub
+transformers
+torch
+whisper