Update app.py

app.py

@@ -9,6 +9,7 @@ import time
 import os 
 import numpy as np
 from sklearn.cluster import AgglomerativeClustering
+from sklearn.metrics import silhouette_score
 
 from pytube import YouTube
 import torch
@@ -191,7 +192,7 @@ def get_youtube(video_url):
     print(abs_video_path)
     return abs_video_path
 
-def speech_to_text(video_file_path, selected_source_lang, whisper_model, num_speakers):
+def speech_to_text(video_file_path, selected_source_lang, whisper_model, min_num_speakers, max_number_speakers):
     """
     # Transcribe youtube link using OpenAI Whisper
     1. Using Open AI's Whisper model to seperate audio into segments and generate transcripts.
@@ -249,8 +250,21 @@ def speech_to_text(video_file_path, selected_source_lang, whisper_model, num_spe
         embeddings = np.nan_to_num(embeddings)
         print(f'Embedding shape: {embeddings.shape}')
 
-        # Assign speaker label
-        clustering = AgglomerativeClustering(num_speakers).fit(embeddings)
+        # Find the best number of speakers
+        if min_num_speakers > max_number_speakers:
+            min_speakers = max_number_speakers
+            max_speakers = min_num_speakers
+        score_num_speakers = {}
+
+        for num_speakers in range(min_speakers, max_speakers+1):
+            clustering = AgglomerativeClustering(num_speakers).fit(embeddings)
+            score = silhouette_score(embeddings, clustering.labels_, metric='euclidean')
+            score_num_speakers[num_speakers] = score
+        best_num_speaker = max(score_num_speakers, key=lambda x:score_num_speakers[x])
+        print(f"The best number of speakers: {best_num_speaker} with {score_num_speakers[best_num_speaker]} score")
+
+        # Assign speaker label   
+        clustering = AgglomerativeClustering(best_num_speaker).fit(embeddings)
         labels = clustering.labels_
         for i in range(len(segments)):
             segments[i]["speaker"] = 'SPEAKER ' + str(labels[i] + 1)
@@ -289,7 +303,7 @@ def speech_to_text(video_file_path, selected_source_lang, whisper_model, num_spe
         save_path = "output/transcript_result.csv"
         df_results = pd.DataFrame(objects)
         df_results.to_csv(save_path)
-        return df_results, system_info, save_path
+        return df_results, system_info, save_pathassuming
     
     except Exception as e:
         raise RuntimeError("Error Running inference with local model", e)
@@ -303,7 +317,8 @@ df_init = pd.DataFrame(columns=['Start', 'End', 'Speaker', 'Text'])
 memory = psutil.virtual_memory()
 selected_source_lang = gr.Dropdown(choices=source_language_list, type="value", value="en", label="Spoken language in video", interactive=True)
 selected_whisper_model = gr.Dropdown(choices=whisper_models, type="value", value="base", label="Selected Whisper model", interactive=True)
-number_speakers = gr.Number(precision=0, value=2, label="Selected number of speakers", interactive=True)
+input_min_number_speakers = gr.Number(precision=0, value=2, label="Select assumed minimum number of speakers", interactive=True)
+input_max_number_speakers = gr.Number(precision=0, value=2, label="Select assumed maximum number of speakers", interactive=True)
 system_info = gr.Markdown(f"*Memory: {memory.total / (1024 * 1024 * 1024):.2f}GB, used: {memory.percent}%, available: {memory.available / (1024 * 1024 * 1024):.2f}GB*")
 download_transcript = gr.File(label="Download transcript")
 transcription_df = gr.DataFrame(value=df_init,label="Transcription dataframe", row_count=(0, "dynamic"), max_rows = 10, wrap=True, overflow_row_behaviour='paginate')
@@ -356,14 +371,17 @@ with demo:
                     gr.Markdown('''
                     ##### Here you can start the transcription process.
                     ##### Please select the source language for transcription.
-                    ##### You should select a number of speakers for getting better results.
+                    ##### You can select a range of assumed numbers of speakers.
                     ''')
                 selected_source_lang.render()
                 selected_whisper_model.render()
-                number_speakers.render()
+                input_min_number_speakers.render()
+                input_min_number_speakers.render()
                 transcribe_btn = gr.Button("Transcribe audio and diarization")
-                transcribe_btn.click(speech_to_text, [video_in, selected_source_lang, selected_whisper_model, number_speakers], [transcription_df, system_info, download_transcript])
-
+                transcribe_btn.click(speech_to_text, 
+                                     [video_in, selected_source_lang, selected_whisper_model, input_min_number_speakers, input_min_number_speakers], 
+                                     [transcription_df, system_info, download_transcript]
+                                    )
                 
         with gr.Row():
             gr.Markdown('''