Update ImageClassifier and FeatureExtractor to return [{label, confidence}] (#292)

* change imageClassifier.predict() to classify(), keep predict until later version

* update imageclassifier.classify results to [{label: .., confidence: ..}]

* update featureExtractor regressor return {value: 0-1}

* update featureextractor.classify to return [{label, confidence}]

Author: yining1023

src/FeatureExtractor/Mobilenet.js

@@ -247,18 +247,21 @@ class Mobilenet {
     }
     await tf.nextFrame();
     this.isPredicting = true;
-    const predictedClass = tf.tidy(() => {
+    const predictedClasses = tf.tidy(() => {
       const imageResize = (imgToPredict === this.video) ? null : [IMAGE_SIZE, IMAGE_SIZE];
       const processedImg = imgToTensor(imgToPredict, imageResize);
       const activation = this.mobilenetFeatures.predict(processedImg);
       const predictions = this.customModel.predict(activation);
-      return predictions.as1D().argMax();
+      return Array.from(predictions.as1D().dataSync());
     });
-    let classId = (await predictedClass.data())[0];
-    if (this.mapStringToIndex.length > 0) {
-      classId = this.mapStringToIndex[classId];
-    }
-    return classId;
+    const results = await predictedClasses.map((confidence, index) => {
+      const label = (this.mapStringToIndex.length > 0 && this.mapStringToIndex[index]) ? this.mapStringToIndex[index] : index;
+      return {
+        label,
+        confidence,
+      };
+    }).sort((a, b) => b.confidence - a.confidence);
+    return results;
   }
 
   /* eslint max-len: ["error", { "code": 180 }] */
@@ -295,7 +298,7 @@ class Mobilenet {
     });
     const prediction = await predictedClass.data();
     predictedClass.dispose();
-    return prediction[0];
+    return { value: prediction[0] };
   }
 
   async load(filesOrPath = null, callback) {

src/ImageClassifier/index.js

@@ -52,7 +52,7 @@ class ImageClassifier {
     return this;
   }
 
-  async predictInternal(imgToPredict, numberOfClasses) {
+  async classifyInternal(imgToPredict, numberOfClasses) {
     // Wait for the model to be ready
     await this.ready;
     await tf.nextFrame();
@@ -62,10 +62,12 @@ class ImageClassifier {
         this.video.onloadeddata = () => resolve();
       });
     }
-    return this.model.classify(imgToPredict, numberOfClasses);
+    return this.model
+      .classify(imgToPredict, numberOfClasses)
+      .then(classes => classes.map(c => ({ label: c.className, confidence: c.probability })));
   }
 
-  async predict(inputNumOrCallback, numOrCallback = null, cb) {
+  async classify(inputNumOrCallback, numOrCallback = null, cb) {
     let imgToPredict = this.video;
     let numberOfClasses = this.topk;
     let callback;
@@ -102,7 +104,11 @@ class ImageClassifier {
       callback = cb;
     }
 
-    return callCallback(this.predictInternal(imgToPredict, numberOfClasses), callback);
+    return callCallback(this.classifyInternal(imgToPredict, numberOfClasses), callback);
+  }
+
+  async predict(inputNumOrCallback, numOrCallback, cb) {
+    return this.classify(inputNumOrCallback, numOrCallback || null, cb);
   }
 }
 

src/ImageClassifier/index_test.js

@@ -48,35 +48,35 @@ describe('imageClassifier', () => {
     expect(classifier.ready).toBeTruthy();
   });
 
-  describe('predict', () => {
+  describe('classify', () => {
     it('Should classify an image of a Robin', async () => {
       const img = await getImage();
-      await classifier.predict(img)
-        .then(results => expect(results[0].className).toBe('robin, American robin, Turdus migratorius'));
+      await classifier.classify(img)
+        .then(results => expect(results[0].label).toBe('robin, American robin, Turdus migratorius'));
     });
 
     it('Should support p5 elements with an image on .elt', async () => {
       const img = await getImage();
-      await classifier.predict({ elt: img })
-        .then(results => expect(results[0].className).toBe('robin, American robin, Turdus migratorius'));
+      await classifier.classify({ elt: img })
+        .then(results => expect(results[0].label).toBe('robin, American robin, Turdus migratorius'));
     });
 
     it('Should support HTMLCanvasElement', async () => {
       const canvas = await getCanvas();
-      await classifier.predict(canvas)
-        .then(results => expect(results[0].className).toBe('robin, American robin, Turdus migratorius'));
+      await classifier.classify(canvas)
+        .then(results => expect(results[0].label).toBe('robin, American robin, Turdus migratorius'));
     });
 
     it('Should support p5 elements with canvas on .canvas', async () => {
       const canvas = await getCanvas();
-      await classifier.predict({ canvas })
-        .then(results => expect(results[0].className).toBe('robin, American robin, Turdus migratorius'));
+      await classifier.classify({ canvas })
+        .then(results => expect(results[0].label).toBe('robin, American robin, Turdus migratorius'));
     });
 
     it('Should support p5 elements with canvas on .elt', async () => {
       const canvas = await getCanvas();
-      await classifier.predict({ elt: canvas })
-        .then(results => expect(results[0].className).toBe('robin, American robin, Turdus migratorius'));
+      await classifier.classify({ elt: canvas })
+        .then(results => expect(results[0].label).toBe('robin, American robin, Turdus migratorius'));
     });
   });
 });