|
import streamlit as st |
|
from faster_whisper import WhisperModel |
|
from pyannote.audio import Pipeline |
|
import pyannote.core |
|
from collections import defaultdict |
|
|
|
|
|
def assign_speakers_to_segments(diarization, transcription_segments): |
|
speaker_segments = [] |
|
|
|
|
|
diarization_dict = defaultdict(list) |
|
for segment, _, speaker in diarization.itertracks(yield_label=True): |
|
diarization_dict[speaker].append((segment.start, segment.end)) |
|
|
|
for transcription_segment in transcription_segments: |
|
start, end = transcription_segment.start, transcription_segment.end |
|
speakers_count = defaultdict(float) |
|
|
|
|
|
for speaker, times in diarization_dict.items(): |
|
for seg_start, seg_end in times: |
|
|
|
overlap_start = max(start, seg_start) |
|
overlap_end = min(end, seg_end) |
|
overlap_duration = max(0, overlap_end - overlap_start) |
|
speakers_count[speaker] += overlap_duration |
|
|
|
|
|
if speakers_count: |
|
speaker = max(speakers_count, key=speakers_count.get) |
|
else: |
|
speaker = "Unknown" |
|
|
|
|
|
speaker_segments.append((speaker, transcription_segment.text)) |
|
|
|
return speaker_segments |
|
|
|
|
|
def main(): |
|
st.title("Aplicación de Diarización y Transcripción de Audio") |
|
|
|
|
|
audio_file = st.file_uploader("Cargar archivo de audio", type=['wav']) |
|
|
|
if audio_file is not None: |
|
|
|
with open("audio_temp.wav", "wb") as f: |
|
f.write(audio_file.getbuffer()) |
|
|
|
|
|
speaker_segments = process_audio("audio_temp.wav") |
|
|
|
|
|
for speaker, text in speaker_segments: |
|
st.write(f"Speaker {speaker}: {text}") |
|
|
|
|
|
def process_audio(audio_path): |
|
|
|
pipeline = Pipeline.from_pretrained( |
|
"pyannote/speaker-diarization-3.1", |
|
use_auth_token=st.secrets["hfapi"] |
|
) |
|
|
|
|
|
model = WhisperModel("large-v3", device="cpu", compute_type="int8") |
|
|
|
|
|
diarization = pipeline(audio_path) |
|
segments, info = model.transcribe(audio_path, vad_filter=True, language="Spanish") |
|
|
|
|
|
return assign_speakers_to_segments(diarization, segments) |
|
|
|
if __name__ == "__main__": |
|
main() |
|
|