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% |
157 Dec 11, 2022
66 Dec 26, 2022
1k Dec 13, 2022
2.5k Jan 6, 2023
202 Jan 6, 2023
2.8k Dec 29, 2022
19 Dec 10, 2022
2.4k Jan 5, 2023
9 Nov 4, 2022
1.1k Dec 24, 2022
18 Oct 6, 2022
53 Jul 4, 2022
148 Dec 21, 2022
78 Dec 27, 2022
419 Jan 3, 2023
122 Dec 7, 2022
82 Nov 29, 2022
165 Nov 22, 2022