ALVINN/ALVINN_f7
1
0
Fork 0
Code Issues 8 Pull Requests Actions Packages Projects Releases 7 Wiki Activity

Optimize detection steps (#129)

Browse Source 
Closes: #128
Reviewed-on: #129
This commit is contained in:
Justin Georgi
2024-03-09 21:06:50 -07:00
parent 1cc76c5950
commit 76b2fe5b0c
1 changed files with 13 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
src/pages/detection-mixin.js
View File

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