README.md

metadata

title: Face Segmentation
emoji: ⚡
colorFrom: indigo
colorTo: red
sdk: streamlit
sdk_version: 1.44.1
app_file: app.py
pinned: false
license: mit
short_description: AI-powered application for face segmentation

Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

Face Segmentation Tool

A deep learning-based tool for precise face segmentation using BiSeNet trained on the CelebAMask-HQ dataset. This tool extracts faces from images with a transparent background, perfect for creating profile pictures, avatars, or creative photo editing.

Features

• Accurate face and hair segmentation with transparent background
• 19-class facial attribute segmentation (skin, eyes, eyebrows, nose, lips, hair, etc.)
• User-friendly Streamlit web interface
• MediaPipe face detection to identify and focus on faces
• Support for downloading the segmented result

Technical Details

This project uses:

• BiSeNet (Bilateral Segmentation Network) for semantic segmentation
• CelebAMask-HQ dataset trained model with 19 facial attribute classes
• MediaPipe for initial face detection and bounding box estimation
• PyTorch for the deep learning components
• Streamlit for the web interface

Installation

Prerequisites

• Python 3.7 or newer
• CUDA-compatible GPU (optional, but recommended for faster processing)

Setup

1. Clone the repository:

   git clone https://github.com/yourusername/face-segmentation.git
   cd face-segmentation
   

2. Create and activate a virtual environment (recommended):

   # On Windows
   python -m venv venv
   venv\Scripts\activate
   
   # On macOS/Linux
   python -m venv venv
   source venv/bin/activate
   

3. Install the required dependencies:

   pip install -r requirements.txt
   

4. Ensure you have the BiSeNet model weights file (bisenet.pth) in the project root directory. (The file should already be included in the repository)

Usage

1. Start the Streamlit app:

   streamlit run app.py
   

2. Open your web browser and go to the URL shown in the console (typically http://localhost:8501)

3. Upload an image with a face

4. Click the "Segment Face" button

5. View and download the segmented result

Class Labels in CelebAMask-HQ

The model recognizes 19 different facial attributes:

ID	Class	Description
0	background	Non-face background areas
1	skin	Face skin
2	nose	Nose
3	eye_g	Eyeglasses
4	l_eye	Left eye
5	r_eye	Right eye
6	l_brow	Left eyebrow
7	r_brow	Right eyebrow
8	l_ear	Left ear
9	r_ear	Right ear
10	mouth	Mouth
11	u_lip	Upper lip
12	l_lip	Lower lip
13	hair	Hair
14	hat	Hat
15	ear_r	Ear rings
16	neck_l	Neck area
17	neck	Neck
18	cloth	Clothing

How It Works

1. The app uses MediaPipe to detect faces in the uploaded image
2. It crops and processes the face region using BiSeNet
3. BiSeNet performs semantic segmentation to classify each pixel into one of 19 classes
4. Selected facial features are preserved while background, neck, and clothes are made transparent
5. The result is an RGBA image with the face and hair intact and a transparent background

Customization

You can modify which facial attributes to keep by adjusting the keep_classes list in the FaceHairSegmenter class:

# Current configuration - keeps all face parts except background, clothes, and neck
self.keep_classes = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 15, 17, 18]

Troubleshooting

• No face detected: Ensure the image contains a clearly visible face.
• Multiple faces detected: The app works best with a single face per image.
• Poor segmentation: For best results, use images with good lighting and a clear face.
• CUDA out of memory: Try using a smaller image or run on CPU if your GPU has limited memory.
• PyTorch class path error: If you encounter an error like "Tried to instantiate class 'path._path', but it does not exist!", try updating your PyTorch version to 1.9.0 or newer using pip install torch==1.9.0 torchvision==0.10.0. This is a known issue with certain PyTorch versions when loading models.

PyTorch Version Compatibility

This project is tested with PyTorch 1.7.0 to 1.13.0. If you encounter model loading issues with newer PyTorch versions (2.x), try downgrading to PyTorch 1.13.0:

pip install torch==1.13.0 torchvision==0.14.0

References

• BiSeNet: Bilateral Segmentation Network for Real-time Semantic Segmentation
• CelebAMask-HQ Dataset

License

This project is licensed under the MIT License - see the LICENSE file for details.