Machine Learning Image Classification - TensorFlow.js

28 Apr 2023  Amiya pattanaik  6 mins read.

Introduction

Machine learning image classification is a powerful technique that enables computers to analyze and categorize images based on patterns and features. It is a subfield of machine learning and computer vision that focuses on training algorithms to automatically recognize and classify images into predefined categories or labels. Image classification has gained significant attention and popularity due to its wide range of applications, including object recognition, medical imaging, autonomous vehicles, and more.

Understanding Image Classification

Image classification involves teaching a machine learning algorithm to identify and assign appropriate labels or categories to images based on their visual characteristics. This process typically involves two main stages: training and inference. During the training stage, the algorithm learns from a large set of labeled images, extracting meaningful features and patterns that differentiate one class from another. In the inference stage, the trained model is applied to new, unseen images to predict their corresponding labels.

The Role of Machine Learning

Machine learning algorithms play a crucial role in image classification. They utilize various techniques, such as deep learning, convolutional neural networks (CNNs), and transfer learning, to learn and extract high-level features from images. These algorithms can automatically learn and identify complex patterns, textures, shapes, and structures in images, enabling accurate classification across a wide range of image types and categories.

Applications of Image Classification

Image classification has diverse applications across various industries and domains. In healthcare, it aids in medical imaging analysis, enabling the identification of diseases, tumors, or anomalies in X-rays, MRIs, and CT scans. In autonomous vehicles, image classification helps in detecting and recognizing objects, pedestrians, traffic signs, and lane markings for safe navigation. E-commerce platforms leverage image classification to automatically categorize and tag products, improving search and recommendation systems. It is also used in security and surveillance systems, quality control in manufacturing, agricultural analysis, and much more.

Use Case

If we were to create a machine learning model to categorize the image with the ID “image” and log the top 5 predicted classes and their probabilities, this is a scenario to consider. Let’s look at how artificial intelligence might be used to address this problem.

Here’s an example use case for machine learning with Tensorflow.js: image classification using a pre-trained model. In this example, we’ll use a pre-trained MobileNet model to classify images into one of 1000 different categories.

Code Snippet

First, we’ll need to include Tensorflow.js in our HTML file:

index.html

<!DOCTYPE html>
<html>
  <head>
    <title>Image Classification with Tensorflow.js</title>
    <script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs@2.0.0/dist/tf.min.js"></script>
    <script src="https://cdn.jsdelivr.net/npm/@tensorflow-models/mobilenet@2.1.0/dist/mobilenet.min.js"></script>
  </head>
  <body>
    <img id="inputImage" src="image.jpg" width="300" height="300">
    <script src="index.js"></script>
  </body>
</html>

Next, we’ll load the pre-trained MobileNet model:

index.js

async function loadModel() {
  const model = await tf.loadGraphModel(
    "https://tfhub.dev/google/imagenet/mobilenet_v2_100_224/classification/4"
  );
  return model;
}

// Once we have the model, we can use it to classify an image

async function classifyImage(model) {
  const img = document.getElementById("image");
  const tensor = tf.browser.fromPixels(img).toFloat();
  const offset = tf.scalar(127.5);
  const normalized = tensor.sub(offset).div(offset);
  const batched = normalized.expandDims(0);
  const predictions = await model.predict(batched).data();
  const top5 = Array.from(predictions)
    .map((p, i) => {
      return {
        probability: p,
        className: IMAGENET_CLASSES[i],
      };
    })
    .sort((a, b) => b.probability - a.probability)
    .slice(0, 5);
  console.log(top5);
}

const IMAGENET_CLASSES = [
  "tench, Tinca tinca",
  "goldfish, Carassius auratus",
  "great white shark, white shark, man-eater, man-eating shark, Carcharodon carcharias",
  // and so on...
];

async function app() {
  const model = await loadModel();
  await classifyImage(model);
}

app();

To run the code, you can follow these steps:

  1. Create a new directory on your computer and create two files inside it: index.html and index.js.

  2. Copy the HTML code I provided above and paste it into index.html.

  3. Copy the JavaScript code I provided and paste it into index.js.

  4. Download an image you want to classify (such as a picture of a cat or a dog) and save it as image.jpg in the same directory as index.html and index.js.

  5. Open index.html in your web browser. You should see the image you downloaded displayed on the page.

  6. Open the developer console in your web browser (usually by pressing F12 or Ctrl+Shift+I) and you should see the top 5 predicted classes and their probabilities printed to the console.

Note that you’ll need to be connected to the internet for the code to work, as it needs to load the pre-trained model from a remote server. In this example, we’re using two script tags to load the latest versions of Tensorflow.js and the MobileNet model from the jsDelivr CDN (Content Delivery Network). The src attribute of the img tag is set to the path of the image file you want to classify, and the src attribute of the script tag is set to the path of your JavaScript code file (index.js).

Challenges and Advancements

Image classification poses several challenges, such as variations in lighting conditions, viewpoints, occlusions, and complex image backgrounds. Overcoming these challenges requires robust preprocessing techniques, advanced feature extraction methods, and large-scale labeled datasets for training. Recent advancements in deep learning, including architectures like ResNet, Inception, and EfficientNet, have significantly improved the accuracy and performance of image classification models. Additionally, transfer learning allows leveraging pre-trained models trained on large-scale datasets, enabling effective image classification with smaller labeled datasets.

Conclusion

Machine learning image classification has revolutionized our ability to analyze and understand images in diverse applications. It enables computers to automatically categorize images based on learned patterns and features. With advancements in deep learning and the availability of powerful computing resources, image classification has become more accurate and efficient. As technology continues to advance, we can expect further improvements in image classification algorithms, leading to enhanced capabilities and exciting new applications in the future.

The above code will classify the image with the ID “image” and log the top 5 predicted classes and their probabilities to the console. Note that we first load the image into a tensor, normalize it, and expand it to a batch size of 1 (since the model expects a batch of images as input). Then we call the predict method on the model to get the predictions, and map the probabilities to the corresponding class names in the IMAGENET_CLASSES array.

Of course, this is just a simple example - in practice, you’d likely want to train your own model or fine-tune a pre-trained model for your specific use case. But hopefully this example gives you a sense of how Tensorflow.js can be used for machine learning in the browser!

We encourage our readers to treat each other respectfully and constructively. Thank you for taking the time to read this blog post to the end. We look forward to your contributions. Let’s make something great together! What do you think? Please vote and post your comments.

Amiya Pattanaik
Amiya Pattanaik

Amiya is a Product Engineering Director focus on Product Development, Quality Engineering & User Experience. He writes his experiences here.