Building a Custom Object Detector for the Web with TensorFlow.js

Object detection is a powerful computer vision technique that allows us to identify and locate objects within images and video. It has a wide range of applications, from monitoring manufacturing defects to analyzing traffic patterns to interactive demos that tag objects in a user‘s environment.

Historically, object detection has required significant expertise in machine learning and access to large compute resources for training. But thanks to the rise of more accessible tools, it‘s now possible for web developers to create and deploy custom object detection models directly in the browser.

In this post, we‘ll walk through how to build a custom object detector using TensorFlow.js, a JavaScript library for machine learning. By the end, you‘ll be able to train a model to locate objects in images and video entirely client-side, no backend infrastructure required. Let‘s dive in!

What is TensorFlow.js?

TensorFlow.js is a JavaScript implementation of the popular TensorFlow machine learning library. It allows you to define, train, and deploy ML models entirely in the browser or in Node.js.

The key advantages of TensorFlow.js are:

1. No environment setup required – users can access ML-powered features without installing anything
2. Reduced latency – no need to communicate with a remote server to get predictions
3. Enhanced privacy – data never leaves the user‘s device
4. Ability to retrain and personalize models based on user interaction

TensorFlow.js supports many of the same model architectures as the Python library, including convolutional neural networks often used for computer vision tasks. It provides a Keras-style API for defining models and built-in functions for common preprocessing and postprocessing ops.

While not as fully-featured as the server-side TensorFlow framework, TensorFlow.js is an excellent choice when you want to add ML capabilities to a web app in a lightweight, privacy-preserving, user-friendly way.

Gathering a Custom Dataset

To train an object detection model, we first need a labeled dataset of images. There are a few paths you can take to generate this dataset:

1. Use an existing open source dataset like COCO, ImageNet, or Open Images
2. Gather images from sources like Google Images, Flickr, or a custom web scraper
3. Take your own photographs of the objects you want to detect

If you go with option 2 or 3, you‘ll then need to manually annotate the images with bounding boxes around each object of interest. This can be a time-consuming process, but luckily there are great open source tools like LabelImg or CVAT that make drawing and labeling boxes relatively painless.

Aim to annotate at least 100-200 images per object class for a small prototype. The more data you‘re able to generate, the better your model will perform. But be careful not to overfit to your training set – make sure to set aside a portion of annotated images for validation.

Converting to TFRecord Format

With our annotated images in hand, the next step is to convert them into a format that TensorFlow expects. The standard format is TFRecord, a binary serialization of structured data.

We‘ll write a script to parse our XML or JSON annotations and create a train.record and val.record file with the encoded image data and bounding box coordinates. Here‘s a simplified example:

def create_tf_example(image_path, annotations):
  with open(image_path, ‘rb‘) as f:
    encoded_image_data = f.read()

  width, height = image.size

  xmins = []
  xmaxs = []
  ymins = []
  ymaxs = []
  classes_text = []
  classes = []

  for annotation in annotations:
    xmins.append(annotation[‘xmin‘] / width)
    xmaxs.append(annotation[‘xmax‘] / width)
    ymins.append(annotation[‘ymin‘] / height)
    ymaxs.append(annotation[‘ymax‘] / height)
    classes_text.append(annotation[‘class‘].encode(‘utf8‘))
    classes.append(class_label_map[annotation[‘class‘]])

  tf_example = tf.train.Example(features=tf.train.Features(feature={
    ‘image/height‘: int64_feature(height),
    ‘image/width‘: int64_feature(width),
    ‘image/encoded‘: bytes_feature(encoded_image_data),
    ‘image/object/bbox/xmin‘: float_list_feature(xmins),
    ‘image/object/bbox/xmax‘: float_list_feature(xmaxs),
    ‘image/object/bbox/ymin‘: float_list_feature(ymins),
    ‘image/object/bbox/ymax‘: float_list_feature(ymaxs),
    ‘image/object/class/text‘: bytes_list_feature(classes_text),
    ‘image/object/class/label‘: int64_list_feature(classes),
  }))
  return tf_example

This function takes an image file path and its associated annotations, encodes the raw image bytes, converts absolute bounding box coordinates to relative ones, maps class labels to integers, and stores everything in a tf.train.Example protobuf message. We then write a series of these examples to a TFRecord file.

Configuring the Model

TensorFlow.js doesn‘t yet support training an object detection model from scratch, so we‘ll leverage a technique called transfer learning to adapt a pre-trained model to our custom dataset.

The first step is to download one of the base models from the TensorFlow detection model zoo. For this example we‘ll use SSD MobileNet v2, which is fast enough to run in real-time in a web browser.

Next, we‘ll define a pipeline configuration file specifying the model architecture, data inputs, and training hyperparameters. The most important settings to modify are:

• num_classes: The number of distinct objects we want to detect
• batch_size: How many images to feed to the model in each training step (reduce if you run out of memory)
• fine_tune_checkpoint: The path to the pre-trained model checkpoint
• fine_tune_checkpoint_type: Set to "detection" for the TF detection zoo models
• label_map_path: The path to a label map file defining the mapping of class names to indices

You can find the complete list of available config options in the sample config files included with the TF object detection API.

Training the Model

With our data and config ready, it‘s time to start training! We‘ll use the model_main_tf2.py script from the Object Detection API for this.

The basic process is:

1. Load the pre-trained model weights as an initial checkpoint
2. Feed batches of images and annotations from our TFRecord files to gradually tune the weights
3. Run the model on a validation set every so often to measure progress
4. Stop training when validation accuracy plateaus

Here‘s an example command to kick off training:

python model_main_tf2.py \
  --pipeline_config_path=configs/pipeline.config \
  --model_dir=training/ \
  --alsologtostderr

Depending on the size of your dataset and compute resources, training can take anywhere from a few minutes to a few hours. Keep an eye on the terminal output to track loss and validation metrics over time.

Exporting for TensorFlow.js

Once we‘re satisfied with our model‘s performance, the final step is to export it in the TensorFlow.js web format. We‘ll use the tensorflowjs_converter utility for this:

tensorflowjs_converter \
  --input_format=tf_saved_model \
  --output_node_names=‘detection_boxes,detection_classes,detection_scores,num_detections‘ \
  --saved_model_tags=serve \
  ./exported-models/my_model/saved_model \
  ./web_model

This command converts the saved model checkpoint to a web-friendly JSON format, optimizes it for inference, and saves the result in a web_model directory.

To make the model easier to use, we‘ll also create a labels file mapping the class indices back to human-readable names:

import json

with open(‘./annotations/label_map.pbtxt‘, ‘r‘) as f:
  label_map = f.readlines()

labels = []
for line in label_map:
  if "name" in line:
    labels.append(line.split(‘"‘)[1])

with open(‘web_model/labels.json‘, ‘w‘) as f:  
  json.dump(labels, f)

Deploying the Model

We now have everything we need to deploy our custom object detector on a website! The simplest way to use the model is to create an HTML page with a canvas element to display the results and a script tag to load the TensorFlow.js library and our model files.

Here‘s a minimal example:

<html>
  <head>
    <script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs"></script>
    <script src="https://cdn.jsdelivr.net/npm/@tensorflow-models/coco-ssd"></script>
  </head>
  <body>
    <img id="image" src="test.jpg"/>
    <canvas id="canvas" width="640" height="480"></canvas>

    <script>
      const model = await tf.loadGraphModel(‘web_model/model.json‘);
      const labels = await fetch(‘web_model/labels.json‘).then(x => x.json());

      const image = document.getElementById(‘image‘);
      const canvas = document.getElementById(‘canvas‘);
      const ctx = canvas.getContext(‘2d‘);

      const predictions = await model.executeAsync(tf.browser.fromPixels(image));

      const boxes = predictions[0].arraySync();
      const scores = predictions[1].arraySync();
      const classes = predictions[2].dataSync();

      ctx.drawImage(image, 0, 0);

      for (let i = 0; i < classes.length; i++) {
        if (scores[i] > 0.5) {
          const [x, y, width, height] = boxes[i];

          ctx.beginPath();
          ctx.rect(
            x * canvas.width, 
            y * canvas.height,
            width * canvas.width,
            height * canvas.height
          );
          ctx.lineWidth = 2;
          ctx.strokeStyle = ‘red‘;
          ctx.fillStyle = ‘red‘;
          ctx.stroke();

          ctx.fillText(
            labels[classes[i]] + ‘: ‘ + scores[i].toFixed(3),
            x * canvas.width, 
            y * canvas.height
          );
        }
      }
    </script>
  </body>
</html>

This loads an image, feeds it through the object detection model, and draws predicted bounding boxes and labels on a canvas. You could further extend it to support uploading images, capturing video frames from a webcam, or sending results to a backend service.

And with that, we have a complete end-to-end custom object detection pipeline running in TensorFlow.js! The model can now be embedded in any website or webapp to locate and tag objects client-side.

Some potential applications of this technology:

• An ecommerce site that lets users search for visually similar products just by taking a picture
• An educational app that helps students learn to identify plants and animals
• An assistive tool that describes the contents of images for visually impaired users
• An industrial system that flags defective parts rolling off an assembly line

The possibilities are endless! What will you build with custom object detection?

Alternatives and Tradeoffs

While running object detection locally in the browser is convenient for many use cases, there are times when you may want to consider alternative approaches.

If you need maximum accuracy and aren‘t worried about network latency, uploading images to a dedicated deep learning service like Google Cloud Vision API, Amazon Rekognition, or Azure Computer Vision will give you state-of-the-art results out of the box – no model training required. These APIs also typically offer additional features like OCR, facial recognition, and explicit content detection.

For applications sensitive to user privacy or running in low-connectivity environments, you can also perform object detection offline on-device using frameworks like TensorFlow Lite, Core ML, or Fritz AI. The downside is users need to install a native app rather than accessing functionality through a web browser.

Another technique to speed up object detection for real-time video applications is to run a tracking algorithm like SORT or Deep SORT after the initial detection step. This allows you to detect objects only every N frames and fill in the gaps with lightweight tracking, rather than running a computationally expensive forward pass on every frame.

Next Steps

There are many ways to further improve and extend our custom object detector:

1. Experiment with different base model architectures – deeper networks will generally achieve better accuracy at the cost of slower inference
2. Gather more training data to cover a wider variety of lighting conditions, camera angles, and object scales
3. Implement techniques like hard negative mining and data augmentation to make the model more robust
4. Set up a continual learning pipeline to automatically retrain the model as users contribute new images
5. Optimize the model for mobile devices using TensorFlow Lite and WebAssembly
6. Integrate the model into a more sophisticated application that takes action based on the detected objects

I hope this post has given you a practical overview of how to build a custom object detector with TensorFlow.js. The complete source code is available on GitHub – feel free to use it as a starting point for your own projects.

What objects will you teach your model to recognize? Let me know in the comments!