Update app.py

app.py

@@ -1,88 +1,89 @@
-import gradio as gr
-import torch
 import os
-from whisperx import load_model, load_align_model, align
-from resemblyzer import preprocess_wav, VoiceEncoder
-from sklearn.cluster import AgglomerativeClustering
-import librosa
-import numpy as np
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-compute_type = "float16" if device == "cuda" else "int8"
-
-whisper_model = load_model("medium", device=device, compute_type=compute_type)
-align_model, metadata = load_align_model(language_code="fr", device=device)
-voice_encoder = VoiceEncoder()
-
-def get_speaker_segments(audio_path, window_size=1.0, step_size=0.5, num_speakers=2):
-    wav, sr = librosa.load(audio_path, sr=16000, mono=True)
-    wav = librosa.util.normalize(wav)
-    duration = librosa.get_duration(y=wav, sr=sr)
-
-    segments = []
-    embeddings = []
-
-    for start in np.arange(0, duration - window_size, step_size):
-        end = start + window_size
-        clip = wav[int(start * sr):int(end * sr)]
-        if len(clip) == 0:
-            continue
-        try:
-            embed = voice_encoder.embed_utterance(clip)
-            embeddings.append(embed)
-            segments.append((start, end))
-        except Exception as e:
-            print(f"⚠️ Skipped segment {start}-{end}s: {e}")
-
-    if len(embeddings) < 2:
-        print("⚠️ Pas assez de segments pour la diarisation. Diarisation annulée.")
-        return [{"start": 0, "end": duration, "speaker": "speaker_00"}]
-
-    clustering = AgglomerativeClustering(n_clusters=num_speakers)
-    labels = clustering.fit_predict(embeddings)
-
-    speaker_segments = []
-    for (start, end), label in zip(segments, labels):
-        speaker_segments.append({"start": start, "end": end, "speaker": f"speaker_{label:02d}"})
-
-    return speaker_segments
-
-def process_audio(audio_file):
-    tmp_path = audio_file
-
-    # Step 1: Transcription
-    result = whisper_model.transcribe(tmp_path, language="fr", word_timestamps=False, verbose=False)
-
-    # Step 2: Diarisation via resemblyzer
-    speaker_segments = get_speaker_segments(tmp_path)
-
-    # Step 3: Alignement mot à mot
-    result_aligned = align(result["segments"], align_model, metadata, tmp_path, return_char_alignments=False)
-
-    # Attribution speaker
-    for segment in result_aligned["segments"]:
-        segment_start = segment["start"]
-        speaker_found = next((sp["speaker"] for sp in speaker_segments if sp["start"] <= segment_start <= sp["end"]), "speaker_??")
-        segment["speaker"] = speaker_found
-
-    # Format final
-    final_output = ""
-    for seg in result_aligned["segments"]:
-        speaker = seg["speaker"]
-        start = f"{seg['start']:.2f}s"
-        end = f"{seg['end']:.2f}s"
-        text = seg['text'].strip()
-        final_output += f"[{start} - {end}] {speaker}: {text}\n"
+import tempfile
+import pandas as pd
+import gradio as gr
+from pydub import AudioSegment
+from faster_whisper import WhisperModel
+from pyannote.audio import Pipeline as DiarizationPipeline
 
-    return final_output
+# Initialisation des modèles
+whisper_model = WhisperModel("large-v2", device="cpu", compute_type="int8")
 
-iface = gr.Interface(
-    fn=process_audio,
-    inputs=gr.Audio(type="filepath", label="Audio (.wav, .mp3...)"),
-    outputs=gr.Textbox(label="Transcription + Diarisation + Alignement"),
-    title="🎙️ Transcription enrichie avec WhisperX + Resemblyzer",
-    description="Transcription française, diarisation légère (sans token), alignement mot à mot."
+diari_pipeline = DiarizationPipeline.from_pretrained(
+    "pyannote/speaker-diarization-3.1",
+    use_auth_token="hf_YOUR_TOKEN_HERE"  # Remplace par ton token Hugging Face perso
 )
 
-if __name__ == "__main__":
-    iface.launch()
+def convert_mp3_to_wav(mp3_path):
+    wav_path = tempfile.mktemp(suffix=".wav")
+    audio = AudioSegment.from_file(mp3_path, format="mp3")
+    audio = audio.set_channels(1).set_frame_rate(16000)
+    audio.export(wav_path, format="wav")
+    return wav_path
+
+def transcribe_and_diarize(audio_file):
+    wav_path = convert_mp3_to_wav(audio_file)
+
+    # Transcription avec Whisper
+    segments, _ = whisper_model.transcribe(wav_path, language="fr", beam_size=5)
+
+    transcript = []
+    for seg in segments:
+        transcript.append({
+            "start": seg.start,
+            "end": seg.end,
+            "text": seg.text.strip()
+        })
+
+    # Diarisation avec pyannote
+    diarization = diari_pipeline(wav_path)
+    speakers = []
+    for turn, _, speaker in diarization.itertracks(yield_label=True):
+        speakers.append({
+            "start": turn.start,
+            "end": turn.end,
+            "speaker": speaker
+        })
+
+    # Fusion transcription + speaker
+    final_output = []
+    for t in transcript:
+        speaker = "Inconnu"
+        for d in speakers:
+            if d["start"] <= t["start"] <= d["end"]:
+                speaker = d["speaker"]
+                break
+        final_output.append({
+            "start": t["start"],
+            "end": t["end"],
+            "speaker": speaker,
+            "text": t["text"]
+        })
+
+    df = pd.DataFrame(final_output)
+
+    # Export TXT format
+    txt_lines = [f"[{row['start']:.2f}s - {row['end']:.2f}s] {row['speaker']} : {row['text']}" for _, row in df.iterrows()]
+    txt_output = "\n".join(txt_lines)
+    txt_path = tempfile.mktemp(suffix=".txt")
+    with open(txt_path, "w", encoding="utf-8") as f:
+        f.write(txt_output)
+
+    # Export CSV format
+    csv_path = tempfile.mktemp(suffix=".csv")
+    df.to_csv(csv_path, index=False)
+
+    return txt_output, csv_path, txt_path
+
+# Interface Gradio
+gr.Interface(
+    fn=transcribe_and_diarize,
+    inputs=gr.Audio(type="filepath", label="Fichier audio MP3"),
+    outputs=[
+        gr.Textbox(label="Transcription avec locuteurs"),
+        gr.File(label="Télécharger le CSV"),
+        gr.File(label="Télécharger le TXT")
+    ],
+    title="Transcription + Diarisation (FR)",
+    description="Charge un fichier MP3. Transcription FR + séparation des locuteurs + export CSV et TXT."
+).launch()