Update app.py

app.py

@@ -192,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, min_num_speakers, max_number_speakers):
+def speech_to_text(video_file_path, selected_source_lang, whisper_model, num_speakers):
     """
     # Transcribe youtube link using OpenAI Whisper
     1. Using Open AI's Whisper model to seperate audio into segments and generate transcripts.
@@ -250,22 +250,19 @@ def speech_to_text(video_file_path, selected_source_lang, whisper_model, min_num
         embeddings = np.nan_to_num(embeddings)
         print(f'Embedding shape: {embeddings.shape}')
 
+        if num_speakers == 0:
         # 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(2, 10+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")
         else:
-            min_speakers = min_num_speakers
-            max_speakers = max_number_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")
-
+            best_num_speaker = num_speakers
+            
         # Assign speaker label   
         clustering = AgglomerativeClustering(best_num_speaker).fit(embeddings)
         labels = clustering.labels_
@@ -320,8 +317,7 @@ 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)
-input_min_number_speakers = gr.Number(precision=0, value=2, label="Select minimum number of speakers", interactive=True)
-input_max_number_speakers = gr.Number(precision=0, value=2, label="Select maximum number of speakers", interactive=True)
+number_speakers = gr.Number(precision=0, value=0, label="Input number of speakers for better results. If value=0, model will automatic find the best 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')
@@ -378,11 +374,10 @@ with demo:
                     ''')
                 selected_source_lang.render()
                 selected_whisper_model.render()
-                input_min_number_speakers.render()
-                input_max_number_speakers.render()
+                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, input_min_number_speakers, input_max_number_speakers], 
+                                     [video_in, selected_source_lang, selected_whisper_model, number_speakers], 
                                      [transcription_df, system_info, download_transcript]
                                     )