Implement speaker diarization and transcription merging pipeline

- Add diarization functionality with configurable speaker count
- Create merge_transcription_and_diarization function to assign speakers to transcribed text
- Enhance app.py with debug mode and dynamic file loading
- Update UI to support advanced speaker configuration
- Modify logo display and add visibility controls for intermediate outputs

app.py

@@ -19,12 +19,14 @@ from ui_config import (
     ERROR_BACKGROUND_COLOR, ERROR_TEXT_COLOR, ERROR_BORDER_COLOR,
     BUTTON_SECONDARY_BACKGROUND_COLOR, BUTTON_SECONDARY_BORDER_COLOR,
     BUTTON_SECONDARY_TEXT_COLOR, RED, GREEN, BLUE,
-    html_social_media, get_html_subtify_logo, html_buy_me_a_coffe
+    html_social_media, get_html_subtify_logo_big, get_html_subtify_logo_small, html_buy_me_a_coffe
 )
 # from url_manager import get_youtube_thumbnail, is_valid_youtube_url, is_valid_twitch_url, is_valid_url
 from slice_audio import slice_audio as slice_audio_main
 from audio import get_audio_from_video
 from transcribe import transcribe, get_language_dict
+from diarize_library import diarize_audio
+import json
 
 NUMBER = 100
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
@@ -56,13 +58,17 @@ YOUTUBE = "youtube"
 TWITCH = "twitch"
 ERROR = "error"
 
+VIEW_OUTPUTS = True
+DEBUG = True
+
 subtify_logo = Image.open("assets/subtify_logo-scaled.png")
 subtify_logo_width, subtify_logo_height = subtify_logo.size
 factor = 4
 new_width = subtify_logo_width // factor
 new_height = subtify_logo_height // factor
 
-html_subtify_logo = get_html_subtify_logo(new_width, new_height)
+html_subtify_logo_big = get_html_subtify_logo_big(new_width, new_height)
+html_subtify_logo_small = get_html_subtify_logo_small(new_width, new_height)
 
 language_dict = union_language_dict()
 
@@ -118,8 +124,10 @@ def change_visibility_texboxes():
     return (
         gr.update(value="Done"), # auxiliar_block1
         gr.update(visible=True), # get_audio_from_video_info
+        gr.update(visible=True), # merged_transcription
         gr.update(visible=True), # video_sliced_progress_info
         gr.update(visible=True), # video_transcribed_progress_info
+        gr.update(visible=True), # diarization_progress_info
         gr.update(visible=True), # transcriptions_concatenated_progress_info
         gr.update(visible=True), # video_translated_progress_info
         gr.update(visible=True), # video_subtitled_progress_info
@@ -128,8 +136,17 @@ def change_visibility_texboxes():
 def get_audio(video_path):
     print('*'*NUMBER)
     print(f"Getting audio from video {video_path}")
-    
+
     audios_folder = "audios"
+
+    if DEBUG:
+        audio_file = f"{audios_folder}/download_audio.mp3"
+        if os.path.exists(audio_file):
+            return [
+                gr.update(value="Loaded"),  # get_audio_from_video_info
+                gr.update(value=audio_file)  # original_audio_path
+            ]
+    
     try:
         audio_path = get_audio_from_video(video_path, audios_folder)
         return [
@@ -162,18 +179,64 @@ def slice_audio(input_audio_path):
         gr.update(value="Ok"),  # video_sliced_progress_info
     )
 
+def diarize(input_audio_path, num_speakers, min_speakers, max_speakers):
+    print('*'*NUMBER)
+    print(f"Diarize {input_audio_path}")
+
+    # Diarization file
+    diarization_file = f"diarization/diarization.json"
+
+    if DEBUG:
+        if os.path.exists(diarization_file):
+            with open(diarization_file, "r") as f:
+                diarization = f.read()
+            return [
+                gr.update(value="Loaded"),
+                gr.update(value=diarization)
+            ]
+
+    # Diarize audio
+    diarization = diarize_audio(input_audio_path, num_speakers, min_speakers, max_speakers, DEVICE)
+
+    # Save diarization
+    with open(diarization_file, "w") as f:
+        json.dump(diarization, f)
+
+    return [
+        gr.update(value="Ok"),
+        gr.update(value=diarization)
+    ]
+
 def trascribe_audio(input_audio_path, source_languaje):
     print('*'*NUMBER)
     print(f"Transcript {input_audio_path}")
 
+    # Transcription file
+    transcription_file = f"transcriptions/transcription_{source_languaje}.json"
+
+    if DEBUG:
+        if os.path.exists(transcription_file):
+            transcription = open(transcription_file, "r").read()
+            transcription = json.loads(transcription)
+            return [
+                gr.update(value="Loaded"),
+                gr.update(value=transcription)
+            ]
+
     # Get language dict
     language_dict = get_language_dict()
 
     # Transcribe audio file
-    transcribe(input_audio_path, language_dict[source_languaje]["transcriber"], DEVICE, CHUNK_SECONDS, CHUNK_OVERLAP_SECONDS)
+    transcription_str, transcription_dict = transcribe(input_audio_path, language_dict[source_languaje]["transcriber"], DEVICE, CHUNK_SECONDS, CHUNK_OVERLAP_SECONDS)
+
+    # Save transcription
+    with open(transcription_file, "w") as f:
+        transcription_json = json.dumps(transcription_dict)
+        f.write(transcription_json)
         
     return (
-        gr.Textbox(value="Ok")
+        gr.update(value="Ok"),
+        gr.update(value=transcription_dict)
     )
 
 def concatenate_transcriptions():
@@ -276,6 +339,16 @@ def process_uploaded_video(video_path):
     videos_folder = "videos"
     if not os.path.exists(videos_folder):
         os.makedirs(videos_folder)
+
+    if DEBUG:
+        video_file = f"{videos_folder}/download_video.mp4"
+        if os.path.exists(video_file):
+            return [
+                gr.update(label="Video uploaded"),  # video_input
+                gr.update(visible=True),  # config_block
+                gr.update(value=video_file),  # original_video_path
+                gr.update(value=html_subtify_logo_small)  # html_subtify_logo_component
+            ]
         
     # Copy uploaded video to videos folder
     new_video_path = os.path.join(videos_folder, "download_video.mp4")
@@ -285,7 +358,103 @@ def process_uploaded_video(video_path):
     return [
         gr.update(label="Video uploaded"),  # video_input
         gr.update(visible=True),  # config_block
-        gr.update(value=new_video_path)  # original_video_path
+        gr.update(value=new_video_path),  # original_video_path
+        gr.update(value=html_subtify_logo_small)  # html_subtify_logo_component
+    ]
+
+def merge_transcription_and_diarization():
+    """
+    Combina la transcripción y la diarización para asignar speakers a cada palabra.
+    
+    Returns:
+        dict: Transcripción combinada con información de speakers
+    """
+    print('*'*NUMBER)
+    print("Merge transcription and diarization")
+    
+    if DEBUG:
+        merged_transcription_path = "merged_transcription_diarization/merged.json"
+        if os.path.exists(merged_transcription_path):
+            with open(merged_transcription_path, 'r') as f:
+                merged_transcription = json.load(f)
+            return [
+                gr.update(value="Loaded"),
+                gr.update(value=merged_transcription)
+            ]
+
+    transcription_path = "transcriptions/transcription_English.json"
+    diarization_path = "diarization/diarization.json"
+    # Cargar los archivos JSON
+    with open(transcription_path, 'r') as f:
+        transcription = json.load(f)
+    with open(diarization_path, 'r') as f:
+        diarization = json.load(f)
+    
+    # Crear una nueva lista para los chunks combinados
+    merged_chunks = []
+    
+    # Para cada palabra en la transcripción
+    for chunk in transcription.get('chunks', []):
+        # Verificar que el chunk tiene timestamps válidos
+        if not (isinstance(chunk.get('start'), (int, float)) and 
+                isinstance(chunk.get('end'), (int, float))):
+            continue
+            
+        word_start = float(chunk['start'])
+        word_end = float(chunk['end'])
+        
+        # Encontrar el speaker correspondiente en la diarización
+        speaker = None
+        for segment in diarization:
+            # Verificar que el segmento tiene timestamps válidos
+            if not (isinstance(segment.get('start'), (int, float)) and 
+                    isinstance(segment.get('end'), (int, float))):
+                continue
+                
+            segment_start = float(segment['start'])
+            segment_end = float(segment['end'])
+            
+            # Si la palabra está dentro del rango de tiempo del segmento
+            if (word_start >= segment_start and word_end <= segment_end):
+                speaker = segment['speaker']
+                break
+            
+            # Si la palabra está mayormente dentro del segmento (más del 50% de su duración)
+            word_duration = word_end - word_start
+            overlap_start = max(word_start, segment_start)
+            overlap_end = min(word_end, segment_end)
+            overlap_duration = max(0, overlap_end - overlap_start)
+            
+            if overlap_duration > word_duration * 0.5:
+                speaker = segment['speaker']
+                break
+        
+        # Crear el nuevo chunk con la información del speaker
+        merged_chunk = {
+            'start': word_start,
+            'end': word_end,
+            'text': chunk['text'],
+            'speaker': speaker if speaker else 'UNKNOWN'
+        }
+        merged_chunks.append(merged_chunk)
+    
+    # Crear el diccionario final
+    merged_transcription = {
+        'text': transcription.get('text', ''),
+        'chunks': merged_chunks
+    }
+    
+    # Crear el directorio si no existe
+    if not os.path.exists(merged_transcription_path):
+        os.makedirs(merged_transcription_path)
+    
+    # Guardar el resultado en el nuevo directorio
+    with open(merged_transcription_path, 'w', encoding='utf-8') as f:
+        json.dump(merged_transcription, f, ensure_ascii=False, indent=2)
+    
+    return [
+        gr.update(value="Ok"),
+        gr.update(value=merged_transcription)
     ]
 
 @spaces.GPU
@@ -336,7 +505,7 @@ def subtify():
         # Layout
         gr.HTML(html_social_media)
         gr.HTML("<h1 style='text-align: center;'>Subtify</h1>")
-        gr.HTML(html_subtify_logo)
+        html_subtify_logo_component = gr.HTML(html_subtify_logo_big)
         
         # Input block, where the user can upload a video and configure the subtify process
         visible = False
@@ -357,21 +526,27 @@ def subtify():
                 with gr.Row():
                     source_languaje = gr.Dropdown(visible=True, label="Source languaje", show_label=True, value="English", choices=language_dict, scale=1, interactive=True, info="Language of the video")
                     target_languaje = gr.Dropdown(visible=True, label="Target languaje", show_label=True, value="Español", choices=language_dict, scale=1, interactive=True, info="Language to translate the subtitles")
-                with gr.Accordion("Advanced settings", open=False, visible=True) as Advanced_setings:
-                    number_of_speakers = gr.Dropdown(visible=True, label="Number of speakers", show_label=True, value=10, choices=num_speaker, scale=1, interactive=True, info="Number of speakers in the video, if you don't know, select 10")
+                with gr.Accordion("Advanced settings", open=True, visible=True) as Advanced_setings:
+                    number_of_speakers = gr.Number(visible=True, label="Number of speakers", show_label=True, value=0, interactive=True, info="Number of speakers in the video, if you don't know, select 0")
+                    min_speakers = gr.Number(visible=True, label="Min speakers", show_label=True, value=0, scale=0, interactive=True, info="Minimum number of speakers in the video")
+                    max_speakers = gr.Number(visible=True, label="Max speakers", show_label=True, value=0, scale=0, interactive=True, info="Maximum number of speakers in the video")
                 subtify_button = gr.Button(size="lg", value="subtify", min_width="10px", scale=0, visible=True)
 
         auxiliar_block1 = gr.Textbox(placeholder="", interactive=False, visible=visible)
         with gr.Row():
             get_audio_from_video_info = gr.Textbox(placeholder="Waiting", label="Get audio from video info", elem_id="get_audio_from_video_info", interactive=False, visible=visible)
             video_transcribed_progress_info = gr.Textbox(placeholder="Waiting", label="Transcribe progress info", elem_id="video_transcribed_progress_info", interactive=False, visible=visible)
+            diarization_progress_info = gr.Textbox(placeholder="Waiting", label="Diarize progress info", elem_id="diarization_progress_info", interactive=False, visible=visible)
+            merged_transcription_progress_info = gr.Textbox(placeholder="Waiting", label="Merge transcription and diarization progress info", elem_id="merged_transcription_progress_info", interactive=False, visible=visible)
             transcriptions_concatenated_progress_info = gr.Textbox(placeholder="Waiting", label="Concatenate progress info", elem_id="transcriptions_concatenated_progress_info", interactive=False, visible=visible)
             video_translated_progress_info = gr.Textbox(placeholder="Waiting", label="Translate progress info", elem_id="transcription_translated_progress_info", interactive=False, visible=visible)
             video_subtitled_progress_info = gr.Textbox(placeholder="Waiting", label="Video subtitle progress info", elem_id="video_subtitled_progress_info", interactive=False, visible=visible)
 
         original_audio_path = gr.Textbox(label="Original audio path", elem_id="original_audio_path", visible=visible)
         original_video_path = gr.Textbox(label="Original video path", visible=visible)
-        original_audio_transcribed_path = gr.Textbox(label="Original audio transcribed", elem_id="original_audio_transcribed", visible=visible)
+        transcription = gr.Textbox(label="transcription", elem_id="transcription", visible=VIEW_OUTPUTS)
+        diarization = gr.Textbox(label="diarization", elem_id="diarization", visible=VIEW_OUTPUTS)
+        merged_transcription = gr.Textbox(label="merged_transcription", elem_id="merged_transcription", visible=VIEW_OUTPUTS)
         original_audio_translated_path = gr.Textbox(label="Original audio translated", elem_id="original_audio_translated", visible=visible)
         subtitled_video = gr.Video(label="Subtitled video", elem_id="subtitled_video", visible=visible, interactive=visible)
         auxiliar_block3 = gr.Textbox(placeholder="Waiting", label="Auxiliar block 3", elem_id="auxiliar_block3", interactive=False, visible=visible)
@@ -400,11 +575,11 @@ def subtify():
         video_input.change(
             fn=process_uploaded_video,
             inputs=[video_input],
-            outputs=[video_input, config_block, original_video_path]
+            outputs=[video_input, config_block, original_video_path, html_subtify_logo_component]
         )
         subtify_button.click(
             fn=change_visibility_texboxes,
-            outputs=[auxiliar_block1, get_audio_from_video_info, video_transcribed_progress_info, transcriptions_concatenated_progress_info, video_translated_progress_info, video_subtitled_progress_info]
+            outputs=[auxiliar_block1, get_audio_from_video_info, merged_transcription_progress_info, video_transcribed_progress_info, diarization_progress_info, transcriptions_concatenated_progress_info, video_translated_progress_info, video_subtitled_progress_info]
         )
         auxiliar_block1.change(
             fn=get_audio,
@@ -414,12 +589,17 @@ def subtify():
         get_audio_from_video_info.change(
             fn=trascribe_audio,
             inputs=[original_audio_path, source_languaje],
-            outputs=[video_transcribed_progress_info]
+            outputs=[video_transcribed_progress_info, transcription]
+        )
+        video_transcribed_progress_info.change(
+            fn=diarize,
+            inputs=[original_audio_path, number_of_speakers, min_speakers, max_speakers],
+            outputs=[diarization_progress_info, diarization]
+        )
+        diarization_progress_info.change(
+            fn=merge_transcription_and_diarization,
+            outputs=[merged_transcription_progress_info, merged_transcription]
         )
-        # video_transcribed_progress_info.change(
-        #     fn=concatenate_transcriptions,
-        #     outputs=[transcriptions_concatenated_progress_info, original_audio_transcribed_path]
-        # )
         # transcriptions_concatenated_progress_info.change(
         #     fn=translate_transcription,
         #     inputs=[original_audio_transcribed_path, source_languaje, target_languaje],