This repository has no description
6.3 kB
200 lines
1#!/usr/bin/env python3
2"""
3Real-time streaming speech-to-text using Vosk.
4Reads audio from stdin and outputs words as they're recognized.
5"""
6
7import argparse
8import json
9import os
10import struct
11import sys
12import zipfile
13from pathlib import Path
14from urllib.request import urlretrieve
15
16from vosk import KaldiRecognizer, Model, SetLogLevel
17
18MODEL_DIR = Path(os.environ.get("XDG_DATA_HOME", Path.home() / ".local/share")) / "vosk"
19
20# Small, fast models for real-time streaming
21MODELS = {
22 "en": "vosk-model-small-en-us-0.15",
23 "de": "vosk-model-small-de-0.15",
24 "fr": "vosk-model-small-fr-0.22",
25 "es": "vosk-model-small-es-0.42",
26 "ru": "vosk-model-small-ru-0.22",
27 "cn": "vosk-model-small-cn-0.22",
28 "ja": "vosk-model-small-ja-0.22",
29}
30
31
32def download_model(lang: str) -> Path:
33 """Download vosk model if not present."""
34 model_name = MODELS.get(lang)
35 if not model_name:
36 # Fall back to English if language not available
37 print(f"No model for '{lang}', using English", file=sys.stderr)
38 model_name = MODELS["en"]
39 lang = "en"
40
41 model_path = MODEL_DIR / model_name
42 if model_path.exists():
43 return model_path
44
45 print(f"Downloading model {model_name}...", file=sys.stderr)
46 MODEL_DIR.mkdir(parents=True, exist_ok=True)
47
48 url = f"https://alphacephei.com/vosk/models/{model_name}.zip"
49 zip_path = MODEL_DIR / f"{model_name}.zip"
50
51 urlretrieve(url, zip_path)
52
53 print("Extracting...", file=sys.stderr)
54 with zipfile.ZipFile(zip_path, "r") as z:
55 z.extractall(MODEL_DIR)
56 zip_path.unlink()
57
58 print("Model ready", file=sys.stderr)
59 return model_path
60
61
62def parse_wav_header(data: bytes) -> tuple[int, int, int, int]:
63 """Parse WAV header, return (sample_rate, channels, bits_per_sample, data_offset)."""
64 if len(data) < 44 or data[:4] != b"RIFF" or data[8:12] != b"WAVE":
65 raise ValueError("Invalid WAV header")
66
67 pos = 12
68 sample_rate = channels = bits_per_sample = 0
69 while pos < len(data) - 8:
70 chunk_id = data[pos : pos + 4]
71 chunk_size = struct.unpack("<I", data[pos + 4 : pos + 8])[0]
72 if chunk_id == b"fmt ":
73 channels = struct.unpack("<H", data[pos + 10 : pos + 12])[0]
74 sample_rate = struct.unpack("<I", data[pos + 12 : pos + 16])[0]
75 bits_per_sample = struct.unpack("<H", data[pos + 22 : pos + 24])[0]
76 elif chunk_id == b"data":
77 return sample_rate, channels, bits_per_sample, pos + 8
78 pos += 8 + chunk_size
79
80 raise ValueError("No data chunk found")
81
82
83def convert_audio(data: bytes, in_rate: int, in_channels: int, in_bits: int) -> bytes:
84 """Convert audio to 16kHz mono 16-bit for Vosk."""
85 import array
86
87 # Align data to sample boundary
88 bytes_per_sample = in_bits // 8
89 frame_size = in_channels * bytes_per_sample
90 aligned_len = (len(data) // frame_size) * frame_size
91 if aligned_len == 0:
92 return b""
93 data = data[:aligned_len]
94
95 # Convert to samples
96 if in_bits == 16:
97 samples = array.array("h", data)
98 elif in_bits == 8:
99 samples = array.array("h", [(s - 128) * 256 for s in data])
100 else:
101 raise ValueError(f"Unsupported bit depth: {in_bits}")
102
103 # Convert to mono
104 if in_channels == 2:
105 mono = array.array("h")
106 for i in range(0, len(samples) - 1, 2):
107 mono.append((samples[i] + samples[i + 1]) // 2)
108 samples = mono
109
110 # Resample to 16kHz if needed
111 if in_rate != 16000:
112 ratio = 16000 / in_rate
113 new_len = int(len(samples) * ratio)
114 resampled = array.array("h")
115 for i in range(new_len):
116 src_idx = min(int(i / ratio), len(samples) - 1)
117 resampled.append(samples[src_idx])
118 samples = resampled
119
120 return samples.tobytes()
121
122
123def main():
124 parser = argparse.ArgumentParser(description="Real-time STT using Vosk")
125 parser.add_argument(
126 "-l",
127 "--language",
128 default=os.environ.get("VOSK_LANG", "en"),
129 help="Language: en, de, fr, es, ru, cn, ja (default: en)",
130 )
131 parser.add_argument(
132 "-w",
133 "--words",
134 action="store_true",
135 help="Output individual words as they're recognized",
136 )
137 args = parser.parse_args()
138
139 # Suppress Vosk logging
140 SetLogLevel(-1)
141
142 model_path = download_model(args.language)
143 model = Model(str(model_path))
144
145 # Read WAV header
146 stdin = sys.stdin.buffer
147 header_buf = stdin.read(4096)
148 sample_rate, channels, bits, data_offset = parse_wav_header(header_buf)
149
150 print(f"Audio: {sample_rate}Hz, {channels}ch, {bits}bit", file=sys.stderr)
151 print("Listening...", file=sys.stderr)
152
153 # Create recognizer for 16kHz audio
154 rec = KaldiRecognizer(model, 16000)
155 rec.SetWords(True)
156
157 # Process initial audio from header buffer
158 initial_audio = header_buf[data_offset:]
159 if initial_audio:
160 converted = convert_audio(initial_audio, sample_rate, channels, bits)
161 rec.AcceptWaveform(converted)
162
163 # Process audio stream
164 chunk_size = int(sample_rate * channels * (bits // 8) * 0.1) # 100ms chunks
165
166 while True:
167 data = stdin.read(chunk_size)
168 if not data:
169 break
170
171 converted = convert_audio(data, sample_rate, channels, bits)
172
173 if rec.AcceptWaveform(converted):
174 result = json.loads(rec.Result())
175 text = result.get("text", "")
176 if text:
177 if args.words:
178 # Show final sentence on stderr with period and newline
179 print(f"\r\033[K{text}.", file=sys.stderr, flush=True)
180 print(text, flush=True)
181 elif args.words:
182 partial = json.loads(rec.PartialResult())
183 text = partial.get("partial", "")
184 if text:
185 # Show partial with corrections on stderr
186 print(f"\r\033[K{text}", end="", file=sys.stderr, flush=True)
187
188 # Final result
189 result = json.loads(rec.FinalResult())
190 text = result.get("text", "")
191 if text:
192 if args.words:
193 print(f"\r\033[K{text}.", file=sys.stderr, flush=True)
194 print(text, flush=True)
195
196 print("Stream ended", file=sys.stderr)
197
198
199if __name__ == "__main__":
200 main()