Chrome browser plugin deployment example (mobilenet)

deployment-examples/README.md

@@ -0,0 +1,14 @@
+# TensorFlow.js Deployment Examples
+
+This directory contains simple examples meant to illustrate how to deploy
+[TensorFlow.js](http://js.tensorflow.org) models in different deployment
+scenarios, such as to the web, to a node.js server, or to a react-native
+app.
+
+Each example directory is standalone so the directory can be copied
+to another project.
+
+
+TODO(bileschi) :
+*  Add instructions how to run & set up
+*  Add link to live example (chrome web store?)

deployment-examples/browser-plugin/README.md

@@ -0,0 +1,5 @@
+# TensorFlow.js Deployment Example : Browser Plugin
+
+This example creates a Chrome plugin, enabling users to click on images within a web page, and perform multi-class object detection on them.  The plugin will report to the console log
+the top predicted classes of the image, as reported by the mobilenet model.
+

deployment-examples/browser-plugin/package.json

@@ -0,0 +1,22 @@
+{
+  "name": "tfjs-basic-chrome-extension",
+  "version": "0.1.0",
+  "description": "Use tfjs model.predict in a chrome extension",
+  "scripts": {
+    "copy": "cp src/content.js dist/src/ && cp src/imagenet_classes.js dist/src/",
+    "build": "parcel build src/background.js -d dist/src/ -o background --no-minify && npm run copy",
+    "watch": "npm run copy && parcel watch src/background.js --hmr-hostname localhost -d dist/src/ -o background"
+  },
+  "license": "Apache 2.0",
+  "devDependencies": {
+    "babel-core": "^6.26.3",
+    "babel-plugin-transform-runtime": "^6.23.0",
+    "babel-polyfill": "^6.26.0",
+    "babel-preset-env": "^1.6.1",
+    "clang-format": "^1.2.3",
+    "parcel-bundler": "^1.7.1"
+  },
+  "dependencies": {
+    "@tensorflow/tfjs": "^1.1.0"
+  }
+}

deployment-examples/browser-plugin/src/background.js

@@ -0,0 +1,153 @@
+import 'babel-polyfill';
+import * as tf from '@tensorflow/tfjs';
+import { IMAGENET_CLASSES } from './imagenet_classes';
+
+// Where to load the model from.
+const MOBILENET_MODEL_PATH = 'https://storage.googleapis.com/tfjs-models/tfjs/mobilenet_v1_0.25_224/model.json';
+const IMAGE_SIZE = 224;
+const TOPK_PREDICTIONS = 2;
+
+function clickMenuFn(info, tab) {
+  imageClassifier.analyzeImage(info.srcUrl, tab.id);
+}
+
+// Add a right-click menu option to trigger classifying the image.
+// Menu option should only appear when right-clicking an image.
+chrome.contextMenus.create({
+  title: "Classify image with TensorFlow.js ", 
+  contexts:["image"], 
+  onclick: clickMenuFn
+});
+
+// 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.
+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 tf.loadLayersModel(MOBILENET_MODEL_PATH);
+      this.model.predict(tf.zeros([1, IMAGE_SIZE, IMAGE_SIZE, 3])).dispose();
+
+      const totalTime = Math.floor(performance.now() - startTime);
+      console.log(`Model loaded and initialized in ${totalTime}ms...`);
+    } catch {
+      console.error(`Unable to load model from URL: ${MOBILENET_MODEL_PATH}`);
+    }
+  }
+
+  async analyzeImage(url, tabId) {
+    if (!tabId) {
+      console.log('No tab.  No prediction.');
+      return;
+    }
+    if (!this.model) {
+      console.log('Waiting for model to load...');
+      setTimeout(() => { this.analyzeImage(url) }, 5000);
+      return;
+    }
+    const img = await this.loadImage(url);
+    if (!img) {
+      console.log('could not load image');
+      return;
+    }
+    const predictions = await this.predict(img);
+    if (!predictions) {
+      console.log('failed to create predictions.');
+      return;
+    }
+    const message = {
+      action: 'IMAGE_CLICK_PROCESSED',
+      url,
+      predictions
+    };
+    chrome.tabs.sendMessage(tabId, message);
+  }
+
+  async loadImage(src) {
+    return new Promise(resolve => {
+      var img = document.createElement('img');
+      img.crossOrigin = "anonymous";
+      img.onerror = function(e) {
+        resolve(null);
+      };
+      img.onload = function(e) {
+        if ((img.height && img.height > 128) || (img.width && img.width > 128)) {
+          // Set image size for tf!
+          img.width = IMAGE_SIZE;
+          img.height = IMAGE_SIZE;
+          resolve(img);
+        }
+        // Let's skip all tiny images
+        resolve(null);
+      }
+      img.src = src;
+    });
+  }
+
+  async getTopKClasses(logits, topK) {
+    const values = await logits.data();
+    const valuesAndIndices = [];
+    for (let i = 0; i < values.length; i++) {
+      valuesAndIndices.push({value: values[i], index: i});
+    }
+    valuesAndIndices.sort((a, b) => {
+      return b.value - a.value;
+    });
+    const topkValues = new Float32Array(topK);
+    const topkIndices = new Int32Array(topK);
+    for (let i = 0; i < topK; i++) {
+      topkValues[i] = valuesAndIndices[i].value;
+      topkIndices[i] = valuesAndIndices[i].index;
+    }
+
+    const topClassesAndProbs = [];
+    for (let i = 0; i < topkIndices.length; i++) {
+      topClassesAndProbs.push({
+        className: IMAGENET_CLASSES[topkIndices[i]],
+        probability: topkValues[i]
+      })
+    }
+    return topClassesAndProbs;
+  }
+
+
+  async predict(imgElement) {
+    console.log('Predicting...');
+    // The first start time includes the time it takes to extract the image
+    // from the HTML and preprocess it, in additon to the predict() call.
+    const startTime1 = performance.now();
+    // The second start time excludes the extraction and preprocessing and
+    // includes only the predict() call.
+    let startTime2;
+    const logits = tf.tidy(() => {
+      const img = tf.browser.fromPixels(imgElement).toFloat();
+      const offset = tf.scalar(127.5);
+      const normalized = img.sub(offset).div(offset);
+      const batched = normalized.reshape([1, IMAGE_SIZE, IMAGE_SIZE, 3]);
+      startTime2 = performance.now();
+      return this.model.predict(batched);
+    });
+
+    // Convert logits to probabilities and class names.
+    const classes = await this.getTopKClasses(logits, TOPK_PREDICTIONS);
+    const totalTime1 = performance.now() - startTime1;
+    const totalTime2 = performance.now() - startTime2;
+    console.log(`Done in ${Math.floor(totalTime1)} ms ` +
+      `(not including preprocessing: ${Math.floor(totalTime2)} ms)`);
+    return classes;
+  }
+
+
+}
+
+var imageClassifier = new ImageClassifier();

deployment-examples/browser-plugin/src/content.js

@@ -0,0 +1,125 @@
+/**
+ * @license
+ * Copyright 2019 Google LLC. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+
+// class name for all text nodes added by this script.
+const TEXT_DIV_CLASSNAME = 'tfjs_mobilenet_extension_text';
+
+// Produces a short text string summarizing the prediction
+// Input prediction should be a list of {className: string, prediction: float}
+// objects.
+function textContentFromPrediction(predictions) {
+  if (!predictions || predictions.length < 1) {
+    return "No prediction 🙁";
+  }
+  // Confident.
+  if (predictions[0].probability >= 0.5) {
+    return "😄 " + predictions[0].className + "!";
+  }
+  // Not Confident.
+  if (predictions[0].probability >= 0.1 && predictions[0].probability < 0.5) {
+    return (predictions[0].className + "?... \n Maybe " + 
+      predictions[1].className + "?");
+  }
+  // Very not confident.
+  if (predictions[0].probability < 0.1) {
+    return "😕  " + predictions[0].className + 
+    "????... \n Maybe " + predictions[1].className + "????";
+  }
+}
+
+// Returns a list of all DOM image elements pointing to the provided url.
+function getImageElementsWithSrcUrl(srcUrl) {
+  const imgElArr = Array.from(document.getElementsByTagName('img'));
+  const filtImgElArr = imgElArr.filter(x => x.src === srcUrl);  
+  return filtImgElArr;
+}
+
+// Removes all text predictions from the DOM.
+//
+// Note:
+// This does not undo the containerization.  A cleaner implementation
+// would move the image node back out of the container div.
+function removeTextElements() {
+  const textDivs = document.getElementsByClassName(TEXT_DIV_CLASSNAME);
+  for (const div of textDivs) {
+    div.parentNode.removeChild(div);
+  }
+}
+
+// Moves the provided imgNode into a container div, and adds a text div as a 
+// peer.  Styles the container div and text div to place the text
+// on top of the image.
+function addTextElementToImageNode(imgNode, textContent) {
+  const originalParent = imgNode.parentElement;
+  const container = document.createElement('div');
+  container.style.position = 'relative';
+  container.style.textAlign = 'center';
+  container.style.colore = 'white';
+  const text = document.createElement('div');
+  text.className = 'tfjs_mobilenet_extension_text';
+  text.style.position = 'absolute';
+  text.style.top = '50%';
+  text.style.left = '50%';
+  text.style.transform = 'translate(-50%, -50%)';
+  text.style.fontSize = '34px';
+  text.style.fontFamily = 'Google Sans,sans-serif';
+  text.style.fontWeight = '700';
+  text.style.color = 'white';
+  text.style.lineHeight = '1em';
+  text.style["-webkit-text-fill-color"] = "white";
+  text.style["-webkit-text-stroke-width"] = "1px";
+  text.style["-webkit-text-stroke-color"] =  "black";
+  // Add the containerNode as a peer to the image, right next to the image.
+  originalParent.insertBefore(container, imgNode);
+  // Move the imageNode to inside the containerNode;
+  container.appendChild(imgNode);
+  // Add the text node right after the image node;
+  container.appendChild(text);
+  text.textContent = textContent; 
+}
+
+// Add a listener to hear from the content.js page when the image is through
+// processing.  The message should contin an action, a url, and predictions (the
+// output of the classifier)
+//
+// message: {action, url, predictions}
+chrome.runtime.onMessage.addListener((message, sender, sendResponse) => {
+  if (message &&
+      message.action === 'IMAGE_CLICK_PROCESSED' &&
+      message.url &&
+      message.predictions) {
+    // Get the list of images with this srcUrl.
+    const imgElements = getImageElementsWithSrcUrl(message.url);
+    if (imgElements.length === 0) {
+      console.log(`Could not find an image with url ${message.url}`)
+    }
+    for (const imgNode of imgElements) {
+      const textContent = textContentFromPrediction(message.predictions); 
+      addTextElementToImageNode(imgNode, textContent);
+    }
+  }
+});
+
+// Set up a listener to remove all annotations if the user clicks
+// the left mouse button.  Otherwise, they can easily cloud up the
+// window.
+window.addEventListener('click', clickHandler, false);
+function clickHandler(mouseEvent) {
+  if (mouseEvent.button == 0) {
+    removeTextElements();
+  }
+}