app.py

#!/usr/bin/env python

from __future__ import annotations

import os
import pathlib
import sys

import gradio as gr
import numpy as np
import torch

sys.path.insert(0, "face_detection")
sys.path.insert(0, "face_parsing")
sys.path.insert(0, "roi_tanh_warping")

from ibug.face_detection import RetinaFacePredictor
from ibug.face_parsing.parser import WEIGHT, FaceParser
from ibug.face_parsing.utils import label_colormap

DESCRIPTION = "# [hhj1897/face_parsing](https://github.com/hhj1897/face_parsing)"


def load_model(model_name: str, device: torch.device) -> FaceParser:
    encoder, decoder, num_classes = model_name.split("-")
    num_classes = int(num_classes)  # type: ignore
    model = FaceParser(device=device, encoder=encoder, decoder=decoder, num_classes=num_classes)
    model.num_classes = num_classes
    return model


device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

detector = RetinaFacePredictor(threshold=0.8, device=device, model=RetinaFacePredictor.get_model("mobilenet0.25"))

model_names = list(WEIGHT.keys())
models = {name: load_model(name, device=device) for name in model_names}


def predict(image: np.ndarray, model_name: str, max_num_faces: int) -> np.ndarray:
    model = models[model_name]
    colormap = label_colormap(model.num_classes)

    # RGB -> BGR
    image = image[:, :, ::-1]

    faces = detector(image, rgb=False)
    if len(faces) == 0:
        raise RuntimeError("No face was found.")
    faces = sorted(list(faces), key=lambda x: -x[4])[:max_num_faces][::-1]
    masks = model.predict_img(image, faces, rgb=False)

    mask_image = np.zeros_like(image)
    for mask in masks:
        temp = colormap[mask]
        mask_image[temp > 0] = temp[temp > 0]

    res = image.astype(float) * 0.5 + mask_image[:, :, ::-1] * 0.5
    res = np.clip(np.round(res), 0, 255).astype(np.uint8)
    return res[:, :, ::-1]


with gr.Blocks(css="style.css") as demo:
    gr.Markdown(DESCRIPTION)
    with gr.Row():
        with gr.Column():
            image = gr.Image(type="numpy", label="Input")
            model_name = gr.Radio(choices=model_names, type="value", value=model_names[1], label="Model")
            max_num_faces = gr.Slider(minimum=1, maximum=20, step=1, value=10, label="Max Number of Faces")
            run_button = gr.Button()
        with gr.Column():
            result = gr.Image(label="Output")
    gr.Examples(
        examples=[[path.as_posix(), model_names[1], 10] for path in pathlib.Path("images").rglob("*.jpg")],
        inputs=[image, model_name, max_num_faces],
        outputs=result,
        fn=predict,
        cache_examples=os.getenv("CACHE_EXAMPLES") == "1",
    )

    run_button.click(
        fn=predict,
        inputs=[image, model_name, max_num_faces],
        outputs=result,
        api_name="predict",
    )

if __name__ == "__main__":
    demo.queue(max_size=20).launch()