mobile-hair-segmentation-pytorch

This repository is part of a program for previewing your own dyeing on mobile device. To do this, you need to separate the hair from the head. And we have to use as light model as MobileNet to use in mobile device in real time. So we borrowed the model structure from the following article.

Real-time deep hair matting on mobile devices

model architecture

This model MobileNet + SegNet.
To do semantic segmentation they transform MobileNet like SegNet. And add additional loss function to capture fine hair texture.

install requirements

pip install -r requirements.txt

model performance (on CPU)

Tip

if you don't apply layer fusion some part, you can get better performance
(delete part on model1 delete part on model2)

preparing datsets

make directory like this

dataset
   |__ images
   |
   |__ masks
   

expected image name
The name of the expected image pair is:

 - dataset/images/1.jpg 
| 
 - dataset/masks/1.jpg  

/dataset
    /images
        /1.jpg
        /2.jpg
        /3.jpg 
         ...
    /masks
        /1.jpg
        /2.jpg
        /3.jpg 
         ...

how to train

after 200 epoch, add other commented augmentation and remove resize
(dataloader/dataloader.py)

There are modelv1 and modelv2 whose backbone are mobilenetv1 and mobilenetv2 each other. Default is mobilenetv2

python main.py --num_epoch [NUM_EPOCH] --model_version [1~2]

If you want to quantize model

python main.py --num_epoch [NUM_EPOCH] --model_version [1~2] --quantize

Or if you want to resume model training

python main.py --num_epoch [NUM_EPOCH] --model_version [1~2] --resume

python main.py --num_epoch [NUM_EPOCH] --model_version [1~2] --model_path [MODEL_PATH]

Test

python webcam.py --model_path [MODEL_PATH]

if you want to use quantized model

python webcam.py --model_path [MODEL_PATH] --quantize

Load model

All you have to do is copy ./model and use it

ex)

import torch
from models import quantized_modelv2

quantize = False
device = torch.device("cuda:0" if torch.cuda.is_available() and not quantize else "cpu")
quantized_modelv2(pretrained=True, device=device).to(device)

Deploy TensorRT

Run docker

The dependency libraries in the container can be found in the release notes.

docker pull nvcr.io/nvidia/tensorrt:<xx.xx>-py<x>
docker run --gpus all -it --rm -v local_dir:container_dir nvcr.io/nvidia/tensorrt:<xx.xx>-py<x>

run torch2tensorrt.py

pip install -r requirements.txt
pip install torch-tensorrt -f https://github.com/pytorch/TensorRT/releases
python torch2tensorrt.py -model_version [1~2]

TensorRt performance