Real-Time semantic segmentation in the browser using TensorFlow.js

Perform semantic segmentation directly from your browser. There is no need for complex libraries or powerful computers.

Semantic segmentation is the problem of detecting and delineating each object of interest appearing in an image. Currently, there are several approaches that solve this problem and produce results as seen below.

Figure 1: Semantic segmentation example ( Source )

This kind of segmentation is predicting every pixel in the image and is known as Dense Prediction as well. It’s important to notice that the instances of the same class are not being separated, the model only cares about the pixel’s category. As shown in Figure 1, the method can say that there are chairs in certain positions, but can not distinguish them.

One of the main applications of this technique is in Autonomous Vehicles, where cars need to understand their environment. Semantic Segmentation is able to assign a meaning to the scenes and put the car in the context, indicating the lane position, if there is some obstruction, as fallen trees or pedestrians crossing the road, and recognizing other cars.

Therefore, apply Semantic Segmentation algorithms in urban street scenes is one of the main Computer Vision challenges nowadays. A popular dataset to evaluate model performance is the Cityscapes . It contains 30 classes from 50 different cities varying the season and wheater conditions.

Figure 2 shows how different algorithms have reached the state of the art in this dataset over time.

Figure 2: Semantic Segmentation on Cityscapes dataset ( Source )

In this project, it was used a mid-level model that can deliver a reasonable precision and run in real-time. The RefineNet [2] was firstly introduced at the end of 2016 from researches of The University of Adelaide and converted to a Light-Weight model in 2018 [3], allowing real-time inferences.

TensorFlow.js

TensorFlow is an open-source library widely-used to create Machine Learning models, mostly in python projects. The emergence of Tensorflow.js allowed that the same development, train and tests of models could be done in the browser.

ML running in the browser means that from a user’s perspective, there’s no need to install any libraries or drivers. Just open a webpage, and the program is ready to run. In addition, it’s ready to run with GPU acceleration. TensorFlow.js automatically supports WebGL, and will accelerate the code behind the scenes when a GPU is available. Users may also open your webpage from a mobile device, in which case your model can take advantage of sensor data, say from a gyroscope or accelerometer. Finally, all data stays on the client, making TensorFlow.js useful for low-latency inference, as well as for privacy-preserving applications [4].

Weights file

TensorFlow.js layers format is a directory containing a model.json file and a set of sharded weights files in a binary format. The model.json file contains both the model topology (aka "architecture" or "graph": a description of the layers and how they are connected) and a manifest of the weight files [5].

The JSON weights are provided into the Light-Weight RefineNet repository [6], but it is also possible to convert PyTorch or Keras model to the TensorFlow.js format [7].

Serving TensorFlow.js models

Tensorflow.js models have to be served through an URL to be loaded in javascript. As mentioned before, the model is based on a main JSON file and shared weights files in a binary format. The idea is to create an HTTP server that will make the model available in a URL allowing requests and be thread as REST API. When loading the model, the TensorFlow.js will do the following requests:

The http-server can be installed using the following command:

npm-install http-server -g

Once you have the program installed, you can serve your model going into your weights folder and running the command below:

http-server -c1 --cors .

Then, the following piece of code will load the tensorflow.js model.

Making predictions

Once the model is loaded, the process of making predictions is pretty simple. You have to load a frame and use the model.predict(frame) command to get predictions from the model. The return is a frame in the PASCAL VOC format that has to be converted and rendered, the code to do it can be found in the project repository .

I decided to create a simple application that read a frame from the webcam and generate a semantic segmentation, as the example below:

Figure 3: Refine Net — Semantic segmentation example

To test the project in a realistic scenario, I downloaded some pictures on my phone to feed the algorithm, the video below shows the application running in realtime in the browser:

The application has a small delay, mostly due to the frame reconstruction time, but the overall results are pretty good. TensorFlow.js unlocked a lot of possibilities allowing that machine learning models that needed powerful computers and complex libraries could be run smoothly in the browser and using few computational resources.

All the source code and instruction to run the project can be found at GitHub . If you have any questions or suggestion you can reach me out at Linkedin .

Thanks for reading :)

References

[1] Romera-Paredes, Bernardino, and Philip Hilaire Sean Torr. “Recurrent instance segmentation.” European conference on computer vision . Springer, Cham, 2016.

[2] Lin, Guosheng, et al. “Refinenet: Multi-path refinement networks for high-resolution semantic segmentation.” Proceedings of the IEEE conference on computer vision and pattern recognition . 2017.

[3] Nekrasov, Vladimir, Chunhua Shen, and Ian Reid. “Light-weight refinenet for real-time semantic segmentation.” arXiv preprint arXiv:1810.03272 (2018).

[4] Medium. 2020. Introducing Tensorflow.Js: Machine Learning In Javascript . [online] Available at: <https://medium.com/tensorflow/introducing-tensorflow-js-machine-learning-in-javascript-bf3eab376db> [Accessed 3 June 2020].

[5] https://www.tensorflow.org/js/tutorials/conversion/import_keras

[6] https://github.com/DrSleep/light-weight-refinenet

[7] https://www.tensorflow.org/js/guide/conversion