# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#
#######################################################################################
#
# This project is one of several repositories exploring image segmentation techniques.
# All related projects and interactive demos can be found at:
# https://huggingface.co/spaces/leonelhs/removators
# Self app: https://huggingface.co/spaces/leonelhs/rembg
#
# Source code is based on or inspired by several projects.
# For more details and proper attribution, please refer to the following resources:
#
# - [face-makeup.PyTorch] - [https://github.com/zllrunning/face-makeup.PyTorch]
# - [BiSeNet] [https://github.com/CoinCheung/BiSeNet]

import gradio as gr
import cv2
import torch
import numpy as np
from PIL import Image
from huggingface_hub import hf_hub_download
import torchvision.transforms as transforms
from bisnet import BiSeNet

REPO_ID = "leonelhs/faceparser"
MODEL_NAME = "79999_iter.pth"

model = BiSeNet(n_classes=19)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model_path = hf_hub_download(repo_id=REPO_ID, filename=MODEL_NAME)
model.load_state_dict(torch.load(model_path, map_location=device))
model.eval()

part_colors = [
    {"part": "background", "color": [255, 0, 0]},
    {"part": "face",       "color": [219,  79,  66]},
    {"part": "right_brow", "color": [255, 170,   0]},
    {"part": "left_brow",  "color": [255,   0,  85]},
    {"part": "right_eye",  "color": [255,   0, 170]},
    {"part": "left_eye",   "color": [  0, 255,   0]},
    {"part": "glasses",    "color": [ 85, 255,   0]},
    {"part": "right_ear",  "color": [170, 255,   0]},
    {"part": "left_ear",   "color": [  0, 255,  85]},
    {"part": "earrings",   "color": [  0, 255, 170]},
    {"part": "nose",       "color": [  0,   0, 255]},
    {"part": "teeth",      "color": [ 85,   0, 255]},
    {"part": "upper_lip",  "color": [170,   0, 255]},
    {"part": "lower_lip",  "color": [  0,  85, 255]},
    {"part": "neck",       "color": [  0, 170, 255]},
    {"part": "collar",     "color": [255, 255,   0]},
    {"part": "cloths",      "color": [255, 255,  85]},
    {"part": "hair",       "color": [199, 21,  133]},
    {"part": "crown",      "color": [255,   0, 255]},
    {"part": "extra20",    "color": [255,  85, 255]},
    {"part": "extra21",    "color": [255, 170, 255]},
    {"part": "extra22",    "color": [  0, 255, 255]},
    {"part": "extra23",    "color": [ 85, 255, 255]},
    {"part": "extra24",    "color": [170, 255, 255]},
]

def image_to_tensor(image):
    return transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
    ])(image)

def parse_face(mask):

    num_of_class = np.max(mask)
    face_parts = []

    for index in range(1, num_of_class + 1):
        face_part = np.where(mask == index)
        canvas = np.full((512, 512, 3), 255, dtype=np.uint8)
        canvas[face_part[0], face_part[1], :] = part_colors[index]["color"]
        canvas = cv2.cvtColor(canvas, cv2.COLOR_BGR2GRAY)
        face_parts.append((canvas, part_colors[index]["part"]))

    return face_parts

def predict(image):
    with torch.no_grad():
        image = image.resize((512, 512), Image.Resampling.BILINEAR)
        input_tensor = image_to_tensor(image)
        input_tensor = torch.unsqueeze(input_tensor, 0)
        if torch.cuda.is_available():
            input_tensor = input_tensor.cuda()
        mask = model(input_tensor)[0]
        mask = mask.squeeze(0).cpu().numpy().argmax(0)
        sections = parse_face(mask)
        return image, sections


aboutme = r"""
# PyTorch Image Face Parser

Extracts facial features (hair, nose, eyes, etc.) from images using image segmentation.

This project is part of a larger collection of repositories exploring image segmentation techniques.  
Related projects and interactive demos are available at: [Removators](https://huggingface.co/spaces/leonelhs/removators)

## Acknowledgments
The source code is based on or inspired by the following projects:
- [face-makeup.PyTorch](https://github.com/zllrunning/face-makeup.PyTorch)
- [BiSeNet](https://github.com/CoinCheung/BiSeNet)

## Contact
For questions, comments, or feedback, please contact:  
📧 leonelhs@gmail.com

"""

with gr.Blocks(title="Face Parser") as app:
    navbar = gr.Navbar(visible=True, main_page_name="Workspace")
    gr.Markdown("## Face Parser Tool")
    with gr.Row():
        with gr.Column(scale=1):
            inp = gr.Image(type="pil", label="Upload Image")
            btn_predict = gr.Button("Parse")
        with gr.Column(scale=2):
            out = gr.AnnotatedImage(label="Face parsed annotated")

    btn_predict.click(predict, inputs=[inp], outputs=[out])


with app.route("About this", "/about"):
    gr.Markdown(aboutme)

app.launch(share=False, debug=True, show_error=True, mcp_server=True, pwa=True)
app.queue()