Upload 3 files

app.py

@@ -0,0 +1,148 @@
+"""
+Gradio Web UI for Speech Annotation Pipeline
+Run: python app.py
+"""
+
+import json
+import csv
+import io
+import os
+import tempfile
+from pathlib import Path
+
+import gradio as gr
+from pipeline import get_pipeline, to_json, to_csv, to_abab_text
+
+UPLOAD_DIR = Path(tempfile.gettempdir()) / "speech_annotation"
+UPLOAD_DIR.mkdir(parents=True, exist_ok=True)
+
+# ── State ─────────────────────────────────────────────────────────────────────
+_last_segments = []     # module-level cache for export
+
+
+# ── Core processing ───────────────────────────────────────────────────────────
+def process_audio(audio_path: str):
+    global _last_segments
+    if audio_path is None:
+        return "⚠️ Please upload an audio file.", "", None, None
+
+    try:
+        pipeline = get_pipeline()
+        segments = pipeline.process(audio_path)
+    except EnvironmentError as e:
+        return f"❌ Config error:\n{e}", "", None, None
+    except Exception as e:
+        return f"❌ Pipeline error:\n{e}", "", None, None
+
+    if not segments:
+        return "⚠️ No speech detected.", "", None, None
+
+    _last_segments = segments
+    abab = to_abab_text(segments)
+
+    # Build HTML table
+    rows = "".join(
+        f"<tr><td><b>{s.speaker}</b></td><td>{s.start_fmt}</td><td>{s.end_fmt}</td><td>{s.text}</td></tr>"
+        for s in segments
+    )
+    table = f"""
+    <table style='width:100%;border-collapse:collapse;font-family:monospace;font-size:14px'>
+      <thead style='background:#1e3a5f;color:white'>
+        <tr>
+          <th style='padding:8px;text-align:left'>Speaker</th>
+          <th style='padding:8px'>Start</th>
+          <th style='padding:8px'>End</th>
+          <th style='padding:8px;text-align:left'>Transcript</th>
+        </tr>
+      </thead>
+      <tbody>
+        {''.join(f"<tr style='background:{'#f0f4ff' if i%2==0 else 'white'}'>" +
+                 f"<td style='padding:6px 8px;font-weight:bold;color:#1e3a5f'>{s.speaker}</td>" +
+                 f"<td style='padding:6px 8px;color:#555'>{s.start_fmt}</td>" +
+                 f"<td style='padding:6px 8px;color:#555'>{s.end_fmt}</td>" +
+                 f"<td style='padding:6px 8px'>{s.text}</td></tr>"
+                 for i, s in enumerate(segments))}
+      </tbody>
+    </table>
+    """
+
+    status = f"✅ Done — {len(segments)} segments detected"
+    return status, abab, table
+
+
+def export_json():
+    if not _last_segments:
+        return None
+    out = UPLOAD_DIR / "annotation_export.json"
+    to_json(_last_segments, str(out))
+    return str(out)
+
+
+def export_csv():
+    if not _last_segments:
+        return None
+    out = UPLOAD_DIR / "annotation_export.csv"
+    to_csv(_last_segments, str(out))
+    return str(out)
+
+
+# ── UI ────────────────────────────────────────────────────────────────────────
+CSS = """
+#title { text-align: center; font-size: 28px; font-weight: bold; color: #1e3a5f; margin-bottom: 4px; }
+#subtitle { text-align: center; color: #666; margin-bottom: 20px; }
+.gr-button-primary { background: #1e3a5f !important; }
+"""
+
+with gr.Blocks(css=CSS, title="Speech Annotation Pipeline") as demo:
+
+    gr.HTML("<div id='title'>🎙️ Speech Annotation Pipeline</div>")
+    gr.HTML("<div id='subtitle'>Upload audio → get speaker-labeled transcripts automatically</div>")
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            audio_input = gr.Audio(
+                label="Upload Audio (.wav / .mp3 / .flac)",
+                type="filepath",
+            )
+            run_btn = gr.Button("▶ Run Annotation", variant="primary")
+            status_box = gr.Textbox(label="Status", interactive=False, lines=1)
+
+            gr.Markdown("### 📥 Export Results")
+            with gr.Row():
+                json_btn = gr.Button("Export JSON")
+                csv_btn = gr.Button("Export CSV")
+            json_file = gr.File(label="JSON Download", visible=True)
+            csv_file = gr.File(label="CSV Download", visible=True)
+
+        with gr.Column(scale=2):
+            gr.Markdown("### 📝 Editable Transcript (ABAB format)")
+            abab_box = gr.Textbox(
+                label="Edit transcript below",
+                lines=20,
+                placeholder="Annotations will appear here after processing…",
+            )
+            gr.Markdown("### 📊 Segment Table")
+            table_html = gr.HTML()
+
+    # Events
+    run_btn.click(
+        fn=process_audio,
+        inputs=[audio_input],
+        outputs=[status_box, abab_box, table_html],
+    )
+
+    json_btn.click(fn=export_json, inputs=[], outputs=[json_file])
+    csv_btn.click(fn=export_csv, inputs=[], outputs=[csv_file])
+
+    gr.Markdown(
+        """
+        ---
+        **Tips:**
+        - Set `HF_TOKEN` env var before running (required for speaker diarization)
+        - GPU strongly recommended for files > 5 min
+        - Edits in the transcript box are yours to save — export when ready
+        """
+    )
+
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=7860, share=False)

pipeline.py

@@ -0,0 +1,206 @@
+"""
+Speech Annotation Pipeline
+- Voice Activity Detection (VAD) via pyannote
+- Speaker Diarization via pyannote.audio
+- Transcription via WhisperX
+- Combines into structured ABAB-style annotations
+"""
+
+import os
+import json
+import csv
+import time
+import logging
+from pathlib import Path
+from dataclasses import dataclass, asdict
+from typing import List, Optional
+
+import torch
+import whisperx
+from pyannote.audio import Pipeline as PyannotePipeline
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# ── Config ────────────────────────────────────────────────────────────────────
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+COMPUTE_TYPE = "float16" if DEVICE == "cuda" else "int8"
+HF_TOKEN = os.environ.get("HF_TOKEN", "")          # Required for pyannote models
+WHISPER_MODEL = os.environ.get("WHISPER_MODEL", "base")   # base / small / medium / large-v2
+
+
+# ── Data model ────────────────────────────────────────────────────────────────
+@dataclass
+class AnnotationSegment:
+    speaker: str          # A, B, C …
+    start: float          # seconds
+    end: float            # seconds
+    text: str
+
+    @property
+    def start_fmt(self) -> str:
+        m, s = divmod(int(self.start), 60)
+        return f"{m:02d}:{s:02d}"
+
+    @property
+    def end_fmt(self) -> str:
+        m, s = divmod(int(self.end), 60)
+        return f"{m:02d}:{s:02d}"
+
+    def to_abab(self) -> str:
+        return f"{self.speaker} {self.start_fmt} {self.text}"
+
+
+# ── Label mapper ──────────────────────────────────────────────────────────────
+def _build_speaker_map(raw_labels: List[str]) -> dict:
+    """Map pyannote speaker IDs (SPEAKER_00, …) to A, B, C, …"""
+    letters = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
+    seen, mapping = [], {}
+    for lbl in raw_labels:
+        if lbl not in seen:
+            seen.append(lbl)
+    for i, lbl in enumerate(seen):
+        mapping[lbl] = letters[i] if i < len(letters) else f"SPK{i}"
+    return mapping
+
+
+# ── Core pipeline ─────────────────────────────────────────────────────────────
+class SpeechAnnotationPipeline:
+    def __init__(self):
+        self._whisper_model = None
+        self._diarization_pipeline = None
+
+    # --- lazy loaders ---
+    def _load_whisper(self):
+        if self._whisper_model is None:
+            logger.info(f"Loading WhisperX model '{WHISPER_MODEL}' on {DEVICE} …")
+            self._whisper_model = whisperx.load_model(
+                WHISPER_MODEL, DEVICE, compute_type=COMPUTE_TYPE
+            )
+        return self._whisper_model
+
+    def _load_diarization(self):
+        if self._diarization_pipeline is None:
+            if not HF_TOKEN:
+                raise EnvironmentError(
+                    "HF_TOKEN environment variable is required for pyannote diarization. "
+                    "Get a token at https://huggingface.co/settings/tokens and accept the "
+                    "model licence at https://hf.co/pyannote/speaker-diarization-3.1"
+                )
+            logger.info("Loading pyannote diarization pipeline …")
+            self._diarization_pipeline = PyannotePipeline.from_pretrained(
+                "pyannote/speaker-diarization-3.1",
+                use_auth_token=HF_TOKEN,
+            )
+            if DEVICE == "cuda":
+                self._diarization_pipeline = self._diarization_pipeline.to(torch.device("cuda"))
+        return self._diarization_pipeline
+
+    # --- steps ---
+    def _transcribe(self, audio_path: str) -> dict:
+        logger.info("Step 1/3 – Transcribing with WhisperX …")
+        model = self._load_whisper()
+        audio = whisperx.load_audio(audio_path)
+        result = model.transcribe(audio, batch_size=16)
+
+        logger.info("Step 1b – Aligning word timestamps …")
+        align_model, metadata = whisperx.load_align_model(
+            language_code=result["language"], device=DEVICE
+        )
+        result = whisperx.align(
+            result["segments"], align_model, metadata, audio, DEVICE,
+            return_char_alignments=False,
+        )
+        return result
+
+    def _diarize(self, audio_path: str) -> list:
+        logger.info("Step 2/3 – Speaker diarization …")
+        pipeline = self._load_diarization()
+        diarization = pipeline(audio_path)
+        segments = []
+        for turn, _, speaker in diarization.itertracks(yield_label=True):
+            segments.append({"start": turn.start, "end": turn.end, "speaker": speaker})
+        return segments
+
+    def _assign_speakers(self, transcript_segments: list, diarization_segments: list) -> List[AnnotationSegment]:
+        """Assign speaker labels to transcript segments using maximum-overlap rule."""
+        logger.info("Step 3/3 – Merging transcription + diarization …")
+
+        raw_labels = [d["speaker"] for d in diarization_segments]
+        speaker_map = _build_speaker_map(raw_labels)
+
+        results = []
+        for seg in transcript_segments:
+            t_start, t_end = seg["start"], seg["end"]
+            text = seg.get("text", "").strip()
+            if not text:
+                continue
+
+            # Find diarization segment with max overlap
+            best_speaker, best_overlap = "A", 0.0
+            for d in diarization_segments:
+                overlap = max(0, min(t_end, d["end"]) - max(t_start, d["start"]))
+                if overlap > best_overlap:
+                    best_overlap = overlap
+                    best_speaker = d["speaker"]
+
+            results.append(AnnotationSegment(
+                speaker=speaker_map.get(best_speaker, best_speaker),
+                start=t_start,
+                end=t_end,
+                text=text,
+            ))
+        return results
+
+    # --- public API ---
+    def process(self, audio_path: str) -> List[AnnotationSegment]:
+        t0 = time.time()
+        audio_path = str(Path(audio_path).resolve())
+
+        transcript = self._transcribe(audio_path)
+        diarization = self._diarize(audio_path)
+        annotations = self._assign_speakers(transcript["segments"], diarization)
+
+        elapsed = time.time() - t0
+        logger.info(f"Done in {elapsed:.1f}s – {len(annotations)} segments")
+        return annotations
+
+
+# ── Export helpers ─────────────────────────────────────────────────────────────
+def to_json(segments: List[AnnotationSegment], path: str):
+    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: List[AnnotationSegment], path: str):
+    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: List[AnnotationSegment]) -> str:
+    return "\n".join(s.to_abab() for s in segments)
+
+
+# ── Singleton instance ─────────────────────────────────────────────────────────
+_pipeline_instance: Optional[SpeechAnnotationPipeline] = None
+
+
+def get_pipeline() -> SpeechAnnotationPipeline:
+    global _pipeline_instance
+    if _pipeline_instance is None:
+        _pipeline_instance = SpeechAnnotationPipeline()
+    return _pipeline_instance

requirements.txt

@@ -0,0 +1,23 @@
+# ── Core ML ───────────────────────────────────────────────────────────────────
+torch>=2.1.0
+torchaudio>=2.1.0
+
+# WhisperX (transcription + alignment)
+whisperx>=3.1.1
+
+# pyannote (diarization + VAD)
+pyannote.audio>=3.1.1
+
+# ── API ───────────────────────────────────────────────────────────────────────
+fastapi>=0.110.0
+uvicorn[standard]>=0.29.0
+python-multipart>=0.0.9      # required for FastAPI file uploads
+
+# ── UI ────────────────────────────────────────────────────────────────────────
+gradio>=4.26.0
+
+# ── Utilities ─────────────────────────────────────────────────────────────────
+pydantic>=2.0.0
+numpy>=1.24.0
+soundfile>=0.12.1
+librosa>=0.10.1