Spaces:
Running
Running
| 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: | |
| ```bash | |
| git clone https://github.com/yourusername/face-segmentation.git | |
| cd face-segmentation | |
| ``` | |
| 2. Create and activate a virtual environment (recommended): | |
| ```bash | |
| # 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: | |
| ```bash | |
| 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: | |
| ```bash | |
| 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: | |
| ```python | |
| # 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: | |
| ```bash | |
| pip install torch==1.13.0 torchvision==0.14.0 | |
| ``` | |
| ## References | |
| - [BiSeNet: Bilateral Segmentation Network for Real-time Semantic Segmentation](https://arxiv.org/abs/1808.00897) | |
| - [CelebAMask-HQ Dataset](https://github.com/switchablenorms/CelebAMask-HQ) | |
| ## License | |
| This project is licensed under the MIT License - see the LICENSE file for details. | |