Add Tensorflow based local detection (#95)

Closes: #12
Reviewed-on: Georgi_Lab/ALVINN_f7#95

package.json

@@ -23,6 +23,7 @@
     "last 5 Firefox versions"
   ],
   "dependencies": {
+    "@tensorflow/tfjs": "^4.17.0",
     "dom7": "^4.0.6",
     "framework7": "^8.3.0",
     "framework7-icons": "^5.0.5",

src/models/thorax_tfwm/classes.json

@@ -0,0 +1,12 @@
+[
+  "Right lung",
+  "Diaphragm",
+  "Heart",
+  "Caudal vena cava",
+  "Cranial vena cava",
+  "Phrenic nerve",
+  "Trachea",
+  "Vagus nerve",
+  "Left Lung",
+  "Aorta"
+]

src/pages/detect.vue

@@ -21,7 +21,7 @@
           :style="chipGradient(result.confidence)"
         />
         <span v-if="numResults == 0 && !detecting" style="height: var(--f7-chip-height); font-size: calc(var(--f7-chip-height) - 4px); font-weight: bolder; margin: 2px;">No results.</span>
-        <f7-preloader v-if="detecting" size="32" style="color: var(--avn-theme-color);" />
+        <f7-preloader v-if="detecting || modelLoading" size="32" style="color: var(--avn-theme-color);" />
       </div>
       <div v-if="showDetectSettings" class="detect-inputs" style="grid-area: detect-settings;">
         <f7-range class="level-slide-horz" :min="0" :max="100" :step="1" @range:change="onLevelChange" v-model:value="detectorLevel" type="range" style="flex: 1 1 100%"/>
@@ -286,9 +286,12 @@
   import SvgIcon from '../components/svg-icon.vue'
 
   import submitMixin from './submit-mixin'
+  import detectMixin from './local-detect'
+
+  import thoraxClasses from '../models/thorax_tfwm/classes.json'
 
   export default {
-    mixins: [submitMixin],
+    mixins: [submitMixin, detectMixin],
     props: {
       f7route: Object,
     },
@@ -302,6 +305,7 @@
         resultData: {},
         selectedChip: -1,
         activeRegion: 4,
+        classesList: [],
         imageLoaded: false,
         imageView: null,
         imageLoadMode: "environment",
@@ -314,7 +318,9 @@
         serverSettings: {},
         isCordova: !!window.cordova,
         uploadUid: null,
-        uploadDirty: false
+        uploadDirty: false,
+        modelLocation: '',
+        modelLoading: false
       }
     },
     setup() {
@@ -325,6 +331,8 @@
         case 'thorax':
           this.activeRegion = 0
           this.detectorName = 'thorax'
+          this.classesList = thoraxClasses
+          this.modelLocation = '../models/thorax_tfwm/model.json'
           break;
         case 'abdomen':
           this.activeRegion = 1
@@ -364,6 +372,12 @@
         }
         
         xhr.send()
+      } else {
+        self.modelLoading = true
+        self.detectorLabels = self.classesList.map( l => { return {'name': l, 'detect': true} } )
+        self.loadModel(self.modelLocation).then(() => {
+          self.modelLoading = false
+        })
       }
       window.onresize = (e) => { this.selectChip('redraw') }
     },
@@ -431,8 +445,11 @@
 
           xhr.send(JSON.stringify(doodsData))
         } else {
-          //TODO
-          f7.dialog.alert('Using built-in model')
+          this.localDetect(this.imageView).then(dets => {
+            self.detecting = false
+            self.resultData = dets
+            self.uploadDirty = true
+          })
         }
       },
       remoteTimeout () {
@@ -526,7 +543,15 @@
         }).then( () => {
           const [imCanvas, _]  = this.resetView()
           imCanvas.style['background-image'] = `url(${this.imageView.src})`
-          this.setData()
+          /******
+           * setTimeout is not a good solution,
+           * but it's the only way I can find to
+           * not cut off drawing of of the progress
+           * spinner
+           ******/
+          setTimeout(() => {
+            this.setData()
+          }, 250)
         }).catch((e) => {
           console.log(e.message)
           f7.dialog.alert(`Error loading image: ${e.message}`)

src/pages/local-detect.js

@@ -0,0 +1,48 @@
+import * as tf  from '@tensorflow/tfjs'
+
+var model = null
+
+export default {
+  methods: {
+    async loadModel(weights) {
+      model = await tf.loadGraphModel(weights).then(graphModel => {
+        return graphModel
+      })
+    },
+    async localDetect(imageData) {
+      const [modelWidth, modelHeight] = model.inputs[0].shape.slice(1, 3);
+
+      const input = tf.tidy(() => {
+        return tf.image.resizeBilinear(tf.browser.fromPixels(imageData), [modelWidth, modelHeight]).div(255.0).expandDims(0)
+      })
+      var results = model.executeAsync(input).then(res => {
+        const [boxes, scores, classes, valid_detections] = res;
+        const boxes_data = boxes.dataSync();
+        const scores_data = scores.dataSync();
+        const classes_data = classes.dataSync();
+        const valid_detections_data = valid_detections.dataSync()[0];
+
+        tf.dispose(res)
+
+        var output = {
+          detections: []
+        }
+        for (var i =0; i < valid_detections_data; i++) {
+          var [dLeft, dTop, dRight, dBottom] = boxes_data.slice(i * 4, (i + 1) * 4);
+          output.detections.push({
+            "top": dTop,
+            "left": dLeft,
+            "bottom": dBottom,
+            "right": dRight,
+            "label": this.detectorLabels[classes_data[i]].name,
+            "confidence": scores_data[i] * 100
+          })
+        }
+
+        return output
+      })
+
+      return results
+    }
+  }
+}