Getting Started with Image Segmentation Using TensorFlow.js

Background removal for people with BodyPix.js

Many powerful and well-equipped image editing tools are available that allow you to remove backgrounds, much like a green screen effect. These tools are packed with powerful features that enable us to make powerful and transformative edits to images—background removal being a key functionality.

But what if we could create our own background removal tool? Interestingly, we can use AI to do just this, automatically segmenting out a person’s body from an image and then cutting out the background.

One of the ways to remove background from a persons’ body is to apply machine learning techniques with the help a computer vision technique known as segmentation.

To briefly describe segmentation—images can be divided into various parts called segments. The unimportant parts are left out, and only the segments containing important information undergo image processing. Then, there is a grouping of pixels that have similar attributes that are separated or segmented from other groupings of pixels.

In this tutorial, we are going to make use of this technique to segment out a persons’ body from the background and remove the unimportant segment (i.e. the background) using a specific image editing JavaScript library.

The tutorial is quick and easy to understand and implement. The idea is to make use of a TensorFlow.js model that enables us to separate and remove the background from an image including a person by using the segmentation package known as BodyPix.

The BodyPix package allows us to segment out an image into pixels that are (and are not) part of a person’s body. Then, we will make use of the jimp package to edit the image and remove the background.

So let’s get started!

Install Necessary Packages

For this project, we’re going to use a Node.js project. First, we need to initialize a Node project and then install the required packages. In order to install the packages, we’re going to use yarn and run the following code in the project terminal:

Next, we need to import all the required packages to the index.js file, as shown in the code snippet below:

Now, we can start by creating a class to contain all the necessary functions. Inside the class, we’re going to create a function that handles the input image.

Loading the Model

The first function we’re going to create will be to load the BodyPix model. For that, we need to define the global variable named bodymodel to handle the bodyPix object. Then, the function loadModel can be used to load the TensorFlow model, as shown in the code snippet below:

Running Prediction

The second function will handle the call of the loadModel function, and also start segmentation on the image to distinguish a person’s body from the background. We need to pass the image in the segmentPersonParts method of the bodymodel instance, as directed in the code snippet below:

Now, in the third and last function, we’re going to remove the background from the processed image. The idea is to get the segmented image and pass the result to the jimp module.

We’re going to iterate over every pixel in the image and remove the color if it is not part of the body. After completion, we will write and output the file of the image with the background removed. The code to implement this background removal function is provided in the code snippet below:

Usage Example

For usage, we read the image using the fs file package and pass it to the removeBg function that we implemented earlier. The process is done asynchronously, and in order to handle it, we’re using async/await, as shown in the code snippet below:

Now, if we test it out, we’ll get the result as displayed in the demo screenshots below:

Hence, we have successfully implemented the background removal tool using TensorFlow.js, BodyPix, and the jimp module.

You can try it out with different images. The image with multiple people should work as well. Try implementing your own UI interface with optimized logic for better results.

Conclusion

This AI-powered segmentation model can be a useful tool for graphic designers or someone trying to remove the background from an image. No need to use complicated photo editing software for this task. How interesting and amazing is that?

The aim was to deliver a tutorial on how to make use of the TensorFlow.js model, BodyPix, and jimp to implement a tool that handles the removal of background from images that include a human subject. It was indeed pretty easy. But this is just the first step.

Now, the challenge is to use other optimized models that help segment other objects types instead of a person’s body. An example model could recognize and segment IT equipment. Make sure to try different things out and learn as you implement.

See you next time! Until then, learn and grow!

Peace out, folks!

Avatar photo

Fritz

Our team has been at the forefront of Artificial Intelligence and Machine Learning research for more than 15 years and we're using our collective intelligence to help others learn, understand and grow using these new technologies in ethical and sustainable ways.

Comments 0 Responses

Leave a Reply

Your email address will not be published. Required fields are marked *