Переглянути джерело

[1.0] Support of WebSocket

Penta 9 місяців тому
батько
коміт
2463f3acbb
1 змінених файлів з 68 додано та 4 видалено
  1. 68 4
      app.py

+ 68 - 4
app.py

@@ -1,16 +1,24 @@
-from fastapi import FastAPI, UploadFile, File
+from fastapi import FastAPI, UploadFile, File, WebSocket, WebSocketDisconnect
 from faster_whisper import WhisperModel
 import tempfile
 import os
+import numpy as np
+import time
 
 app = FastAPI()
 
+# =========================
+# Whisper model (shared)
+# =========================
 model = WhisperModel(
     "small",
     device="cuda",
-    compute_type="float32"
+    compute_type="float32"  # GTX 1050 Ti compatible
 )
 
+# =========================
+# REST API (Swagger / tests)
+# =========================
 @app.post("/transcribe")
 async def transcribe(file: UploadFile = File(...)):
     with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp:
@@ -25,9 +33,65 @@ async def transcribe(file: UploadFile = File(...)):
 
     os.remove(tmp_path)
 
-    text = " ".join([seg.text for seg in segments])
+    text = " ".join(seg.text for seg in segments)
 
     return {
         "text": text,
         "language": info.language
-    }
+    }
+
+# =========================
+# WebSocket STT (streaming)
+# =========================
+
+SAMPLE_RATE = 16000
+MIN_AUDIO_SEC = 0.5
+MAX_AUDIO_SEC = 10.0
+
+@app.websocket("/ws/transcribe")
+async def ws_transcribe(ws: WebSocket):
+    await ws.accept()
+
+    audio_buffer = []
+    start_time = time.time()
+
+    try:
+        while True:
+            data = await ws.receive_bytes()
+
+            # PCM int16 → float32
+            pcm = np.frombuffer(data, dtype=np.int16).astype(np.float32) / 32768.0
+            audio_buffer.append(pcm)
+
+            audio_len = sum(len(chunk) for chunk in audio_buffer) / SAMPLE_RATE
+
+            # Sécurité : reset si trop long
+            if audio_len > MAX_AUDIO_SEC:
+                audio_buffer.clear()
+                start_time = time.time()
+                continue
+
+            # Transcription dès qu'on a assez d'audio
+            if audio_len >= MIN_AUDIO_SEC:
+                audio = np.concatenate(audio_buffer)
+                audio_buffer.clear()
+
+                segments, _ = model.transcribe(
+                    audio,
+                    language="fr",
+                    vad_filter=True
+                )
+
+                text = " ".join(seg.text for seg in segments).strip()
+
+                if text:
+                    await ws.send_json({
+                        "type": "final",
+                        "text": text
+                    })
+
+    except WebSocketDisconnect:
+        pass
+    except Exception as e:
+        await ws.send_json({"error": str(e)})
+        await ws.close()