Make ImageNet's predict promise correct

ml5js · Feb 17, 2018 · a4fdbd6 · a4fdbd6
1 parent 123ae1c
commit a4fdbd6
Showing 1 changed file with 9 additions and 9 deletions.
diff --git a/src/ImageNet/index.js b/src/ImageNet/index.js
@@ -9,7 +9,7 @@ import { MobileNet } from './../utils/mobileNet';
 class ImageNet {
  constructor(model) {
  this.model = model;
- this.ready = false;
+ this.readyPromise = null;
  this.math = ENV.math;
  if (this.model === 'SqueezeNet') {
  this.net = new SqueezeNet(this.math);
@@ -22,14 +22,11 @@ class ImageNet {
  }
 
  async predict(img, num, callback) {
- if (this.ready) {
- this.getClasses(img, num, callback);
- } else {
- ImageNet.loadModel(this.net).then(() => {
- this.ready = true;
- this.getClasses(img, num, callback);
- });
+ if (!this.readyPromise) {
+ this.readyPromise = ImageNet.loadModel(this.net);
  }
+ await this.readyPromise;
+ return this.getClasses(img, num, callback);
  }
 
  // Private Method
@@ -45,7 +42,10 @@ class ImageNet {
  });
  });
  results.sort((a, b) => b.probability - a.probability);
- callback(results);
+ if (callback) {
+ callback(results);
+ }
+ return results;
  }
 
  static async loadModel(model) {