From 76b2fe5b0c098f16daeccb42c25d23833e4828cd Mon Sep 17 00:00:00 2001
From: Justin Georgi <justin.georgi@gmail.com>
Date: Sat, 9 Mar 2024 21:06:50 -0700
Subject: [PATCH] Optimize detection steps (#129)

Closes: #128
Reviewed-on: https://gitea.azgeorgis.net/ALVINN/ALVINN_f7/pulls/129
---
 src/pages/detection-mixin.js | 23 +++++++++++++----------
 1 file changed, 13 insertions(+), 10 deletions(-)

diff --git a/src/pages/detection-mixin.js b/src/pages/detection-mixin.js
index 10420e0..b361c3a 100644
--- a/src/pages/detection-mixin.js
+++ b/src/pages/detection-mixin.js
@@ -1,13 +1,11 @@
 import * as tf  from '@tensorflow/tfjs'
 import { f7 } from 'framework7-vue'
-import { nextTick } from 'vue'
 
 var model = null
 
 export default {
   methods: {
     async loadModel(weights) {
-      await nextTick()
       model = await tf.loadGraphModel(weights)
       const [modelWidth, modelHeight] = model.inputs[0].shape.slice(1, 3)
       const dummyT = tf.ones([1,modelWidth,modelHeight,3])
@@ -24,17 +22,18 @@ export default {
 
       console.time('run prediction')
       const res = model.predict(input)
+      const rawRes = tf.transpose(res,[0,2,1]).arraySync()[0]
       console.timeEnd('run prediction')
 
       console.time('post-process')
-      const detectAttempts = res.shape[2]
       const outputSize = res.shape[1]
-      const rawRes = tf.transpose(res,[0,2,1]).dataSync()
       let rawBoxes = []
       let rawScores = []
 
-      for (var i = 0; i < detectAttempts; i++) {
-        var getBox = rawRes.slice((i * outputSize),(i * outputSize) + 4)
+      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,
@@ -42,20 +41,22 @@ export default {
           (getBox[1] + (getBox[3] / 2)) / modelHeight,
         ]
         rawBoxes.push(boxCalc)
-        rawScores.push(rawRes.slice((i * outputSize) + 4,(i + 1) * outputSize))
+        rawScores.push(getScores)
       }
       const tBoxes = tf.tensor2d(rawBoxes)
       let tScores = null
+      let structureScores = null
       let boxes_data = []
       let scores_data = []
       let classes_data = []
       for (var c = 0; c < outputSize - 4; c++) {
-        tScores = rawScores.map(x => x[c])
-        var validBoxes = await tf.image.nonMaxSuppressionAsync(tBoxes,tf.tensor1d(tScores),10,0.5,.05)
+        structureScores = rawScores.map(x => x[c])
+        tScores = tf.tensor1d(structureScores)
+        var validBoxes = await tf.image.nonMaxSuppressionAsync(tBoxes,tScores,10,0.5,.05)
         validBoxes = validBoxes.dataSync()
         if (validBoxes) {
           boxes_data.push(...rawBoxes.filter( (_, idx) => validBoxes.includes(idx)))
-          var outputScores = tScores.filter( (_, idx) => validBoxes.includes(idx))
+          var outputScores = structureScores.filter( (_, idx) => validBoxes.includes(idx))
           scores_data.push(...outputScores)
           classes_data.push(...outputScores.fill(c))
         }
@@ -79,6 +80,8 @@ export default {
 
       tf.dispose(res)
       tf.dispose(tBoxes)
+      tf.dispose(tScores)
+      tf.dispose(input)
       console.timeEnd('post-process')
 
       return output