Portrait-Segmentation

Real-time Automatic Deep Matting For Mobile Devices With Mobile-Unet

Portrait segmentation refers to the process of segmenting a person in an image from its background. Here we use the concept of semantic segmentation to predict the label of every pixel (dense prediction) in an image.

Here we limit ourselves to binary classes (person or background) and use only plain portrait-selfie images for matting.

Dependencies

Tensorflow(>=1.14.0), Python 3
Keras(>=2.2.4), Kito
Opencv(>=3.4), PIL, Matplotlib

pip uninstall -y tensorflow
pip install -U tf-nightly
pip install keras
pip install kito

Prerequisites

Download training data-set
GPU with CUDA support
Download caffe harmonization model

Dataset

The dataset consists of 18698 human portrait images of size 128x128 in RGB format, along with their masks(ALPHA). Here we augment the PFCN dataset with (handpicked) portrait images form supervisely dataset. Additionaly, we download random selfie images from web and generate their masks using state-of-the-art deeplab-xception model for semantic segmentation.

Now to increase the size of dataset, we perform augmentation like cropping, brightness alteration and flipping. Since most of our images contain plain background, we create new synthetic images using random backgrounds (natural) using the default dataset, with the help of a python script.

Besides the aforesaid augmentation techniques, we normalize(also standardize) the images and perform run-time augmentations like flip, shift and zoom using keras data generator and preprocessing module.

Model Architecture

Here we use Mobilent v2 with depth multiplier 0.5 as encoder (feature extractor).

For the decoder part, we have two variants. You can use a upsampling block with either Transpose Convolution or Upsample2D+Convolution. In the former case we use a stride of 2, whereas in the later we use resize bilinear for upsampling, along with Conv2d. Ensure proper skip connections between encoder and decoder parts for better results.

Additionaly, we use dropout regularization to prevent overfitting.It also helps our network to learn more robust features during training.

Here is the snapshot of the upsampled version of model.

How to run

Download the dataset from the above link and put them in data folder. After ensuring the data files are stored in the desired directorires, run the scripts in the following order.

1. python train.py # Train the model on data-set
2. python eval.py checkpoints/up_super_model-102-0.06.hdf5 # Evaluate the model on test-set
3. python export.py checkpoints/up_super_model-102-0.06.hdf5 # Export the model for deployment
4. python test.py test/four.jpeg # Test the model on a single image
5. python webcam.py test/beach.jpg # Run the model on webcam feed
6. python segvideo.py test/sunset.jpg # Apply blending filters on video

You may also run the Jupyter Notebook (ipynb) in google colaboratory, after downloading the training dataset.

In case you want to train with a custom dataset, check out the scripts in utils directory for data preparation.

Training graphs

Since we are using a pretrained mobilentv2 as encoder for a head start, the training quickly converges to 90% accuracy within first couple of epochs. Also, here we use a flexible learning rate schedule (ReduceLROnPlateau) for training the model.

Training Loss

Training Accuracy

Validation Loss

Validation Accuracy

Learning Rate Schedule

Demo

Result

Here the inputs and outputs are images of size 128x128. The first row represents the input and the second row shows the corresponding cropped image obtained by cropping the input image with the mask output of the model.

Accuracy: 96%

Android Application

Real-time portrait video in android application

(Shot on OnePlus 3 😉)

Fun With Filters (Python)

Let's add some filters to harmonize our output image with the background. Our aim is to give a natural blended feel to the output image i.e the edges should look smooth and the lighting(colour) of foreground should match(or blend) with its background.

The first method is the traditional alpha blending, where the foreground images are blended with background using the blurred(gaussian) version of the mask.In the smooth-step filter, we clamp the blurred edges and apply a polynomial function to give a curved appearence to the foreground image edges.Next, we use the colour transfer algorithm to transfer the global colour to the foreground image.Also, opencv(computational photography) provides a function called seamless clone to blend an image onto a new background using an alpha mask.Finally, we use the dnn module of opencv to load a colour harmonization model(deep model) in caffe and transfer the background style to the foreground.

Here are some sample results:-

For live action, checkout the script segvideo.py to see the effects applied on a webcam video.

Also download the caffe model and put it inside models/caffe folder.

Keyboard Controls:-

Hold down the following keys for filter selection.

C- Colour transfer
S- Seamless clone
M- Smooth step
H- Colour harmonize

Move the slider to change the background image.

Tensorflowjs: Running the model on a browser

To ensure that your applications runs in a platform independent way(portabe), the easiest way is to implement them as a web-application and run it using a browser.You can easily convert the trained model to tfjs format and run them using javascript with the help of tensorflowjs conversion tools.If you are familiar with React/Vue js , you can easily incorporate the tfjs into you application and come up with a really cool AI webapp, in no time!!!

Here is the link to the portrait segmentation webapp: CVTRICKS

If you want to run it locally, start a local server using python SimpleHTTPServer. Initially configure the hostname, port and CORS permissions and then run it using your browser.

NB: The application is computaionally intensive and resource heavy.

Key Insights and Drawbacks

Always start experimentation with standard/pretrained networks. Also try out default/standard hyperparameter settings before experimentation.
Make sure your ground truth is correct/uncorrupted and is in desired format before training (even standard dataset).
For mobile devices, make sure you use a mobile-friendly architecture (like mobilenet) for training and deployment.
Using google colaboratory along with google drive for training was EASY & FUN.It provides high end GPU (RAM also) for free.
Some of the mobile optimization tools(even TF) are still experimental (GPU deegate, NNAPI, FP16 etc.) and are buggy.They support only limited operations and edge devices.
Even state-of-the art segmenation models(deeplab-xception) seems to suffer from false positives (even at higher sizes), when we test them on a random image.
The segmentaion maps produced at this low resolution (128x128) have coarse or sharp edges (stair-case effect), especially when we resize them to higher resolution.
To tackle the problem of coarse edges, we apply a blur filter (also antialiasing at runtime) using opencv and perform alpha blending with the background image. Other approach was to threshold the blurred segmentation map with smooth-step function using GLSL shaders.
In android we can use tensorflow-lite gpu-delegate to speed up the inference.It was found that flattening the model output into a rank 1 (or 2) tensor helped us to reduce the latency due to GPU-CPU data transfer.Also this helped us to post-process the mask without looping over a multi-dimensional array.
Using opencv (Android NEON) for post-processing helped us to improve the speed of inference.But this comes at the cost of additional memory for opencv libraray in the application.
Still, there is a scope for improving the latency of inference by performing all the postprocessing in the GPU, without transfering the data to CPU. This can be acheived by using opengl shader storge buffers (SSBO). We can configure the GPU delegate to accept input from SSBO and also access model output from GPU memory for further processing (wihout CPU) and subsequent rendering.
The difference between the input image frame rate and output mask generation frame rate may lead to an output(rendering), where the segmentation mask lags behind current frame.This stale mask phenemena arises due to the model(plus post-processing) taking more than 40ms (corr. to 25 fps input) per frame (real-time video). The solution is to render then output image in accordance to the mask generation fps (depends on device capability) or reduce the input frame rate.
If your segmentaion-mask output contains minor artifacts, you can clean them up using morphological operations like opening or closing. However it can be slightly expensive if your output image size is large, especially if you perform them on every frame output.
If the background consists of noise, clutter or objects like clothes, bags etc. the model fails miserably.
Even though the stand-alone running time of exported (tflite) model is low(around 100 ms),other operations like pre/post-processing, data loading, data-transfer etc. consumes significant time in a mobile device.
The models trained with resize bilinear(default parameters) in tensorflow seems to suffer from a problem of mask shifting.This problem occurs if the image size is even (i.e bilinear_128 model in our case).The pixels in the output mask seems to be sligtly shifted horizontaly in one direction(left/right).
Opencv dnn module provides support for running models trained on popular platforms like Caffe,Tensorflow, Torch etc.It supports acceleration through OpenCL, Vulkan, Intel IE etc.It also supports variety of hardwares like CPU,GPU and VPU.Finally, we can also run smaller FP16 models for improved speed.
Once you are familiar with tensorflow, it is fairly easy to train and perform inference using tensorflowjs.It also comes with support of WebGL backend for accelerating the inference and training process.Th main advantage is the portability of the application i.e it can be run on PC, phones or tablet without any modifications.

TODO

Port the code to TF 2.0
Use a bigger image for training(224x224)
Try quantization-aware training
Train with mixed precision (FP16)
Optimize the model by performing weight pruning
Improve accuracy & reduce artifacts at runtime
Incroporate depth information and boundary refinement techniques
Apply photorealistic style transfer on foreground based on background image