ALVINN_f7/src/assets/detect-worker.js

import * as tf  from '@tensorflow/tfjs'

let model = null

onmessage = function (e) {
  switch (e.data.call) {
    case 'loadModel':
      loadModel('.' + e.data.weights,e.data.preload).then(() => {
        postMessage({success: 'model'})
      }).catch((err) => {
        postMessage({error: true, message: err.message})
      })
      break
    case 'localDetect':
      localDetect(e.data.image).then((dets) => {
        postMessage({success: 'detection', detections: dets})
      }).catch((err) => {
        //throw (err)
        postMessage({error: true, message: err.message})
      })
      e.data.image.close()
      break
    case 'videoFrame':
      videoFrame(e.data.image).then((frameDet) =>{
        postMessage({succes: 'frame', coords: frameDet.cds, modelWidth: frameDet.mW, modelHeight: frameDet.mH})
      }).catch((err) => {
        postMessage({error: true, message: err.message})
      })
      e.data.image.close()
      break
    default:
      console.log('Worker message incoming:')
      console.log(e)
      postMessage({result1: 'First result', result2: 'Second result'})
      break
  }
}

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}
}