Update pipeline.py

pipeline.py

@@ -1,68 +1,116 @@
-import tempfile
-import traceback
+import os
+import json
+import csv
+import time
+import logging
+import httpx
 from pathlib import Path
-import gradio as gr
-from pipeline import get_pipeline, to_json, to_csv, to_abab_text
+from dataclasses import dataclass
+from typing import List, Optional
 
-UPLOAD_DIR = Path(tempfile.gettempdir()) / "speech_annotation"
-UPLOAD_DIR.mkdir(parents=True, exist_ok=True)
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
-_last_segments = []
+DEEPGRAM_API_KEY = os.environ.get("DEEPGRAM_API_KEY", "")
 
-def process_audio(audio_path, num_speakers):
-    global _last_segments
-    if audio_path is None:
-        return "⚠️ Please upload an audio file first.", ""
-    try:
-        pipeline = get_pipeline()
-        n = int(num_speakers) if num_speakers and int(num_speakers) > 0 else 0
-        segments = pipeline.process(audio_path, num_speakers=n)
-    except Exception as e:
-        return f"❌ Error: {e}\n{traceback.format_exc()}", ""
-    if not segments:
-        return "⚠️ No speech detected.", ""
-    _last_segments = segments
-    abab = to_abab_text(segments)
-    unique = len(set(s.speaker for s in segments))
-    return f"✅ Done — {len(segments)} segments, {unique} speaker(s) detected", abab
+@dataclass
+class AnnotationSegment:
+    speaker: str
+    start: float
+    end: float
+    text: str
 
-def export_json():
-    if not _last_segments:
-        return None
-    out = str(UPLOAD_DIR / "annotation.json")
-    to_json(_last_segments, out)
-    return out
+    @property
+    def start_fmt(self):
+        m, s = divmod(int(self.start), 60)
+        return f"{m:02d}:{s:02d}"
 
-def export_csv():
-    if not _last_segments:
-        return None
-    out = str(UPLOAD_DIR / "annotation.csv")
-    to_csv(_last_segments, out)
-    return out
+    @property
+    def end_fmt(self):
+        m, s = divmod(int(self.end), 60)
+        return f"{m:02d}:{s:02d}"
 
-with gr.Blocks(title="Speech Annotation Pipeline") as demo:
-    gr.Markdown("# 🎙️ Speech Annotation Pipeline")
-    audio_input = gr.Audio(label="Upload Audio (.wav/.mp3/.flac)", type="filepath")
-    num_speakers = gr.Slider(minimum=0, maximum=10, step=1, value=0,
-                             label="Number of speakers (0 = auto-detect)")
-    run_btn = gr.Button("▶ Run Annotation", variant="primary")
-    status_box = gr.Textbox(label="Status", value="Ready.", interactive=False)
-    abab_box = gr.Textbox(label="Transcript (ABAB)", lines=15, value="",
-                          placeholder="Annotations will appear here...")
-    with gr.Row():
-        json_btn = gr.Button("Export JSON")
-        csv_btn  = gr.Button("Export CSV")
-    json_file = gr.File(label="JSON")
-    csv_file  = gr.File(label="CSV")
+    def to_abab(self):
+        return f"{self.speaker} {self.start_fmt} {self.text}"
 
-    run_btn.click(
-        fn=process_audio,
-        inputs=[audio_input, num_speakers],
-        outputs=[status_box, abab_box],
-        queue=False
-    )
-    json_btn.click(fn=export_json, inputs=[], outputs=[json_file])
-    csv_btn.click(fn=export_csv, inputs=[], outputs=[csv_file])
+SPEAKER_LETTERS = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
 
-if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)
+def _speaker_label(idx):
+    return SPEAKER_LETTERS[idx] if idx < len(SPEAKER_LETTERS) else f"SPK{idx}"
+
+class SpeechAnnotationPipeline:
+    def process(self, audio_path, num_speakers=0, **kwargs):
+        if not DEEPGRAM_API_KEY:
+            raise EnvironmentError("DEEPGRAM_API_KEY not set.")
+        logger.info(f"Sending to Deepgram: {audio_path}")
+        t0 = time.time()
+        audio_bytes = Path(audio_path).read_bytes()
+        ext = Path(audio_path).suffix.lower().lstrip(".")
+        mime = {"wav":"audio/wav","mp3":"audio/mpeg","mp4":"audio/mp4",
+                "m4a":"audio/mp4","flac":"audio/flac","ogg":"audio/ogg"}.get(ext,"audio/wav")
+        params = {"model":"nova-2","diarize":"true","punctuate":"true",
+                  "utterances":"true","smart_format":"true"}
+        headers = {"Authorization": f"Token {DEEPGRAM_API_KEY}", "Content-Type": mime}
+        response = httpx.post("https://api.deepgram.com/v1/listen",
+                              params=params, headers=headers,
+                              content=audio_bytes, timeout=300)
+        response.raise_for_status()
+        data = response.json()
+        utterances = data.get("results", {}).get("utterances", [])
+        if not utterances:
+            words = data.get("results",{}).get("channels",[{}])[0]\
+                        .get("alternatives",[{}])[0].get("words",[])
+            utterances = self._words_to_utterances(words)
+        segments = []
+        for utt in utterances:
+            segments.append(AnnotationSegment(
+                speaker=_speaker_label(utt.get("speaker", 0)),
+                start=utt.get("start", 0),
+                end=utt.get("end", 0),
+                text=utt.get("transcript", "").strip(),
+            ))
+        logger.info(f"Done in {time.time()-t0:.1f}s — {len(segments)} segments")
+        return segments
+
+    def _words_to_utterances(self, words):
+        if not words:
+            return []
+        utterances = []
+        current = {"speaker": words[0].get("speaker",0), "start": words[0].get("start",0),
+                   "end": words[0].get("end",0),
+                   "transcript": words[0].get("punctuated_word", words[0].get("word",""))}
+        for word in words[1:]:
+            if word.get("speaker") == current["speaker"]:
+                current["end"] = word.get("end", current["end"])
+                current["transcript"] += " " + word.get("punctuated_word", word.get("word",""))
+            else:
+                utterances.append(current)
+                current = {"speaker": word.get("speaker",0), "start": word.get("start",0),
+                           "end": word.get("end",0),
+                           "transcript": word.get("punctuated_word", word.get("word",""))}
+        utterances.append(current)
+        return utterances
+
+def to_json(segments, path):
+    data = [{"speaker": s.speaker, "start": round(s.start,3), "end": round(s.end,3),
+             "start_fmt": s.start_fmt, "end_fmt": s.end_fmt, "text": s.text} for s in segments]
+    with open(path, "w", encoding="utf-8") as f:
+        json.dump(data, f, indent=2, ensure_ascii=False)
+
+def to_csv(segments, path):
+    with open(path, "w", newline="", encoding="utf-8") as f:
+        writer = csv.writer(f)
+        writer.writerow(["speaker","start","end","text"])
+        for s in segments:
+            writer.writerow([s.speaker, s.start_fmt, s.end_fmt, s.text])
+
+def to_abab_text(segments):
+    return "\n".join(s.to_abab() for s in segments)
+
+_instance = None
+
+def get_pipeline():
+    global _instance
+    if _instance is None:
+        _instance = SpeechAnnotationPipeline()
+    return _instance