Update app.py

app.py

@@ -3,7 +3,8 @@ import pyannote.audio
 import streamlit as st
 from transformers import pipeline
 from dotenv import load_dotenv
-from utils_hf import extract_audio_segment, transcribe_audio, download_youtube_audio, get_audio_duration, StreamlitProgressHook
+from utils_hf import ( extract_audio_segment, transcribe_audio, download_youtube_audio, get_audio_duration, 
+                      StreamlitProgressHook, diarize_short_segment )
 
 load_dotenv()
 
@@ -30,38 +31,17 @@ if st.button("Transcribir y Diarizar"):
                 st.error(f"Error al procesar el video de YouTube: {e}")
                 st.stop()
 
-        with st.spinner("Realizando diarización..."):
+        with st.spinner("Realizando diarización en un segmento..."):
             try:
-                pipeline = pyannote.audio.Pipeline.from_pretrained("pyannote/speaker-diarization-3.1",
-                                                                   use_auth_token=token)
+                diarization_segments, diarization_text, speakers = diarize_short_segment(audio_path, token, start_time=0, end_time=600)
 
-                with StreamlitProgressHook() as hook:
-                    diarization = pipeline(audio_path, hook=hook)
-
-                diarization_segments = []
-                diarization_text = ""
-                speakers = set()
-                for segment, _, speaker in diarization.itertracks(yield_label=True):
-                    diarization_segments.append((segment.start, segment.end, speaker))
-                    diarization_text += f"{segment.start:.2f} - {segment.end:.2f}: {speaker}\n"
-                    speakers.add(speaker)
-
-                num_speakers = len(speakers)
-
-                path = 'diarization.txt'
-                if os.path.exists(path):
-                    os.remove(path)
-                    st.success("Borrado archivo diarization.txt anterior")
-
-                with open("diarization.txt", "w") as file:
-                    file.write(diarization_text)
-
-                st.text_area("Diarización", diarization_text, height=200)
-                st.success(
-                    f"Diarización completada y guardada en diarization.txt. Número de hablantes detectados: {num_speakers}")
+                if diarization_text:
+                    st.text_area("Diarización", diarization_text, height=200)
+                    st.success(f"Diarización completada. Número de hablantes detectados: {len(speakers)}")
+                else:
+                    st.error("No se pudo realizar la diarización.")
             except Exception as e:
                 st.error(f"Error al realizar la diarización: {e}")
-                st.stop()
 
         with st.spinner("Transcribiendo segmentos de audio..."):
             try: