开源软件名称（OpenSource Name）：

开源软件地址(OpenSource Url)：

开源编程语言(OpenSource Language)：

开源软件介绍(OpenSource Introduction)：

Portrait-Segmentation

Real-time Automatic Deep Matting For Mobile Devices

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. This technique is widely used in computer vision applications like background replacement and background blurring on mobile devices.

Here we limit ourselves to binary classes (person or background) and use only plain portrait-selfie images for matting. We experimented with the following architectures for implementing a real-time portrait segmentation model for mobile devices.

1. Mobile-Unet
2. DeeplabV3+
3. Prisma-Net
4. Portrait-Net
5. Slim-Net
6. SINet

The models were trained with standard(and custom) portrait datasets and their performance was compared with the help of standard evaluation metrics and benchmarking tools. Finally, the models were deployed on edge devices, using popular embedded(mobile) machine-learning platforms for real-time inference.

Dependencies

• Tensorflow(>=1.14.0), Python 3
• Keras(>=2.2.4), Kito, Scipy, Dlib
• 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
• Download portrait dataset: AISegment

Datasets

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 and model robustness, we perform augmentation like cropping, brightness alteration, flipping, curve filters, motion blur etc.. 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.

AISegment: It is a human matting dataset for binary segmentation of humans and their background. This dataset is currently the largest portrait matting dataset, containing 34,427 images and corresponding matting results. The data set was marked by the high quality of Beijing Play Star Convergence Technology Co., Ltd., and the portrait soft segmentation model trained using this data set has been commercialized.

Dataset Links:-

1. Portseg_128
2. Portrait_256
3. PFCN
4. AISegment
5. Baidu_Aug
6. Supervisely
7. Pascal_Person
8. Supervisely Portrait

Also checkout the datset: UCF Selfie

Annotation Tools

A good dataset is always the first step for coming up with a robust and and accurate model, especially in the case of semantic segmentation. There are many standard datsets available for portrait(or person) segmentation like PFCN, MSCOCO Person, PascalVOV Person, Supervisely etc. But it seems that either the quality or quantity of the images are still insufficient for our use case. So, it would be a good idea to collect custom images for our training process. It is easy to collect images and create ground truth for tasks like classification or object detection; but for semantic segmentation we need to be extra-careful regarding the quality of masks. Also, data collection and annotation takes a lot of time and effort, compared to other computer vision tasks.

Here are some tools for annotation and data collection which i found to be useful in this regard:-

1. Offline Image Editors - Pros: Free, High Accuracy; Cons: Manual Annotation Time; Eg: GIMP, Photoshop etc.
2. Pretrained Models: Pros - Fast, Easy to Use; Cons: Limited Accuracy; Eg: Deeplab Xception, MaskRCNN etc.
3. Online Annotation Tools - Pros: Automated, Easy to Use, Flexible; Cons: Price; Eg: Supervisely, Remove.bg.
4. Crowd Sourcing Tools - Pros: Potential Size and Variety, Less Effort; Cons: Time, Quality; Eg: Amazon MTurk.

If you are planning to use the model on mobile phones specifically for portrait selfies, it would be a good idea to include lots of such portrait images captured using mobile phones in your dataset.

The following are some examples of the tools/models which offers reasonable accuracy and flexibility.

• unscreen.com: An automatic online tool for removing backgrounds from videos(paid).
• MODNet : A real-time portrait video matting model with very high accuracy(open-source).

Mobile-Unet 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.

The other two architectures that we've experimented include mobilenet_v3 and prisma-net ,and their block diagrams are provided in the pictures directory.

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 tflite_webcam.py  # Run the model using tflite interpreter
7.  python segvideo.py test/sunset.jpg # Apply blending filters on video
8.  python bgsub.py # Perform static background subtraction
9.  python portrait_video.py # Use portrait-net for video segmentation
10. python3 tftrt_convert.py # Convert keras model to tf-trt format
11. python3 tftrt_infer.py test/beach.jpg # Perform inference on jetson tx2

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

1. Model Type - 1

Here the inputs and outputs are images of size 128x128.The backbone is mobilenetv2 with depth multiplier 0.5. 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%, FPS: 10-15

2. Model Type - 2

Here the inputs and outputs are images of size 224x224. The backbone is mobilenetv3 with depth multiplier 1.0. 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: 97%, FPS: 10-15

3. Model Type - 3

Here the inputs and outputs are images of size 256x256. The prisma-net architecture is based on unet and uses residual blocks with depthwise separable convolutions instead of regular convolutional blocks(Conv+BN+Relu). Also,it uses elementwise addition instead of feature concatenation in the decoder part.

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%, FPS: 8-10

NB: Accuracy measured on a predefined test data-set and FPS on the android application, using Oneplus3.

Failure cases

When there are objects like clothes, bags etc. in the background the model fails to segment them properly as background, especially if they seem connected to the foreground person. Also if there are variations in lighting or illumination within the image, there seems to be a flickering effect on the video resulting in holes in the foreground object.

Portrait-net

The decoder module consists of refined residual block with depthwise convolution and up-sampling blocks with transpose convolution. Also,it uses elementwise addition instead of feature concatenation in the decoder part. The encoder of the model is mobilnetev2 and it uses a four channel input, unlike the ohter models, for leveraging temporal consistency. As a result, the output video segmentaion appears more stabilized compared to other models. Also, it was observed that depthwise convolution and elementwise addition in decoder greatly improves the speed of the model.

Here is the model summary:-

1. Model 1

• Dataset: Portrait-mix (PFCN+Baidu+Supervisely)
• Device: Oneplus3 (GPU: Adreno 530)
• Size: 224x224

2. Model 2

• Dataset: AISegment
• Device: Redmi Note 8 Pro (GPU: Mali-G76 MC4)
• Size: 256x256

NB: Accuracy measured on a random test data-set and executiom time(GPU) using tflite benchmark tool.

Deeplab, Quantization Aware Training and ML Accelerators

DeepLab is a state-of-art deep learning model for semantic image segmentation. It was originally used in google pixel phones for implementing the portrait mode in their cameras. Later, it was shown to produce great results with popular segmentation datsets like pascal voc, cityscapes, ADE20K etc. Here, we use a deeplab model with mobilenetv2 as backbone for portrait segmentation. The deeplabv3+ model uses features like depthwise convolution, atrous convolution(dilated) and relu6 activation function, for fast and accurate segmentation. Here, we do not use ASPP blocks for training, for the sake of higher inference speeds. The model is trained with output stride 16 and exported for inference with output stride of 8.

Initially, we need to train a non-quantized model with AISegment dataset until convergence. For this purpose, we use a pretrained deeplab model trained on coco and pascal voc with mobilenetv2 backbone and depth multiplier 0.5. For the initial checkpoint, the number of classes was 21 and input size was 513x513. We may reuse all the layer weights from the pretrained model for finetuning; except the logits. In the official training script, re-configure the settings for dataset(AISegment), final logits layer, number of classes(2) and initial checkpoint(dm 0.5). Once the model is completely trained, fine-tune this float model with quantization using a small learning rate. Finally, convert this model to tflite format for inference on android.

The quantized models can be run on android phone using it's CPU; whereas GPU typically reguires a float model. It was found that the float model on gpu and quantized model on cpu (4 threads) took about the same time for inference. Even though the quantized model seems to suffer from small accuracy loss, there were no visible differences on inference. On the other hand, gpu consumes less power than cpu for the same task. If the inference is carried out using gpu rather than cpu asynchronously, the cpu can run other operations in parallel, for maximizig the overall throughput.

A third alternative is to offload the computations to specialized devices like DSP or NPU's. These chips like the Hexagon 685 DSP(Redmi Note 7 Pro), which are sometimes referred to as “neural processing units”, “neural engines”, or “machine learning cores”, are tailored specifically to AI algorithms’ mathematical needs. They’re much more rigid in design than a traditional CPUs, and contain special instructions and arrangements (HVX architecture, Tensor Cores etc.) that accelerate certain scalar and vector operations, which become noticeable in large-scale implementations. They consume less power than CPU's or GPU's and usually runs models in FP16 or UINT8 format in a faster and efficient manner. In android we can run the quantized models on DSP using NNAPI delegate or Hexagon Delegate. The hexagon delegate leverages the Qualcomm Hexagon library to execute quantized kernels on the DSP.The delegate is intended to complement NNAPI functionality, particularly for devices where NNAPI DSP acceleration is unavailable.

Here are the benchmark results:-

Deeplab Training:-

• Tensorflow version: 1.15
• Float Train: 150k epochs,
• Quant Train: 30k
• Eval mIOU: 97%

Deeplab Testing:-

• Device: Redmi Note 7 Pro, Snapdragon 675
• Model: Deeplab_Mobilenetv2_DM05
• Input SIze: 513x513
• Output Classes: 2

The automatic mixed precision feature in TensorFlow, PyTorch and MXNet provides deep learning researcher and engineers with AI training speedups of up to 3X on NVIDIA Volta and Turing GPUs with adding just a few lines of code (Automatic Mixed Precision Training). On recent nvidia GPU's, they use tensor cores with half precision(FP16) computations to speed up the training process, while maintaining the same accuracy. Using other techniques like TensorRT and XLA we can further improve the inference speed on such devices. However, tensor cores which provide mix precision(float and int8), requires certain dimensions of tensors such as dimensions of your dense layer, number of filters in Conv layers, number of units in RNN layer to be a multiple of 8. Also, currently they support very few operations and are still in experimental stage.

Here is a bar chart comparing the INT8 performance of ML accelerators across various platforms.

SlimNet: Real-time Portrait Segmentation on High Resolution Images

Slim-net is a light weight CNN for performing real-time portrait segmentation on mobile devices, using high resolution images. We were able to achieve 99% training accuracy on the aisegment portrait dataset and run the model(1.5MB) on a mid-range android smartphone at 20 fps on deployment. Using the high resolution input images, we were able to preserve fine details and avoid sharp edges on segmentation masks, during inference .The architecture is heavily inspired from the mediapipe hair-segmentation model for android and the tflite model runs on any android device without additional API's.

The following is a brief summary of the architectural features of the model:-

1. The model is based on encoder-decoder architecture and uses PReLU activation throught the network. It hepls us to achieve faster convergence and improved accuracy.

2. The inputs are initially downsampled from a size of 512 to 128 (i,e 1/4'th). This helps us to reduce the overall computation costs; while preseving the details.

3. It uses skip connections between the encoder and decoder blocks (like unet) and helps us to extract fine details and improves gradient flow across layers.

4. Further, it uses bottleneck layers (like resnet) with depthwise convolutions for faster inference.

5. Also, it uses dilated convolution(like deeplab) and helps us to maintain larger receptive field with same computation and memory costs, while also preserving resolution.

6. Finally, the features are upsampled to full resolution(512) with the help of transposed convolutions.

Benchmark and Test Summary:-

• Dataset: AISegment
• Device: Redmi Note 8 Pro (GPU: Mali-G76 MC4)
• Size: 512x512

The slim-net model for portrait segmentation was successfully trained using tensorflow 1.15 and exported to tflite format. The new dataset consist of 55082 images for training and 13770 images for testing. It includes portrait images fromm AISegment dataset and sythetic images with custom backgrounds. This model with input size of 512x512 took about 2 days for training on a GTX 1080 Ti with batch size of 64. Finally, a test accuracy of 99% was obtained on the test-set after 300 epochs, using a minimal learning rate of 1e^-6(after decay).

The model seems to perform well on still images; but on videos in mobile it shows some flickering effect.

SINet: Extreme Lightweight Portrait Segmentation

SINet is an lightweight portrait segmentaion dnn architecture for mobile devices. The model which contains around 86.9 K parameters is able to run at 100 FPS on iphone (input size -224) , while maintaining the accuracy under an 1% margin from the state-of-the-art portrait segmentation method. The proposed portrait segmentation model conatins two new modules for fast and accurate segmentaion viz. information blocking decoder structure and spatial squeeze modules.

1. Information Blocking Decoder: It measures the confidence in a low-resolution feature map, and blocks the influence of high-resolution feature maps in highly confident pixels. This prevents noisy information to ruin already certain areas, and allows the model to focuson regions with high uncertainty.

2. Spatial Squeeze Modules: The S2 module is an efficient multipath network for feature extraction. Existing multi-path structures deal with the various size of long-range dependencies by managing multiple receptive fields. However, this increases latency in real implementations, due to having unsuitable structure with regard to kernel launching and synchronization. To mitigate this problem, they squeeze the spatial resolution from each feature map by average pooling, and show that this is more effective than adopting multi-receptive fields.

Besides the aforementioned features, the SINet architecture uses depthwise separable convolution and PReLU actiavtion in the encoder modules. They also use Squeeze-and-Excitation (SE) blocks that adaptively recalibrates channel-wise feature responses by explicitly modelling interdependencies between channels, for improving the model accuracy. For training, they used cross entropy loss with additional boundary refinement. In general it is faster and smaller than most of the portrait segmentaion models; but in terms of accuracy it falls behind portrait-net model by a small margin. The model seems to be faster than mobilentv3 in iOS; but in android it seems likely to make only a marginal difference(due to optimized tflite swish operator).

We trained the sinet model with aisegment + baidu portrait dataset using input size 320 and cross entropy loss function, for 600 epochs and achieved an mIOU of 97.5%. The combined dataset consists of around 80K images(train+val), after data augmentaion. The final trained model has a size of 480kB and 86.91K parameters. Finally, the pytorch model was exported to ONNX and CoreML formats for mobile deployment.

In practice the model works well with simple portrait images; but for videos with more background regions the model produces artefacts on the output during inference. Unfortunaltely both the original models and aisegment retrained models suffer from this problem, even after acheiving 95% mIOU after training. In the worst case scenario, we may need to run a localizer over the image and crop out the tightly bound roi region containing person before running the segmentation model or apply morphological opening/closing over the output binary mask. But this comes with additional cost and would nullify the overall advantage of the light weight segmenation model.

Quantizing MobilenetV3 Models With Keras API

MobileNetV3 is the next generation of on-device deep vision model from google. It is twice as fast as MobileNetV2 with equivalent accuracy, and advances the state-of-the-art for mobile computer vision networks. Here we use minimalistic version of mobilenetv3 with input size 256 as the encoder part of the network. In the decoder module we use transition blocks along with upsampling layers , similar to the decoder modules in the portrait-net architecture. There are two branches in this block: one branch contains two depthwise separable convolutions and the other contains a single 1×1 convolution to adjust the number of channels. For upsampling we use bilinear resize along with the transition blocks in the decoder module. In the case of skip connections between encoder anmd decoder, we use element-wise addition instead of concatenation for faster inference speed.

During training, initially we freeze all the layers of encoder and train it for 10 epochs. After that, we unfreeze all the layers and train the model for additional 10 epochs. Finally, we perform quantization aware training on the float model, and convert all of the the models to tflite format.

Observations:-

1. Using the pretrained mobilnetv3 as the encoder during training greatly improved the convergence speed. Also, the input images were normalized to [-1.0, 1.0] range before passing to the model. The float model convegerd to 98% validation accuracy within first 20 epochs.
2. Using the latest tensorflow built from source and aisegment dataset with 68852 images, the training process took about 2 hours for completing 20 epochs, on a Tesla P100 GPU in google colaboratory.
3. In the current tensorflow 2.3 and tf model optimization library, some layers like Rescale, Upsampling2D, Conv2DTranspose are not supported by the tf.keras Quantization Aware Training API's. For this purpose you have to install the latest nightly version or build the same from source. Similarly the mobilenetv3 pretrained models are only available in tf-nightly(currently).
4. Using elementwise additon instead of concatenation on skip connection bewteen encoder and decoder greatly helped us to decrease the model size and improve it's inference time.
5. After quantization aware training, even though the model size was reduced by 3x, there was no considerable loss in model accuracy.
6. On POCO X3 android phone, the float model takes around 17ms on CPU and 9ms on it's GPU (>100 FPS), whereas the quantized model takes around 15ms on CPU (2 threads). We were unable to run the fully quantized models(UINT8) using nnapi opr hexagon delegate since some of the layers were not fully supported. However we can run them partially on such accelerators with decreased performance(comparatively).

Android Tflite Benchmark

To measure the performance of a tflite model on android devices, we can use the native binary benchmark tool. It provides a summary of average execution time and memory consumption of individual operators on the device(CPU, GPU and DSP). To benchmark the models on your own android device using a linux system, perform the following steps :-

1. Install adb tool on your system.

sudo apt-get install android-tools-adb android-tools-fastboot.

2. Connect your phone to system and copy the benchmark tool.

adb push benchmark_model_latest /data/local/tmp

3. Make the binary executable.

adb shell chmod +x /data/local/tmp/benchmark_model

4. Copy all the tflite models into the device.

adb push mnv3_post_quant.tflite /data/local/tmp
adb push mnv3_seg_float.tflite /data/local/tmp
adb push mnv3_seg_quant.tflite /data/local/tmp

5. Optionally, copy all the hexagon delegate library files.

adb push libhexagon_interface.so /data/local/tmp
adb push libhexagon_nn_skel*.so /data/local/tmp

6. Finally run the benchmarks on the device.

adb shell /data/local/tmp/benchmark_model_latest \
  --graph=/data/local/tmp/mnv3_seg_float.tflite \
  --num_threads=2 \
  ----enable_op_profiling=true

adb shell /data/local/tmp/benchmark_model_latest \
  --graph=/data/local/tmp/mnv3_seg_quant.tflite \
  --num_threads=2 \
  --enable_op_profiling=true

adb shell /data/local/tmp/benchmark_model_latest \
  --graph=/data/local/tmp/mnv3_seg_float.tflite \
  --use_gpu=true \
  --enable_op_profiling=true

adb shell /data/local/tmp/benchmark_model_latest \
  --graph=/data/local/tmp/mnv3_post_quant.tflite \
  --use_hexagon=true \
  --hexagon_profiling=true \
  --enable_op_profiling=true

Note: The benchmark binary and hexagon library files are stored in the directory - libraries and binaries.

Android Application

SegMe_V0

This version of android demo application uses the nightly experimental gpu delegate for on-device inferencing and GLSurfaceView for displaying the output on screen.

Real-time portrait video in android application

(Shot on OnePlus 3