DEV Community

loading...
Cover image for Image recognition with TensorFlow.js

Image recognition with TensorFlow.js

basilebong profile image Basile Bong Originally published at basilebong.com Updated on ・9 min read

In this post I will show you how to create a simple image classifier, without any machine learning knowledge using a pretrained model form the TensorFlow team.

Checkout the demo and the source code.

Table of contents

What you need

  1. Knowledge of JavaScript, CSS and HTML
  2. A code editor (I recommend VS Code)
  3. A local server (I recommend live server VS Code extension).

Let's start!

Initializing the app

Create a new folder and add 3 files:

.
├── app.css
├── app.js
└── index.html
Enter fullscreen mode Exit fullscreen mode

Edit index.html and add the following code:

<!DOCTYPE html>
<html lang="en">
  <head>
    <meta charset="UTF-8" />
    <meta name="viewport" content="width=device-width, initial-scale=1.0" />
    <title>My web app</title>

    <!-- Custom style -->
    <link rel="stylesheet" href="app.css" />

    <!-- Google font -->
    <link rel="preconnect" href="https://fonts.gstatic.com" />
    <link
      href="https://fonts.googleapis.com/css2?family=Source+Sans+Pro&display=swap"
      rel="stylesheet"
    />
  </head>
  <body>
    <script src="app.js"></script>
  </body>
</html>
Enter fullscreen mode Exit fullscreen mode

File uploader

In order to classify an image we first need to let the user upload a picture.

Edit index.html and add the following code inside <body></body>:

<main>
  <section class="image-section">
    <img src="" id="image" />
  </section>

  <section class="file-section">
    <div class="file-group">
      <label for="file-input">Upload a picture</label>
      <input type="file" id="file-input" />
    </div>
  </section>
</main>
Enter fullscreen mode Exit fullscreen mode

Edit app.css to enhance the look:

body {
  font-family: "Source Sans Pro", sans-serif;
}

main {
  width: 100%;
  min-height: 100vh;
  display: flex;
  justify-content: center;
  align-items: center;
  flex-direction: column;
}

section {
  margin: 2rem 1rem;
}

.file-group {
  padding: 1rem;
  background: #efefef;
  border-radius: 1rem;
}

#image {
  max-width: 100%;
  width: 400px;
  height: auto;
  display: block;
  margin: auto;
}

.image-section {
  display: none;
  position: relative;
}

.image-loaded .image-section {
  display: block;
}
Enter fullscreen mode Exit fullscreen mode

The next step is to create the JavaScript code that will handle the file upload and display the image on the page.

To help us manipulate the image and the file input, we are going to save those two DOM elements into some variables.

const fileInput = document.getElementById("file-input");
const image = document.getElementById("image");
Enter fullscreen mode Exit fullscreen mode

When the user uploads a new image, the getImage() function is triggered.

fileInput.addEventListener("change", getImageDataUrl);
Enter fullscreen mode Exit fullscreen mode

The goal is to display the uploaded image inside our web application. To do so create a new function getImage() and write it before the event listener.

function getImage() {
  // ...
}
Enter fullscreen mode Exit fullscreen mode

First we have to check if the file has been uploaded with success. So add the following code inside the getImage() function.

function getImage() {
  if (!fileInput.files[0]) throw new Error("Image not found");
  const file = fileInput.files[0];
}
Enter fullscreen mode Exit fullscreen mode

Then we need to read the file that has been uploaded with FileReader. You can find more information on the mozilla.org webpage.

To display the image inside our web app, we need a URL that can be set as src attribute of the <img id="image"/> tag. This URL will be generated by the readAsDataURL(file) method that returns a data URL.

Data URLs, URLs prefixed with the data: scheme, allow content creators to embed small files inline in documents. They were formerly known as "data URIs" until that name was retired by the WHATWG. -mozilla.org

const reader = new FileReader();
Enter fullscreen mode Exit fullscreen mode

The FileReader is asynchronous. We have to wait for the result with onload before we can display the image.

reader.onload = function (event) {
  image.setAttribute("src", event.target.result);
  document.body.classList.add("image-loaded");
};

reader.readAsDataURL(file);
Enter fullscreen mode Exit fullscreen mode

Finally, your app.js file should look like this:

const fileInput = document.getElementById("file-input");
const image = document.getElementById("image");

/**
 * Get the image from file input and display on page
 */
function getImage() {
  // Check if an image has been found in the input
  if (!fileInput.files[0]) throw new Error("Image not found");
  const file = fileInput.files[0];

  // Get the data url form the image
  const reader = new FileReader();

  // When reader is ready display image.
  reader.onload = function (event) {
    image.setAttribute("src", event.target.result);
    document.body.classList.add("image-loaded");
  };

  // Get data url
  reader.readAsDataURL(file);
}

/**
 * When user uploads a new image, display the new image on the webpage
 */
fileInput.addEventListener("change", getImage);
Enter fullscreen mode Exit fullscreen mode

Image classification

Thanks to TensorFlow and its pretrained model, the classification of images becomes very easy. A model is a file that has been trained over a set of data in order to recognize certain patterns. I will not deep dive into this subject, but if you want to know more I recommend you to read the Microsoft documentation.

To start using TenserFlow.js and it's pretrained image classification model (mobilenet) we will have to edit the index.html file and add the following lines into the <head></head>:

<!-- TensorFlow-->
<script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs@1.0.1"></script>
<!-- TensorFlow pretrained model-->
<script src="https://cdn.jsdelivr.net/npm/@tensorflow-models/mobilenet@1.0.0"></script>
Enter fullscreen mode Exit fullscreen mode

Loading

To avoid that the web application is used before the model is fully loaded, we will display a loader inside our web application.

Edit index.html, add the class .loading to the <body></body>, and the HTML markup of the loader.

<!-- Add loading class -->
<body class="loading">
  <main>
    <!-- Add this -->
    <div class="loader">
      <h2>Loading ...</h2>
    </div>

    <!-- ... -->
  </main>

  <script src="app.js"></script>
</body>
Enter fullscreen mode Exit fullscreen mode

Next we will have to hide the file input during the loading process. To do so edit app.css and add the following code:

.loading .loader {
  display: block;
}

.loader {
  display: none;
}

.loading .image-section,
.loading .file-section {
  display: none;
}
Enter fullscreen mode Exit fullscreen mode

Next we will have to load the model in our app.js file. Add the following code at the end of your file.

// Async loading
mobilenet.load().then(function (m) {
  // Save model
  model = m;

  // Remove loading class from body
  document.body.classList.remove("loading");

  // When user uploads a new image, display the new image on the webpage
  fileInput.addEventListener("change", getImage);
});
Enter fullscreen mode Exit fullscreen mode

As you can see addEventListener has been moved inside the loading function. We also need to add an empty model variable at the beginning of our code:

const fileInput = document.getElementById("file-input");
const image = document.getElementById("image");
let model;

// ...
Enter fullscreen mode Exit fullscreen mode

Finally, your code should look like this:

const fileInput = document.getElementById("file-input");
const image = document.getElementById("image");
let model;

/**
 * Get the image from file input and display on page
 */
function getImage() {
  // Check if an image has been found in the input
  if (!fileInput.files[0]) throw new Error("Image not found");
  const file = fileInput.files[0];

  // Get the data url form the image
  const reader = new FileReader();

  // When reader is ready display image
  reader.onload = function (event) {
    const dataUrl = event.target.result;
    image.setAttribute("src", dataUrl);
    document.body.classList.add("image-loaded");
  };

  // Get data URL
  reader.readAsDataURL(file);
}

/**
 * Load model
 */
mobilenet.load().then(function (m) {
  // Save model
  model = m;

  // Remove loading class from body
  document.body.classList.remove("loading");

  // When user uploads a new image, display the new image on the webpage
  fileInput.addEventListener("change", getImage);
});
Enter fullscreen mode Exit fullscreen mode

Now the UI is only displayed when the model is fully loaded.

Using the model

The mobilenet model needs an <img /> HTML element as parameter that has a defined width and height. Currently this two attributes are missing. To add them we will have to edit the getImage() function inside the app.js file.

To get the size of the image we will use the Image class.

The Image() constructor creates a new HTMLImageElement instance. It is functionally equivalent to document.createElement('img'). -mozilla.org

function getImage() {
  // ...
  reader.onload = function (event) {
    // ...

    // Create image object
    const imageElement = new Image();
    imageElement.src = dataUrl;

    // When image object is loaded
    imageElement.onload = function () {
      // Set <img /> attributes
      image.setAttribute("src", this.src);
      image.setAttribute("height", this.height);
      image.setAttribute("width", this.width);

      // Classify image
      classifyImage();
    };

    // ...
  };

  //..
}
Enter fullscreen mode Exit fullscreen mode

The classifyImage() function does not exist yet.
Now your getImage() function should look like this:

function getImage() {
  // Check if an image has been found in the input
  if (!fileInput.files[0]) throw new Error("Image not found");
  const file = fileInput.files[0];

  // Get the data url form the image
  const reader = new FileReader();

  // When reader is ready display image
  reader.onload = function (event) {
    // Ge the data url
    const dataUrl = event.target.result;

    // Create image object
    const imageElement = new Image();
    imageElement.src = dataUrl;

    // When image object is loaded
    imageElement.onload = function () {
      // Set <img /> attributes
      image.setAttribute("src", this.src);
      image.setAttribute("height", this.height);
      image.setAttribute("width", this.width);

      // Classify image
      classifyImage();
    };

    // Add the image-loaded class to the body
    document.body.classList.add("image-loaded");
  };

  // Get data URL
  reader.readAsDataURL(file);
}
Enter fullscreen mode Exit fullscreen mode

After a lot of preparation we can finally use the model with only a view lines of code. First we will create a new function called classifyImage().

function classifyImage() {
  model.classify(image).then(function (predictions) {
    console.log("Predictions: ");
    console.log(predictions);
  });
}
Enter fullscreen mode Exit fullscreen mode

Run the application and you should see the predictions in your developer console!

Display the prediction

The last thing we want to do is to display a sentence that describes the picture.
First we need to add a place in our HTML code where the description can be placed.
Edit index.html:

<!-- ... -->
<section class="image-section">
  <img src="" id="image" />
  <div class="image-prediction"></div>
</section>
<!-- ... -->
Enter fullscreen mode Exit fullscreen mode

Then add the necessary CSS in app.css:

/* Black overlay over the image */
.image-section::before {
  content: "";
  z-index: 2;
  position: absolute;
  height: 100%;
  width: 100%;
  background: linear-gradient(transparent, transparent, #000000);
}

.image-prediction {
  position: absolute;
  bottom: 1rem;
  text-align: center;
  font-size: 18px;
  color: #fff;
  left: 0;
  right: 0;
  z-index: 3;
}
Enter fullscreen mode Exit fullscreen mode

Then open app.js and change the classifyImage() function:

function classifyImage() {
  model.classify(image).then((predictions) => {
    displayDescription(predictions);
  });
}
Enter fullscreen mode Exit fullscreen mode

The predictions are an array of predictions. Each prediction contains a className and a probability.

[
  {
    className: "chow, chow chow",
    probabilty: 0.856542315,
  },
];
Enter fullscreen mode Exit fullscreen mode

The first thing we are going to do is to sort the results and only keep the predictions with the height probability. In this case the probability needs to be at least 20% (which is super low). If it is lower we display an error message.

function displayDescription(predictions) {
  const result = predictions.sort((a, b) => a > b)[0];

  if (result.probability > 0.2) {
    const probability = Math.round(result.probability * 100);

    // Display result
    description.innerText = `${probability}% shure this is a ${result.className.replace(
      ",",
      " or"
    )} 🐶`;
  } else description.innerText = "I am not shure what I should recognize 😢";
}
Enter fullscreen mode Exit fullscreen mode

Finally, your code should look like this:

const fileInput = document.getElementById("file-input");
const image = document.getElementById("image");
const description = document.getElementById("prediction");

let model;

/**
 * Display the result in the page
 */
function displayDescription(predictions) {
  // Sort by probability
  const result = predictions.sort((a, b) => a > b)[0];

  if (result.probability > 0.2) {
    const probability = Math.round(result.probability * 100);

    // Display result
    description.innerText = `${probability}% shure this is a ${result.className.replace(
      ",",
      " or"
    )} 🐶`;
  } else description.innerText = "I am not shure what I should recognize 😢";
}

/**
 * Classify with the image with the mobilenet model
 */
function classifyImage() {
  model.classify(image).then((predictions) => {
    displayDescription(predictions);
  });
}

/**
 * Get the image from file input and display on page
 */
function getImage() {
  // Check if an image has been found in the input
  if (!fileInput.files[0]) throw new Error("Image not found");
  const file = fileInput.files[0];

  // Get the data url form the image
  const reader = new FileReader();

  // When reader is ready display image
  reader.onload = function (event) {
    // Ge the data url
    const dataUrl = event.target.result;

    // Create image object
    const imageElement = new Image();
    imageElement.src = dataUrl;

    // When image object is loaded
    imageElement.onload = function () {
      // Set <img /> attributes
      image.setAttribute("src", this.src);
      image.setAttribute("height", this.height);
      image.setAttribute("width", this.width);

      // Classify image
      classifyImage();
    };

    // Add the image-loaded class to the body
    document.body.classList.add("image-loaded");
  };

  // Get data URL
  reader.readAsDataURL(file);
}

/**
 * Load model
 */
mobilenet.load().then((m) => {
  // Save model
  model = m;

  // Remove loading class from body
  document.body.classList.remove("loading");

  // When user uploads a new image, display the new image on the webpage
  fileInput.addEventListener("change", getImage);
});
Enter fullscreen mode Exit fullscreen mode

Congratulations

Congratulations, you did it!

Note that this application is not fully finished:

  • We didn't check if the uploaded file is an image
  • We didn't check if the image is a dog
  • We didn't check for upload errors

Credits

Discussion

pic
Editor guide