Class ImageClassifier

chrome-extension/src/service_worker.js:69–123

* Async loads a mobilenet on construction. Subsequently handles * requests to classify images through the .analyzeImage API. * Successful requests will post a chrome message with * 'IMAGE_CLICK_PROCESSED' action, which the content.js can * hear and use to manipulate the DOM.

 * hear and use to manipulate the DOM.
 */
class ImageClassifier {
  constructor() {
    this.loadModel();
  }

  /**
   * Loads mobilenet from URL and keeps a reference to it in the object.
   */
  async loadModel() {
    console.log('Loading model...');
    const startTime = performance.now();
    try {
      this.model = await mobilenet.load({ version: 2, alpha: 1.00 });
      // Warms up the model by causing intermediate tensor values
      // to be built and pushed to GPU.
      tf.tidy(() => {
        this.model.classify(tf.zeros([1, IMAGE_SIZE, IMAGE_SIZE, 3]));
      });
      const totalTime = Math.floor(performance.now() - startTime);
      console.log(`Model loaded and initialized in ${totalTime} ms...`);
    } catch (e) {
      console.error('Unable to load model', e);
    }
  }

  /**
   * Triggers the model to make a prediction on the image referenced by the
   * image data. After a successful prediction a IMAGE_CLICK_PROCESSED message
   * when complete, for the content.js script to hear and update the DOM with
   * the results of the prediction.
   *
   * @param {ImageData} imageData ImageData of the image to analyze.
   * @param {string} url url of image to analyze.
   * @param {number} tabId which tab the request comes from.
   */
  async analyzeImage(imageData, url, tabId) {
    if (!tabId) {
      console.error('No tab.  No prediction.');
      return;
    }
    if (!this.model) {
      console.log('Waiting for model to load...');
      setTimeout(
        () => { this.analyzeImage(imageData, url, tabId) }, FIVE_SECONDS_IN_MS);
      return;
    }
    console.log('Predicting...');
    const startTime = performance.now();
    const predictions = await this.model.classify(imageData, TOPK_PREDICTIONS);
    const totalTime = performance.now() - startTime;
    console.log(`Done in ${totalTime.toFixed(1)} ms `);
    const message = { action: 'IMAGE_CLICK_PROCESSED', url, predictions };
    chrome.tabs.sendMessage(tabId, message);
  }
}

const imageClassifier = new ImageClassifier();