add voice activity detection

2023-05-30 21:07:36 +05:30
parent 69eb8ae421
commit e37e7f5522
7 changed files with 75 additions and 8 deletions
@@ -170,7 +170,8 @@ async function stopCapture() {

  if (optionTabId) {
    res = await sendMessageToTab(currentTabId, {
-      type: "STOP"
+      type: "STOP",
+      data: { currentTabId: currentTabId },
    });
    await removeChromeTab(optionTabId);
  }
@@ -188,3 +189,15 @@ chrome.runtime.onMessage.addListener((message) => {
    stopCapture();
  }
 });
+
+
+/**
+ * Listens for if the tab is reloaded.
+ * @param {Object} message - The message received from the runtime.
+ */
+chrome.tabs.onUpdated.addListener(async (tabId, changeInfo, tab) => {
+  if (changeInfo.status === 'complete') {
+    await executeScriptInTab(tabId, "content.js");
+    await delayExecution(500);
+  }
+});
@@ -4,7 +4,9 @@
    "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"
@@ -4,15 +4,14 @@
 <head>
    <meta charset="UTF-8">
    <title>Audio Transcription Options</title>
-    <script src="options.js"></script>
    <meta http-equiv="X-UA-Compatible" content="IE=edge">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
-    <link rel="stylesheet" href="options.css" type="text/css">
+    <link rel="stylesheet" href="style.css" type="text/css">
 </head>

 <body>
-    <div class="header"><img src="./collabora.png"/> <h1>Audio Transcription</h1></div>
-    <div class="button" id="stop">Stop Capture</div>
+    <script src="ort.min.js"></script>
+    <script src="options.js"></script>
 </body>

 </html>
@@ -38,12 +38,43 @@ function sendMessageToTab(tabId, data) {
 */
 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");
@@ -51,7 +82,8 @@ async function startRecord(option) {

    socket.onmessage = async (event) => {
      // console.log(event.data);
-      await sendMessageToTab(option.currentTabId, {
+      res = await sendMessageToTab(option.currentTabId, {
+        type: "transcript",
        data: event.data,
      });
    };
@@ -67,7 +99,22 @@ async function startRecord(option) {
      const inputData = event.inputBuffer.getChannelData(0);

      audioDataCache.push(inputData);
-      socket.send(inputData);
+
+      // voice activity detection inference
+      const audioBuffer = new ort.Tensor('float32', inputData, [1, inputData.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(inputData);
+        }
+        else
+          console.log("no speech found: " + speech_prob)
+      }
    };

    // Prevent page mute