TensorFlow.js – Remove or restore segments of an image in the browser using DeepLab model

This article is a continuation of the previous article on Semantic Image Segmentation in the browser using TensorFlow.js pre-trained DeepLab model. Here, we will use the DeepLab model’s output to manipulate specific parts of the image based on user selection. Essentially, we will allow the user to remove or restore a specific segment of an image that corresponds to an object class identified by the model.

Since I have covered the basics of the DeepLab model in the previous article, let’s dive straight into the details of manipulating the image using the prediction output of the model.

Below is the complete code for the HTML file.

When the “Segment Image” button is clicked, the source image is segmented using the DeepLab model and the output segmented image is displayed besides the original image along with the legends at the bottom.

Below the legends section, there are two buttons – “Remove Selected Objects” and “Restore Selected Objects”. Users can click and select one or more of the legends and then click the “Remove Selected Objects” button to remove the objects corresponding to the selected legend from the image. Clicking on the “Restore Selected Objects” will restore those sections of the image that were previously removed.

How it is done

Before we get into the details, let’s take a look at the DeepLab model’s prediction output.

Prediction object

A sample output prediction object is given below.

{
    "legend": {
        "background": [ 0, 0, 0 ],
        "person": [ 192, 128, 128 ],
        "motorbike": [ 64, 128, 128 ]
    },
    "height": 385,
    "width": 513,
    "segmentationMap": {
        "0": 0,
        "1": 0,
        "2": 0,
        "3": 255,
        "4": 0,
        "5": 0,
        "6": 0,
        "7": 255,
        "8": 0,
        ..
        "790016": 0,
        "790017": 0,
        "790018": 0,
        "790019": 255
    }
}

The prediction object contains the following 4 attributes within it.

1. legend
The legend attribute refers to an object that contains the r, g, b color values for each of the identified object classes.

2. height and width
The height and width attribute refers to the height and width of the image based on which the segmentation is performed.

3. segmentationMap
The segmentationMap attribute refers to an object that contains the pixel wise labelling by providing the color code for each pixel in the image based on the object identified at that pixel. Basically the object corresponds to the data attribute of the ImageData object.

For every pixel, the ImageData  object contains 4 values

• R – the color red (0 to 255)
• G – the color green (0 to 255)
• B – the color blue (0 to 255)
• A – The alpha channel (0 to 255; 0 is transparent and 255 is fully visible)

As a summary, in the above example,

• The output image data was of the size 513 X 385 (width X height). 
• Total pixels = 197505 (by multiplying, 513 * 385)
• The segmentationMap object contains 790020 values. 4 values (RGBA) for each of the 197505 pixels (197505 * 4 = 790020)

Image manipulation

The idea here is to take the color of the selected legend, specifically it’s r, g, b values and match it against the r, g, b values for each pixel in the image using the segmentationMap object, and set the alpha value of that pixel to either 0 or 255 to make that particular pixel transparent or visible respectively.

Objects to remove / restore from the image

When a user clicks on one or more of the legends, the storeObjectColor() function is invoked. This function is provided as the event handler to the span elements that are used to display the legends.

In this function, we mainly extract the background color of the span element, essentially the element’s r, g, b color values, and store it in an object named objectColors, which is in the global scope.

Remove / Restore segments from the image

The removeOrRestoreSelectedObjects() function is responsible for removing or restoring specific segments of the image. This function is provided as the onclick event handler for both the Remove Selected Objects and Restore Selected Objects button.

Using the event object, generated by the button’s onclick event, determine which button was clicked and set the alphaValueToSet variable to 0 if Remove Selected Objects was clicked, 255 otherwise.

As a next step, get the individual objects from the output prediction object which is stored in the global variable – prediction.

Create a new Canvas element to display the manipulated image output. Set the width and height of canvas using the height and width values from the prediction output. Draw the source image onto this canvas.

Now, get the ImageData using the canvas’ ctx.getImageData() function. This provides pixel wise data for the image and for each pixel, the ImageData.data object contains the 4 values – r, g, b and a – any value between 0 and 255.

As we saw above, the segmentationMap object contains four values for each pixel. So, for each iteration, we will be reading 4 consecutive values from this object, hence the i+=4 increment in the for loop.

In each iteration (for each pixel), compare the r, g, b values from the segmentationMap object to the r, g, b values for each of the objects present in the objectColors object. If it matches, then set the alpha value of the pixel in the ImageData.data object (imageData obtained from the canvas in the above step) using the variable alphaValueToSet which was determined earlier based on the button that was clicked to invoke this function.

Finally, invoke the canvas’ ctx.putImageData() function and pass the manipulated ImageData object to draw the output image. Display the canvas besides the original image.

The source code for this example is present in this Github repo.

Hope this was useful. Happy Learning !

References

1. Github repo – Repository of the DeepLab v3 model in TensorFlow.js pre-trained model
2. Github repo – Repository of DeepLab v2 model in TensorFlow library
3. Image Segmentation – Wikipedia