thangtran480 / hair-segmentation

hair segmentation in mobile device

Hair Segmentation Realtime using Keras

The architecture was inspired by Real-time deep hair matting on mobile devices

Prerequisites

python 3.6 tensorflow-gpu==1.13.1, opencv-python==4.1.0.25, Keras==2.2.4, numpy==1.16.4, scikit-image==0.15.0

install environment in conda:

conda env create -f environment.yml

Dataset

• CelebAMask-HQ (contain 29.300 images and masks hair segmentation)
• Figaro-1k
• Lft

Download data CelebAMask-HQ and use preprocess in ./data/pre-process-data-CelebAMask-HQ.ipynb to create dataset

Data structure training

├── my-data
│   ├── train
│   │   ├──image
│   │   │   ├── 1.jpg
│   │   │   ├── 2.jpg
│   │   │   ├── 3.jpg
...
│   │   ├──mask
│   │   │   ├── 1.jpg
│   │   │   ├── 2.jpg
│   │   │   ├── 3.jpg
...
│   ├── val
...
│   ├── test
...

Train model

python train.py [--datadir PATH_FOLDER] [--batch_size BATCH_SIZE] [epochs EPOCHS] [--lr LEARNING_RATE] []

optional arguments:
    --datadir:        path to folder dataset, default ./data
    --batch_size:     batch size training, default 4
    --epochs:         number of eposchs, default 5
    --lr:             learning rate, default 1e-4
    --image_size:     size image input, default (224, 224)
    --use_pretrained: use pretrained, default false
    --path_model:     directory is saved checkpoint, default ./checkpoints
    --device:         device training model, default 0 (GPU:0), 1(GPU:1), -1(CPU)          

Evaluate model

python evaluate.py

Run pretrain model

# Run test.py
python demo.py

You will see the predicted results of test image in test/data

Result

Convert to Tensorflow Lite

• Convert

# Convert Model to Mobile
python convert_to_tflite.py

• Show shape model tflite

# Shape input and output shape model tflite 
python shape_input_output_tflite.py

About Keras

Keras is a minimalist, highly modular neural networks library, written in Python and capable of running on top of either TensorFlow or Theano. It was developed with a focus on enabling fast experimentation. Being able to go from idea to result with the least possible delay is key to doing good research.

Use Keras if you need a deep learning library that:

allows for easy and fast prototyping (through total modularity, minimalism, and extensibility). supports both convolutional networks and recurrent networks, as well as combinations of the two. supports arbitrary connectivity schemes (including multi-input and multi-output training). runs seamlessly on CPU and GPU. Read the documentation Keras.io

Keras is compatible with: Python 3.6.

TODO

• Implement model using Keras
• Convert model to Tensorflow Lite
• Implement model to Android (DOING)

License

Copyright (c) 2019 Thang Tran Van

Licensed under the MIT License. You may not use this file except in compliance with the License

hair-segmentation's Issues

Models are missing

Hi,

Thank you for your hard work, but in the commit of
"add train model with GPU
@thangtran480
thangtran480 committed on 6 Jun"

you deleted the folder called models with its files, can you reupload them in the latest version please ? :)

ml model is not working in android

hi dev.. how to get mask image from this ml model.pls give some advice

Is your pretrained model weights(hairnet_matting_30.hdf5) the best model or we need further training?

I need more quality, if I train further, will I get better model weights?

color of hair

I want to classify hair color.

My first naive solution is cv2.mean(image, hair_mask) and find similarity to set of predefined colors.

What do you think about?

ValueError: zero-size array to reduction operation maximum which has no identity

Hi @thangtran480
Every time I ran python train.py, I jus get the error "ValueError: zero-size array to reduction operation maximum which has no identity".

Train all CelebAMask-HQ photos takes several hours, so I reduced the number of images to identify something wrong, but nothing.

I made data struture like this:

Are there any more steps I need to take?

Android sample not work

After select an image app will hang and get an anr warning.

Deploy model to Android with Tensorflow Lite

You save the pixel in a text file. I set these pixels in a bitmap, but it shows a black screen. How can I solve this problem?

Android model?

Do you happen to know when you will release an Android model and demo? That could do hair segmentation in realtime.

How to apply color on Hair

I downloaded the source code, but when I pick an image from the gallery it is taking too much time to respond, and the other thing is there is no option to add color to Hair,
Please guide ...

Extract hair region

How to extract the hair region from the original image after the hair is segmented? Once we predict the hair segmentation part, we are getting results like the purple color on hair which shows the detection part. but I want only segmented hair (with the original color of hair) and remove face and background. Can you guide me on how to manipulate pixels with OpenCV? @thangtran480 @gungui98

Data to train

Hello @thangtran480
Your project seems to be very nice!
I've tried to train my own model there is no problem by using Figaro1K, but how did you use lfw images ?
Thanks :)

iOS example would be great, unable to load tflite model on iOS with Tensorflow-lite

Was trying to load the converted tflite model with the example here:
Most are using pod file with 1.13.1 version of Tensorflow-lite.
Currently it looks like 1.14+ is a requirement for your model.
https://github.com/toniz/deeplab-on-ios

Bad result

result is so bad

how to prepare image to predict mask

Hello.

As I see model learned on aligned faces.
All faces in center of image and most of it is normalized (see in one direction).

But what we need to do with any photo with face?
Do we need find, rotate, crop area of face before predict?
Or this library do it by themself?

Use of multiple color

Hello,

Thanks for sharing this. We found your code very useful. Is it possible to have a different color other than purple for hair? Where can I find the exact piece of code for the same?

Regards,
Vidhi