msterdc / human-segmentation-pytorch

Human segmentation models, training/inference code, and trained weights, implemented in PyTorch

Human-Segmentation-PyTorch

Human segmentation models, training/inference code, and trained weights, implemented in PyTorch.

Supported networks

• UNet: backbones MobileNetV2 (all aphas and expansions), ResNetV1 (all num_layers)
• DeepLab3+: backbones ResNetV1 (num_layers=18,34,50,101), VGG16_bn
• BiSeNet: backbones ResNetV1 (num_layers=18)
• PSPNet: backbones ResNetV1 (num_layers=18,34,50,101)
• ICNet: backbones ResNetV1 (num_layers=18,34,50,101)

To assess architecture, memory, forward time (in either cpu or gpu), numper of parameters, and number of FLOPs of a network, use this command:

python measure_model.py

Dataset

Portrait Segmentation (Human/Background)

• Automatic Portrait Segmentation for Image Stylization: 1800 images
• Supervisely Person: 5711 images

Set

• Python3.6.x is used in this repository.
• Clone the repository:

git clone --recursive https://github.com/AntiAegis/Human-Segmentation-PyTorch.git
cd Human-Segmentation-PyTorch
git submodule sync
git submodule update --init --recursive

• To install required packages, use pip:

workon humanseg
pip install -r requirements.txt
pip install -e models/pytorch-image-models

Training

• For training a network from scratch, for example DeepLab3+, use this command:

python train.py --config config/config_DeepLab.json --device 0

where config/config_DeepLab.json is the configuration file which contains network, dataloader, optimizer, losses, metrics, and visualization configurations.

• For resuming training the network from a checkpoint, use this command:

python train.py --config config/config_DeepLab.json --device 0 --resume path_to_checkpoint/model_best.pth

• One can open tensorboard to monitor the training progress by enabling the visualization mode in the configuration file.

Inference

There are two modes of inference: video and webcam.

python inference_video.py --watch --use_cuda --checkpoint path_to_checkpoint/model_best.pth
python inference_webcam.py --use_cuda --checkpoint path_to_checkpoint/model_best.pth

Benchmark

• Networks are trained on a combined dataset from the two mentioned datasets above. There are 6627 training and 737 testing images.
• Input size of model is set to 320.
• The CPU and GPU time is the averaged inference time of 10 runs (there are also 10 warm-up runs before measuring) with batch size 1.
• The mIoU is measured on the testing subset (737 images) from the combined dataset.
• Hardware configuration for benchmarking:

CPU: Intel(R) Core(TM) i7-7700HQ CPU @ 2.80GHz
GPU: GeForce GTX 1050 Mobile, CUDA 9.0

Model	Parameters	FLOPs	CPU time	GPU time	mIoU
UNet_MobileNetV2 (alpha=1.0, expansion=6)	4.7M	1.3G	167ms	17ms	91.37%
UNet_ResNet18	16.6M	9.1G	165ms	21ms	90.09%
DeepLab3+_ResNet18	16.6M	9.1G	133ms	28ms	91.21%
BiSeNet_ResNet18	11.9M	4.7G	88ms	10ms	87.02%
PSPNet_ResNet18	12.6M	20.7G	235ms	666ms	---
ICNet_ResNet18	11.6M	2.0G	48ms	55ms	86.27%