Browse Source

Added audio save

Efren Yevale Varela 1 month ago
parent
commit
e34b0c9fb2
8 changed files with 166 additions and 9 deletions
  1. 1 0
      .gitignore
  2. 22 7
      README.md
  3. 2 0
      config.example.toml
  4. 4 0
      src/stt_daemon/config.py
  5. 6 1
      src/stt_daemon/daemon.py
  6. 44 1
      src/stt_daemon/output.py
  7. 30 0
      tests/test_config.py
  8. 57 0
      tests/test_output.py

+ 1 - 0
.gitignore

@@ -3,6 +3,7 @@ venv/
 __pycache__/
 *.pyc
 transcripts/
+audio/
 models/
 *.onnx
 *.wav

+ 22 - 7
README.md

@@ -117,8 +117,10 @@ loopback_latency_msec = 1                      # Module latency (ms)
 monitor_output        = false                  # Loop sink.monitor → default sink (OFF = no feedback risk)
 
 [output]
-stdout    = true                      # Print transcripts to stdout
-jsonl_dir = "transcripts"             # Write JSONL files (one per day) to this dir
+stdout     = true                     # Print transcripts to stdout
+jsonl_dir  = "transcripts"            # Write JSONL files (one per day) to this dir
+save_audio = false                    # Also write each transcribed utterance as a 16 kHz mono WAV
+audio_dir  = "audio"                  # WAV destination (YYYY-MM-DD subfolders); only used if save_audio
 
 [sink]
 sink_name             = "VirtualMicSink"    # Name of the null sink to create
@@ -331,17 +333,30 @@ uv run stt-daemon clear --yes
 
 **Stdout** (human-readable):
 ```
-[14:23:45] (es) Hola, cómo estás?
-[14:23:52] (en) What time is it?
-[14:24:10] (es) Gracias por tu ayuda.
+[14:23:45] (es) Hola, cómo estás?  [rms=0.041 end_of_turn]
+[14:23:52] (en) What time is it?  [rms=0.038 end_of_turn]
+[14:24:10] (es) Gracias por tu ayuda.  [rms=0.052 pause]
 ```
 
 **JSONL** (one line per utterance):
 ```json
-{"ts": "2025-09-20T14:23:45.123456", "lang": "es", "lang_prob": 0.95, "duration_s": 2.5, "text": "Hola, cómo estás?"}
-{"ts": "2025-09-20T14:23:52.654321", "lang": "en", "lang_prob": 0.98, "duration_s": 1.8, "text": "What time is it?"}
+{"ts": "2025-09-20T14:23:45.123456", "lang": "es", "lang_prob": 0.95, "duration_s": 2.5, "rms": 0.04123, "end_reason": "end_of_turn", "text": "Hola, cómo estás?"}
+{"ts": "2025-09-20T14:23:52.654321", "lang": "en", "lang_prob": 0.98, "duration_s": 1.8, "rms": 0.03811, "end_reason": "pause", "text": "What time is it?"}
 ```
 
+`end_reason` is one of `end_of_turn` (full `silence_ms` elapsed), `pause` (chunked at
+`chunk_pause_ms` because it passed `soft_max_s`), `hard_cap` (`max_utterance_s`
+reached), or `flush` (shutdown mid-utterance).
+
+**Audio files** — when `save_audio = true`, each transcribed utterance is also
+written as a 16 kHz mono WAV under `audio_dir/YYYY-MM-DD/`, and the record gains
+an `"audio"` field with its path:
+```json
+{..., "end_reason": "pause", "text": "Hola", "audio": "audio/2025-09-20/142345-123-pause.wav"}
+```
+16 kHz mono 16-bit is ~32 KB/s (~1.9 MB/minute), so enable this only if you need
+the clips for auditing.
+
 ## Virtual Sink Provisioning: How It Works
 
 The daemon uses PulseAudio's module system to create a virtual microphone from multiple sources:

+ 2 - 0
config.example.toml

@@ -31,6 +31,8 @@ avg_logprob_min    = -1.0             # drop segments with confidence below this
 [output]
 stdout    = true
 jsonl_dir = "transcripts"             # one JSONL file per day: YYYY-MM-DD.jsonl
+save_audio = false                    # write each transcribed utterance as a 16 kHz mono WAV
+audio_dir  = "audio"                  # WAV destination (YYYY-MM-DD subfolders); used only if save_audio
 
 [sink]
 sink_name             = "VirtualMicSink"    # Name of the null sink to create

+ 4 - 0
src/stt_daemon/config.py

@@ -49,6 +49,8 @@ class OutputConfig:
     """Output configuration."""
     stdout: bool = True
     jsonl_dir: str | None = "transcripts"
+    save_audio: bool = False
+    audio_dir: str = "audio"
 
 
 @dataclass
@@ -151,6 +153,8 @@ def load_config(path: str | None = None) -> Config:
             cfg.output = OutputConfig(
                 stdout=output_data.get("stdout", cfg.output.stdout),
                 jsonl_dir=output_data.get("jsonl_dir", cfg.output.jsonl_dir),
+                save_audio=output_data.get("save_audio", cfg.output.save_audio),
+                audio_dir=output_data.get("audio_dir", cfg.output.audio_dir),
             )
 
         # Sink section

+ 6 - 1
src/stt_daemon/daemon.py

@@ -94,6 +94,9 @@ class Daemon:
         self.sink = TranscriptSink(
             stdout=self.cfg.output.stdout,
             jsonl_dir=self.cfg.output.jsonl_dir,
+            save_audio=self.cfg.output.save_audio,
+            audio_dir=self.cfg.output.audio_dir,
+            sample_rate=self.cfg.audio.sample_rate,
         )
 
         self.logger.info("Daemon ready. Listening...")
@@ -157,7 +160,9 @@ class Daemon:
                 if self.transcriber and self.sink:
                     try:
                         result = self.transcriber.transcribe(utt.audio)
-                        self.sink.emit(result, utt.duration_s, utt.rms, utt.reason)
+                        self.sink.emit(
+                            result, utt.duration_s, utt.rms, utt.reason, utt.audio
+                        )
                     except Exception as e:
                         self.logger.error(f"Transcription error: {e}")
 

+ 44 - 1
src/stt_daemon/output.py

@@ -1,24 +1,50 @@
 """Output sinks for transcripts."""
 
 import json
+import wave
 from datetime import datetime
 from pathlib import Path
 
+import numpy as np
+
 from .transcriber import TranscriptResult
 
 
+def write_wav(path: Path, audio: np.ndarray, sample_rate: int) -> None:
+    """Write float32 mono audio in [-1, 1] to a 16-bit PCM WAV file."""
+    pcm = (np.clip(audio, -1.0, 1.0) * 32767.0).astype(np.int16)
+    with wave.open(str(path), "wb") as wav:
+        wav.setnchannels(1)
+        wav.setsampwidth(2)
+        wav.setframerate(sample_rate)
+        wav.writeframes(pcm.tobytes())
+
+
 class TranscriptSink:
     """Emits transcripts to stdout and/or JSONL files."""
 
-    def __init__(self, stdout: bool = True, jsonl_dir: str | None = None):
+    def __init__(
+        self,
+        stdout: bool = True,
+        jsonl_dir: str | None = None,
+        save_audio: bool = False,
+        audio_dir: str = "audio",
+        sample_rate: int = 16000,
+    ):
         """Initialize output sinks.
 
         Args:
             stdout: If True, emit to stdout.
             jsonl_dir: If set, append JSONL lines to daily files in this directory.
+            save_audio: If True, write each transcript's audio as a WAV file.
+            audio_dir: Directory for WAV files (dated subfolders).
+            sample_rate: Audio sample rate, written into the WAV header.
         """
         self.stdout = stdout
         self.jsonl_dir = Path(jsonl_dir) if jsonl_dir else None
+        self.save_audio = save_audio
+        self.audio_dir = Path(audio_dir)
+        self.sample_rate = sample_rate
 
     def emit(
         self,
@@ -26,6 +52,7 @@ class TranscriptSink:
         duration_s: float,
         rms: float = 0.0,
         reason: str = "unknown",
+        audio: np.ndarray | None = None,
     ) -> None:
         """Emit a transcript result.
 
@@ -34,10 +61,15 @@ class TranscriptSink:
             duration_s: Duration of the utterance in seconds.
             rms: Root-mean-square level of the utterance audio.
             reason: Why the utterance ended (end_of_turn/pause/hard_cap/flush).
+            audio: Utterance samples (float32, mono); saved when save_audio is enabled.
         """
         if not result.text:
             return
 
+        audio_path = None
+        if self.save_audio and audio is not None:
+            audio_path = self._save_audio(audio, reason)
+
         # Stdout line: [HH:MM:SS] (lang) text  [rms=0.041 end_of_turn]
         if self.stdout:
             now = datetime.now()
@@ -60,11 +92,22 @@ class TranscriptSink:
                 "end_reason": reason,
                 "text": result.text,
             }
+            if audio_path is not None:
+                record["audio"] = str(audio_path)
 
             with open(jsonl_path, "a") as f:
                 json.dump(record, f, ensure_ascii=False)
                 f.write("\n")
 
+    def _save_audio(self, audio: np.ndarray, reason: str) -> Path:
+        now = datetime.now()
+        day_dir = self.audio_dir / now.strftime("%Y-%m-%d")
+        day_dir.mkdir(parents=True, exist_ok=True)
+        name = f"{now.strftime('%H%M%S')}-{now.strftime('%f')[:3]}-{reason}.wav"
+        path = day_dir / name
+        write_wav(path, audio, self.sample_rate)
+        return path
+
     def close(self) -> None:
         """Clean up (no-op for file-based output)."""
         pass

+ 30 - 0
tests/test_config.py

@@ -186,3 +186,33 @@ def test_chunking_defaults():
     assert cfg.vad.trailing_trim_ms == 200
     assert cfg.whisper.no_speech_prob_max == 0.6
     assert cfg.whisper.avg_logprob_min == -1.0
+
+
+def test_output_defaults_disable_audio_saving():
+    with tempfile.TemporaryDirectory() as tmpdir:
+        import os
+        old_cwd = os.getcwd()
+        try:
+            os.chdir(tmpdir)
+            cfg = load_config()
+        finally:
+            os.chdir(old_cwd)
+
+    assert cfg.output.save_audio is False
+    assert cfg.output.audio_dir == "audio"
+
+
+def test_output_save_audio_fields():
+    config_toml = """
+[output]
+save_audio = true
+audio_dir = "clips"
+"""
+    with tempfile.TemporaryDirectory() as tmpdir:
+        config_path = Path(tmpdir) / "config.toml"
+        config_path.write_text(config_toml)
+
+        cfg = load_config(str(config_path))
+
+    assert cfg.output.save_audio is True
+    assert cfg.output.audio_dir == "clips"

+ 57 - 0
tests/test_output.py

@@ -1,6 +1,9 @@
 """Tests for transcript output sinks."""
 
 import json
+import wave
+
+import numpy as np
 
 from stt_daemon.output import TranscriptSink
 from stt_daemon.transcriber import TranscriptResult
@@ -41,3 +44,57 @@ def test_empty_text_is_skipped(tmp_path, capsys):
 
     assert capsys.readouterr().out == ""
     assert list(tmp_path.glob("*.jsonl")) == []
+
+
+def test_saves_wav_when_enabled(tmp_path):
+    audio = np.zeros(16000, dtype=np.float32)
+    audio[:8000] = 0.5
+    logs = tmp_path / "logs"
+
+    sink = TranscriptSink(
+        stdout=False,
+        jsonl_dir=str(logs),
+        save_audio=True,
+        audio_dir=str(tmp_path / "audio"),
+        sample_rate=16000,
+    )
+    sink.emit(_result(), 1.0, rms=0.3, reason="pause", audio=audio)
+
+    wavs = list((tmp_path / "audio").rglob("*.wav"))
+    assert len(wavs) == 1
+
+    with wave.open(str(wavs[0]), "rb") as wf:
+        assert wf.getnchannels() == 1
+        assert wf.getsampwidth() == 2
+        assert wf.getframerate() == 16000
+        assert wf.getnframes() == 16000
+
+    record = json.loads(next(logs.glob("*.jsonl")).read_text().splitlines()[0])
+    assert record["audio"] == str(wavs[0])
+
+
+def test_no_wav_when_disabled(tmp_path):
+    sink = TranscriptSink(
+        stdout=False,
+        jsonl_dir=None,
+        save_audio=False,
+        audio_dir=str(tmp_path / "audio"),
+        sample_rate=16000,
+    )
+    sink.emit(_result(), 1.0, rms=0.3, reason="pause", audio=np.zeros(1600, dtype=np.float32))
+
+    assert not (tmp_path / "audio").exists()
+
+
+def test_save_audio_without_samples_is_safe(tmp_path):
+    sink = TranscriptSink(
+        stdout=False,
+        jsonl_dir=None,
+        save_audio=True,
+        audio_dir=str(tmp_path / "audio"),
+        sample_rate=16000,
+    )
+    sink.emit(_result(), 1.0, rms=0.3, reason="pause", audio=None)
+
+    assert not (tmp_path / "audio").exists()
+