peterzhousz / hairnets

Hair Segmentation and Classification with Unet and GoogleNet

Hair segmentation and classification with Unet and GoogleNet

This repository contains the implementation of a deep learning algorithm to classify hair types from images. It consists of two separate CNNs:

• One CNN is used to segment hair in face images. This is a binary classification task: the neural network predicts if each pixel in the image is either hair or non-hair. This neural network structure is derived from the U-Net architecture, described in this paper. The performance of this segmentation network is tested on the LFW | Part Labels Database and achieve an accuracy of 92%, that is the best score from papers we have read so far.

• One other CNN is used to classify hair segment into type a, b or c. This is a GoogleNet architecture.

The folder ./weights contains the pre-trained weights for segmentation and classification. The folder ./libs contains the functions used by main files.

Requirements:

• Tensorflow >= 1.12
• Keras
• Skimage
• Opencv
• PIL >= 1.1.7

Part I: Segmentation

-You need to create a folder datasets and insert three folders for the 'funneled images', 'Ground Truth Images' and 'Ground Truth Labels' that you will download from this link.

• Then run the file create_dataset.py to create and process the training data.
• Run train.segmentation.py to train the network for segmentation
• Run test_segmentation.py to test the segmentation on test images. The test images should be 224x224x3 and you need to store hair segment in a folder for data augmentation and hair classification

Part II: Data Augmentation

Use data_augmentation.py and the hair segment obtained from Part I to apply random transformations and increase the volume of hair segments. Store the files in a folder

Part II: Classification

Use train_classification.py to train the network for classifying hair type. You need to specify the location of each hair type folder.