Cleanup detection mixin

Signed-off-by: Justin Georgi <justin.georgi@gmail.com>
This commit is contained in:
2024-03-08 21:26:40 -07:00
parent 251751c7bf
commit bcae1a9474

View File

@@ -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,19 +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]).arraySync()[0]
let rawBoxes = []
let rawScores = []
const filteredRes = rawRes.filter( r => r.slice(4).some( s => s > .05))
for (var i = 0; i < filteredRes.length; i++) {
var getScores = filteredRes[i].slice(4)
var getBox = filteredRes[i].slice(0,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,
@@ -48,16 +45,18 @@ export default {
}
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))
}
@@ -81,6 +80,8 @@ export default {
tf.dispose(res)
tf.dispose(tBoxes)
tf.dispose(tScores)
tf.dispose(input)
console.timeEnd('post-process')
return output