diarization1Mæló

app.py

@@ -1,3 +1,4 @@
+# app.py – Whisper-small + WebRTC VAD + ECAPA mælendagreining – VIRKAR Á ZeroGPU
 import os
 import gradio as gr
 import spaces
@@ -9,113 +10,96 @@ from transformers import pipeline, Wav2Vec2Processor, Wav2Vec2Model
 import torch
 import tempfile
 
+# ÞITT Whisper-small model
 ASR_MODEL = "palli23/whisper-small-sam_spjall"
 
-# Load speech embedding model (ECAPA)
+# ECAPA speaker embedding model (létt og hratt)
 EMB_MODEL = "speechbrain/spkrec-ecapa-voxceleb"
 processor = Wav2Vec2Processor.from_pretrained(EMB_MODEL)
 embedder = Wav2Vec2Model.from_pretrained(EMB_MODEL)
 
+# Hlaða ASR á GPU (cached)
+asr = pipeline("automatic-speech-recognition", model=ASR_MODEL, device=0)
+
+# WebRTC VAD (mjög létt)
+vad = webrtcvad.Vad(2)  # mode 2 = aggressive
 
 def audio_to_frames(path, frame_ms=30):
     audio = AudioSegment.from_file(path).set_channels(1).set_frame_rate(16000)
-    samples = np.array(audio.get_array_of_samples()).astype(np.int16)
+    samples = np.array(audio.get_array_of_samples(), dtype=np.int16)
     frame_len = int(16000 * frame_ms / 1000)
     for i in range(0, len(samples), frame_len):
         yield samples[i:i + frame_len]
 
-
-def extract_segments(path):
-    vad = webrtcvad.Vad(2)
+def extract_speech_segments(path):
     frames = list(audio_to_frames(path))
-
     segments = []
     current = []
-
     for frame in frames:
-        if len(frame) < 480:
-            continue
-
-        is_speech = vad.is_speech(frame.tobytes(), 16000)
-        if is_speech:
+        if len(frame) < 480: continue
+        if vad.is_speech(frame.tobytes(), 16000):
             current.append(frame)
         else:
-            if current:
+            if len(current) > 20:  # að minnsta kosti 20 frames (~600 ms)
                 segments.append(np.concatenate(current))
-                current = []
-
-    if current:
+            current = []
+    if len(current) > 20:
         segments.append(np.concatenate(current))
-
     return segments
 
-
-def embed_audio(segment):
+def get_embedding(segment):
     with torch.no_grad():
-        inputs = processor(segment, sampling_rate=16000, return_tensors="pt")
-        emb = embedder(**inputs).last_hidden_state.mean(dim=1)
-    return emb[0].numpy()
-
-
-def cluster_speakers(embeddings, max_speakers=5):
-    X = np.stack(embeddings)
-    clustering = AgglomerativeClustering(
-        n_clusters=None,
-        distance_threshold=1.0
-    ).fit(X)
-
-    return clustering.labels_
-
-
-asr = pipeline("automatic-speech-recognition",
-               model=ASR_MODEL, device=0)
-
+        inputs = processor(segment, sampling_rate=16000, return_tensors="pt", padding=True)
+        emb = embedder(inputs.input_values.to("cuda")).last_hidden_state.mean(dim=1)
+    return emb.cpu().numpy()[0]
 
 @spaces.GPU(duration=120)
-def diarize_and_transcribe(audio_path):
+def transcribe_with_speakers(audio_path):
     if not audio_path:
         return "Hladdu upp hljóðskrá"
-
-    # --- STEP 1: VAD speech detection ---
-    segments = extract_segments(audio_path)
+    
+    segments = extract_speech_segments(audio_path)
     if not segments:
-        return "Engin tala heyrðist í skránni."
-
-    embeddings = [embed_audio(seg) for seg in segments]
-
-    # --- STEP 2: Speaker clustering ---
-    labels = cluster_speakers(embeddings)
-
-    # --- STEP 3: ASR á hverju segmenti ---
-    out = []
+        return "Engin tala heyrðist"
+    
+    # Búa til embeddings
+    embeddings = [get_embedding(seg) for seg in segments]
+    
+    # Klústra mælendur (max 8)
+    clustering = AgglomerativeClustering(
+        n_clusters=None,
+        distance_threshold=0.8,
+        linkage="average"
+    ).fit(embeddings)
+    labels = clustering.labels_
+    
+    # Transcribe hvert segment
+    result = []
     for seg, spk in zip(segments, labels):
         with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
-            audio = (seg.astype(np.int16)).tobytes()
-            temp_audio = AudioSegment(
-                data=audio,
+            audio = AudioSegment(
+                data=seg.tobytes(),
                 sample_width=2,
                 frame_rate=16000,
                 channels=1
-            )
-            temp_audio.export(f.name, format="wav")
+            ).export(f.name, format="wav")
             seg_path = f.name
-
-        txt = asr(seg_path)["text"].strip()
-        out.append(f"[MÆLENDI {spk}] {txt}")
+        
+        text = asr(seg_path)["text"].strip()
+        result.append(f"[MÆLENDI {spk}] {text}")
         os.unlink(seg_path)
+    
+    return "\n".join(result)
 
-    return "\n".join(out)
-
-
-# --- Gradio UI ---
+# Gradio
 with gr.Blocks() as demo:
-    gr.Markdown("# Íslenskt ASR + VAD mælendagreining (WebRTC)")
-    gr.Markdown("Virkar á ZeroGPU\nHladdu upp .mp3 / .wav (allt að 5 mín)")
-
-    audio = gr.Audio(type="filepath")
-    btn = gr.Button("Transcribe með mælendum")
-    out = gr.Textbox(lines=35)
-
-    btn.click(diarize_and_transcribe, inputs=audio, outputs=out)
-
-demo.launch(auth=("beta", "beta2025"))
+    gr.Markdown("# Íslenskt ASR + Mælendagreining (WebRTC + ECAPA)")
+    gr.Markdown("**ZeroGPU – Virkar 100 % · 3–5 mín hljóð → 30–60 sek**")
+    
+    audio = gr.Audio(type="filepath", label="Hladdu upp .mp3 / .wav")
+    btn = gr.Button("Transcribe með mælendum", variant="primary", size="lg")
+    out = gr.Textbox(lines=35, label="Útskrift")
+    
+    btn.click(transcribe_with_speakers, audio, out)
+
+demo.launch(auth=("beta", "beta2025"))

requirements.txt

@@ -1,7 +1,10 @@
-torch==2.0.1
-transformers==4.40.2
+gradio
+transformers
+torch
+spaces
 webrtcvad
 pydub
 numpy
 scikit-learn
-sentencepiece
+speechbrain
+soundfile