Deploy TensorFlow.js in a Chrome extension

In this tutorial, you'll build and install a Chrome extension that lets you right-click on an image in a web page and perform multiclass object detection on the image. The extension applies a MobileNetV2 classifier to the image and then labels the image with the predicted class.

The example code is available on GitHub.

Prerequisites

To complete this tutorial, you need the following installed in your development environment:

Build the extension

Get the source code and build the extension:

  1. Clone or download the tfjs-examples repository.
  2. Change into the chrome-extension directory: cd tfjs-examples/chrome-extension.
  3. Install dependencies: yarn.
  4. Run the build script: yarn build.

After running the build script, you should see the following new files:

  • dist/src/content.js
  • dist/src/service_worker.js
  • dist/src/service_worker.js.map

Install the extension

Install the extension in Chrome:

  1. In the Chrome browser, navigate to chrome://extensions.
  2. Turn on Developer mode using the toggle on the right side of the browser.
  3. Select Load unpacked and select the tfjs-examples/chrome-extension/dist directory. This directory contains the manifest.json file and the src/*.js files packaged by the build.

You should see a new card for TF.js mobilenet in a Chrome extension.

Use the extension

With the extension installed, you can classify images in the browser:

  1. Navigate to a site with images on it. For example, navigate to google.com, search for "tigers", and select Images in the results page. You should see a page of tiger images.
  2. Right-click on an image and select Classify image with TensorFlow.js. There's a warm up period, so the first time you run the app, inference will be slower. (In your own applications, you can prime the model by feeding it dummy data.)

The extension runs the model on the image and then superimposes text indicating the prediction.

Remove the extension

When you're done experimenting with the extension, you can remove it:

  1. In Chrome, navigate to chrome://extensions.
  2. In the TF.js mobilenet in a Chrome extension card, select Remove and confirm that you want to remove the extension.

How the extension works

This section describes how the extension works, at a high level.

The manifest file, manifest.json, specifies a service worker that Chrome will run in the background:

"background": {
   "service_worker": "src/service_worker.js"
},

The service worker script, service_worker.js, imports the TensorFlow.js package and the mobilenet model.

import * as mobilenet from '@tensorflow-models/mobilenet';
import * as tf from '@tensorflow/tfjs';

The build script in package.json uses a bundler, Parcel, to bundle everything together so no external scripts are loaded at runtime.

"build": "parcel build src/service_worker.js --dist-dir dist/src/ && npm run copy",

This is to comply with Chrome Manifest V3, which prohibits remotely hosted code. Note that a service worker can still load external resources, such as TensorFlow.js models.

The service worker script creates a context menu item that operates on images. Then the script listens for clicks.

/**
 * Adds a right-click menu option to trigger classifying the image.
 * The menu option should only appear when right-clicking an image.
 */
chrome.runtime.onInstalled.addListener(() => {
  chrome.contextMenus.create({
    id: 'contextMenu0',
    title: 'Classify image with TensorFlow.js ',
    contexts: ['image'],
  });
});

chrome.contextMenus.onClicked.addListener(clickMenuCallback);

When the user selects the menu item, a callback sends a message containing the current tab ID and the URL of the right-clicked image. (Note that in a service worker, DOM objects are not available.)

function clickMenuCallback(info, tab) {
  const message = { action: 'IMAGE_CLICKED', url: info.srcUrl };
  chrome.tabs.sendMessage(tab.id, message, (resp) => {
    if (!resp.rawImageData) {
      console.error(
        'Failed to get image data. ' +
        'The image might be too small or failed to load. ' +
        'See console logs for errors.');
      return;
    }
    const imageData = new ImageData(
      Uint8ClampedArray.from(resp.rawImageData), resp.width, resp.height);
    imageClassifier.analyzeImage(imageData, info.srcUrl, tab.id);
  });
}

The content script, content.js, listens for messages and handles the IMAGE_CLICKED action. The script receives the image URL, loads the image, renders the image on an OffscreenCanvas, gets the image data from the canvas, and sends the data back to the service worker.

After the service worker receives the image data, it runs the mobilenet model with the data and gets the prediction results. In the clickMenuCallback function above, imageClassifier is an instance of the ImageClassifier class, which loads the model and gets predictions. The service worker script then sends results back to the content script for display. After the content script receives the results, it overlays the results on top of the original image.

The service worker thread goes idle when no activity occurs for approximately 30 seconds. For more information on managing service worker events, see the Chrome documentation.

What's next

This tutorial showed how to deploy a Chrome extension that uses TensorFlow.js and a pre-trained MobileNet model to classify images. To learn more about pre-trained models for TensorFlow.js, see the pre-trained model repository.