Add detection worker (#187)

Closes: #186

This PR shifts much of the tensorflow function to a shared worker for multithreading performance.

Reviewed-on: #187

src/assets/detect-worker.js

@@ -0,0 +1,182 @@
+import * as tf  from '@tensorflow/tfjs'
+import { f7 } from 'framework7-vue'
+
+let model = null
+
+self.onconnect = (e) => {
+  const port = e.ports[0];
+
+  port.onmessage = function (e) {
+    switch (e.data.call) {
+      case 'loadModel':
+        loadModel('.' + e.data.weights,e.data.preload).then(() => {
+          port.postMessage({success: 'model'})
+        }).catch((err) => {
+          port.postMessage({error: true, message: err.message})
+        })
+        break
+      case 'localDetect':
+        localDetect(e.data.image).then((dets) => {
+          port.postMessage({success: 'detection', detections: dets})
+        }).catch((err) => {
+          port.postMessage({error: true, message: err.message})
+        })
+        e.data.image.close()
+        break
+      case 'videoFrame':
+        videoFrame(e.data.image).then((frameDet) =>{
+          port.postMessage({succes: 'frame', coords: frameDet.cds, modelWidth: frameDet.mW, modelHeight: frameDet.mH})
+        }).catch((err) => {
+          port.postMessage({error: true, message: err.message})
+        })
+        e.data.image.close()
+        break
+      default:
+        console.log('Worker message incoming:')
+        console.log(e)
+        port.postMessage({result1: 'First result', result2: 'Second result'})
+        break
+    }
+  }
+
+  port.start()
+}
+
+async function loadModel(weights, preload) {
+  if (model && model.modelURL == weights) {
+    return model
+  } else if (model) {
+    tf.dispose(model)
+  }
+  model = await tf.loadGraphModel(weights)
+  const [modelWidth, modelHeight] = model.inputs[0].shape.slice(1, 3)
+  /*****************
+   * If preloading then run model
+   * once on fake data to preload
+   * weights for a faster response
+   *****************/
+  if (preload) {
+    const dummyT = tf.ones([1,modelWidth,modelHeight,3])
+    model.predict(dummyT)
+  }
+  return model
+}
+
+async function localDetect(imageData) {
+  console.time('pre-process')
+  const [modelWidth, modelHeight] = model.inputs[0].shape.slice(1, 3)
+  let gTense = null
+  const input = tf.tidy(() => {
+    gTense = tf.image.rgbToGrayscale(tf.image.resizeBilinear(tf.browser.fromPixels(imageData), [modelWidth, modelHeight])).div(255.0).expandDims(0)
+    return tf.concat([gTense,gTense,gTense],3)
+  })
+  tf.dispose(gTense)
+  console.timeEnd('pre-process')
+
+  console.time('run prediction')
+  const res = model.predict(input)
+  const tRes = tf.transpose(res,[0,2,1])
+  const rawRes = tRes.arraySync()[0]
+  console.timeEnd('run prediction')
+
+  console.time('post-process')
+  const outputSize = res.shape[1]
+  let rawBoxes = []
+  let rawScores = []
+
+  for (var i = 0; i < rawRes.length; i++) {
+    var getScores = rawRes[i].slice(4)
+    if (getScores.every( s => s < .05)) { continue }
+    var getBox = rawRes[i].slice(0,4)
+    var boxCalc = [
+      (getBox[0] - (getBox[2] / 2)) / modelWidth,
+      (getBox[1] - (getBox[3] / 2)) / modelHeight,
+      (getBox[0] + (getBox[2] / 2)) / modelWidth,
+      (getBox[1] + (getBox[3] / 2)) / modelHeight,
+    ]
+    rawBoxes.push(boxCalc)
+    rawScores.push(getScores)
+  }
+
+  if (rawBoxes.length > 0) {
+    const tBoxes = tf.tensor2d(rawBoxes)
+    let tScores = null
+    let resBoxes = null
+    let validBoxes = []
+    let structureScores = null
+    let boxes_data = []
+    let scores_data = []
+    let classes_data = []
+    for (var c = 0; c < outputSize - 4; c++) {
+      structureScores = rawScores.map(x => x[c])
+      tScores = tf.tensor1d(structureScores)
+      resBoxes = await tf.image.nonMaxSuppressionAsync(tBoxes,tScores,10,0.5,.05)
+      validBoxes = resBoxes.dataSync()
+      tf.dispose(resBoxes)
+      if (validBoxes) {
+        boxes_data.push(...rawBoxes.filter( (_, idx) => validBoxes.includes(idx)))
+        var outputScores = structureScores.filter( (_, idx) => validBoxes.includes(idx))
+        scores_data.push(...outputScores)
+        classes_data.push(...outputScores.fill(c))
+      }
+    }
+
+    validBoxes = []
+    tf.dispose(tBoxes)
+    tf.dispose(tScores)
+    tf.dispose(tRes)
+    const valid_detections_data = classes_data.length
+    var output = {
+      detections: []
+    }
+    for (var i =0; i < valid_detections_data; i++) {
+      var [dLeft, dTop, dRight, dBottom] = boxes_data[i]
+      output.detections.push({
+        "top": dTop,
+        "left": dLeft,
+        "bottom": dBottom,
+        "right": dRight,
+//        "label": this.detectorLabels[classes_data[i]].name,
+        "label": classes_data[i],
+        "confidence": scores_data[i] * 100
+      })
+    }
+  }
+  tf.dispose(res)
+  tf.dispose(input)
+  console.timeEnd('post-process')
+
+  return output || { detections: [] }
+}
+
+async function videoFrame (vidData) {
+  const [modelWidth, modelHeight] = model.inputs[0].shape.slice(1, 3)
+  console.time('frame-process')
+  let rawCoords = []
+  try {
+    const input = tf.tidy(() => {
+      return tf.image.resizeBilinear(tf.browser.fromPixels(vidData), [modelWidth, modelHeight]).div(255.0).expandDims(0)
+    })
+    const res = model.predict(input)
+    const rawRes = tf.transpose(res,[0,2,1]).arraySync()[0]
+
+    if (rawRes) {
+      for (var i = 0; i < rawRes.length; i++) {
+        let getScores = rawRes[i].slice(4) 
+        if (getScores.some( s => s > .5)) {
+          let foundTarget = rawRes[i].slice(0,2)
+          foundTarget.push(Math.max(...getScores))
+          rawCoords.push(foundTarget)
+        }
+      }
+
+    }
+    tf.dispose(input)
+    tf.dispose(res)
+    tf.dispose(rawRes)
+  } catch (e) {
+    console.log(e)
+  }
+  console.timeEnd('frame-process')
+  return {cds: rawCoords, mW: modelWidth, mH: modelHeight}
+}