143 lines
4.5 KiB
Python
143 lines
4.5 KiB
Python
#!/usr/bin/env python3
|
|
import vosk
|
|
import sys
|
|
import json
|
|
import struct
|
|
import numpy as np
|
|
from queue import Queue
|
|
from threading import Thread
|
|
import soundfile as sf
|
|
import tempfile
|
|
import os
|
|
|
|
# Global recognizer
|
|
recognizer = None
|
|
audio_queue = Queue()
|
|
result_queue = Queue()
|
|
|
|
def initialize_vosk():
|
|
global recognizer
|
|
model_path = "vosk-model" # Update this path
|
|
|
|
if not os.path.exists(model_path):
|
|
return {"success": False, "error": "Model not found"}
|
|
|
|
try:
|
|
vosk.SetLogLevel(-1)
|
|
model = vosk.Model(model_path)
|
|
recognizer = vosk.KaldiRecognizer(model, 16000)
|
|
return {"success": True}
|
|
except Exception as e:
|
|
return {"success": False, "error": str(e)}
|
|
|
|
def audio_worker():
|
|
global recognizer
|
|
while True:
|
|
audio_data, request_id = audio_queue.get()
|
|
|
|
try:
|
|
# Write to temp file and read with soundfile
|
|
with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as f:
|
|
f.write(audio_data)
|
|
temp_path = f.name
|
|
|
|
try:
|
|
data, samplerate = sf.read(temp_path, dtype='float32')
|
|
|
|
# Resample if needed
|
|
if samplerate != 16000:
|
|
duration = len(data) / samplerate
|
|
data = np.interp(
|
|
np.linspace(0, len(data)-1, int(duration * 16000)),
|
|
np.arange(len(data)),
|
|
data
|
|
)
|
|
|
|
# Convert to 16-bit PCM
|
|
data = (data * 32767).astype('int16')
|
|
|
|
# Process with Vosk
|
|
if recognizer.AcceptWaveform(data.tobytes()):
|
|
text = json.loads(recognizer.Result()).get('text', '')
|
|
is_final = True
|
|
else:
|
|
text = json.loads(recognizer.PartialResult()).get('partial', '')
|
|
is_final = False
|
|
|
|
result_queue.put(({
|
|
"success": True,
|
|
"text": text,
|
|
"is_final": is_final,
|
|
"requestId": request_id
|
|
}, request_id))
|
|
|
|
finally:
|
|
os.unlink(temp_path)
|
|
|
|
except Exception as e:
|
|
result_queue.put(({
|
|
"success": False,
|
|
"error": str(e),
|
|
"requestId": request_id
|
|
}, request_id))
|
|
|
|
def main():
|
|
# Initialize Vosk
|
|
init_result = initialize_vosk()
|
|
if not init_result["success"]:
|
|
error = json.dumps({
|
|
"success": False,
|
|
"error": init_result["error"],
|
|
"requestId": 0
|
|
}).encode()
|
|
sys.stdout.buffer.write(struct.pack('>I', len(error)))
|
|
sys.stdout.buffer.write(error)
|
|
sys.stdout.buffer.flush()
|
|
return
|
|
|
|
# Start worker thread
|
|
Thread(target=audio_worker, daemon=True).start()
|
|
|
|
while True:
|
|
try:
|
|
# Read message length (4 bytes)
|
|
length_bytes = sys.stdin.buffer.read(4)
|
|
if not length_bytes:
|
|
break
|
|
length = struct.unpack('>I', length_bytes)[0]
|
|
|
|
# Read request ID (4 bytes)
|
|
id_bytes = sys.stdin.buffer.read(4)
|
|
if not id_bytes:
|
|
break
|
|
request_id = struct.unpack('>I', id_bytes)[0]
|
|
|
|
# Read audio data
|
|
audio_data = sys.stdin.buffer.read(length)
|
|
if len(audio_data) != length:
|
|
break
|
|
|
|
# Add to processing queue
|
|
audio_queue.put((audio_data, request_id))
|
|
|
|
# Check for results
|
|
while not result_queue.empty():
|
|
result, res_id = result_queue.get()
|
|
response = json.dumps(result).encode()
|
|
sys.stdout.buffer.write(struct.pack('>I', len(response)))
|
|
sys.stdout.buffer.write(struct.pack('>I', res_id)))
|
|
sys.stdout.buffer.write(response)
|
|
sys.stdout.buffer.flush()
|
|
|
|
except Exception as e:
|
|
error = json.dumps({
|
|
"success": False,
|
|
"error": str(e),
|
|
"requestId": request_id if 'request_id' in locals() else 0
|
|
}).encode()
|
|
sys.stdout.buffer.write(struct.pack('>I', len(error)))
|
|
sys.stdout.buffer.write(error)
|
|
sys.stdout.buffer.flush()
|
|
|
|
if __name__ == "__main__":
|
|
main() |