Merge pull request #6 from makaveli10/vad_on_server

VAD on server.

Audio-Transcription/manifest.json

@@ -4,9 +4,7 @@
     "name": "Audio Transcription",
     "version": "1.0.0",
     "description": "This extension captures the audio on the current tab, sends it to a server for transcription and shows the transcription in Real-time.",
-    "content_security_policy": {
-        "extension_pages": "script-src 'self' 'wasm-unsafe-eval'; object-src 'self';"
-    },
+    
     "options_page": "options.html",
     "background": {
         "service_worker": "background.js"

Audio-Transcription/options.js

@@ -73,43 +73,13 @@ function resampleTo16kHZ(audioData, origSampleRate = 44100) {
  */
 async function startRecord(option) {
   const stream = await captureTabAudio();
-  var doVad = true;
+
   if (stream) {
     // call when the stream inactive
     stream.oninactive = () => {
       window.close();
     };
 
-    // create onnx model
-    // initialize onnx model
-    const session = await ort.InferenceSession.create('./silero_vad.onnx');
-    var h = new Array(128);
-    for (let i = 0; i < h.length; i++) {
-      h[i] = 0;
-    }
-    var c = new Array(128);
-    for (let i = 0; i < h.length; i++) {
-      c[i] = 0;
-    }
-    
-    const sr = new BigInt64Array(1)
-    sr[0] = BigInt(16000);
-    const srate = new ort.Tensor('int64', sr, [1]);
-    let speech_prob = undefined;
-    const vad_infer = async (feed_dict) => {
-      // feed inputs and run
-      try{
-        const results = await session.run(feed_dict);
-
-        // update states
-        h = results.hn.data
-        c = results.cn.data
-        speech_prob = results.output.data
-      } catch(e) {
-        console.log(e)
-      }
-    }
-
     const socket = new WebSocket("ws://localhost:9090/");
     socket.onopen = function(e) { 
       socket.send("handshake");
@@ -136,21 +106,8 @@ async function startRecord(option) {
 
       audioDataCache.push(inputData);
 
-      // voice activity detection inference
-      const audioBuffer = new ort.Tensor('float32', audioData16kHz, [1, audioData16kHz.length]);
-      const hh = new ort.Tensor('float32', h, [2, 1, 64]);
-      const hc = new ort.Tensor('float32', c, [2, 1, 64]);
-      const feeds = { input: audioBuffer, sr: srate, h: hh, c: hc};
-      
       // feed inputs and run
-      if (doVad) {
-        vad_infer(feeds)
-        if (speech_prob > 0.4) {
-          socket.send(audioData16kHz);
-        }
-        else
-          console.log("no speech found: " + speech_prob)
-      }
+      socket.send(audioData16kHz);
     };
 
     // Prevent page mute

README.md

@@ -30,24 +30,26 @@ Unlike traditional speech recognition systems that rely on continuous audio stre
 - On the client side
     - To transcribe an audio file:
     ```bash
-     python client.py --audio "audio.wav"
+     python client.py --audio "audio.wav" --host "localhost" --port "9090"
     ```
 
     - To transcribe from microphone:
     ```bash
-     python client.py
+     python client.py --host "localhost" --port "9090"
     ```
 
 
 ## Transcribe audio from browser
-- Run the websocket-server
+- Run the server
 ```bash
- python websocket_server.py
+ python server.py
 ```
 This would start the websocket server on port ```9090```.
 
+
 - Head over to ```Audio-Transcription``` module to unpack and load a chrome extension to capture any audio in the browser(only Chrome for now) and send it to the websocket server to transcribe the audio in the current tab.
 
+
 ## Future Work
 - [ ] Update Documentation.
 - [ ] Keep only a single server implementation i.e. websockets and get rid of the socket implementation in ```server.py```. Also, update ```client.py``` to websockets-client implemenation.

server.py

@@ -1,4 +1,3 @@
-# import asyncio
 import websockets
 import pickle, struct, time, pyaudio
 import threading
@@ -12,6 +11,7 @@ logging.basicConfig(level = logging.INFO)
 from collections import deque
 from dataclasses import dataclass
 
+import torch
 import numpy as np
 from websockets.sync import server
 from websockets.sync.server import serve
@@ -33,8 +33,7 @@ def recv_audio(websocket):
             if isinstance(frame_data, str):
                 logging.info(frame_data)
                 continue
-            else:
-                frame_np = np.frombuffer(frame_data, np.float32)
+            frame_np = np.frombuffer(frame_data, np.float32)
             clients[websocket].add_frames(frame_np)
             
         except Exception as e:
@@ -51,6 +50,15 @@ class ServeClient:
         self.data = b""
         self.frames = b""
         self.transcriber = WhisperModel("small.en", compute_type="float16", local_files_only=False)
+        
+        # voice activity detection model
+        self.vad_model, _ = torch.hub.load(repo_or_dir='snakers4/silero-vad',
+                                           model='silero_vad',
+                                           force_reload=True,
+                                           onnx=True
+                                           )
+        self.vad_threshold = 0.4
+        
         self.timestamp_offset = 0.0
         self.frames_np = None
         self.frames_offset = 0.0
@@ -100,6 +108,15 @@ class ServeClient:
         return wrapped
     
     def add_frames(self, frame_np):
+        try:
+            speech_prob = self.vad_model(torch.from_numpy(frame_np), self.RATE).item()
+            if speech_prob < self.vad_threshold:
+                return
+            
+        except Exception as e:
+            logging.error(e)
+            return
+        
         if self.frames_np is not None and self.frames_np.shape[0] > 45*self.RATE:
             self.frames_offset += 45.0
             self.frames_np = self.frames_np[int(30*self.RATE):]