init app

.gitignore

@@ -0,0 +1,2 @@
+.idea/
+__pycache__/

README.md

@@ -1,5 +1,5 @@
 ---
-title: Faceparser
+title: Face Parser
 emoji: 👁
 colorFrom: blue
 colorTo: green

app.py

@@ -0,0 +1,79 @@
+import os
+
+import gradio as gr
+import numpy as np
+import torch
+from PIL import Image
+from bisnet import BiSeNet
+from huggingface_hub import snapshot_download
+
+from utils import vis_parsing_maps, decode_segmentation_masks, image_to_tensor
+
+os.system("pip freeze")
+
+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")
+snapshot_folder = snapshot_download(repo_id=REPO_ID)
+model_path = os.path.join(snapshot_folder, MODEL_NAME)
+model.load_state_dict(torch.load(model_path, map_location=device))
+model.eval()
+
+
+def makeOverlay(image, mask):
+    prediction_mask = np.asarray(mask)
+    image = image.resize((512, 512), Image.BILINEAR)
+    dark_map, overlay = vis_parsing_maps(image, prediction_mask)
+    colormap = decode_segmentation_masks(dark_map)
+    return overlay, colormap
+
+
+def predict(image):
+    with torch.no_grad():
+        image = image.resize((512, 512), Image.BILINEAR)
+        input_tensor = image_to_tensor(image)
+        input_tensor = torch.unsqueeze(input_tensor, 0)
+        if torch.cuda.is_available():
+            input_tensor = input_tensor.cuda()
+        output = model(input_tensor)[0]
+        return output.squeeze(0).cpu().numpy().argmax(0)
+
+
+def inference(image):
+    mask = predict(image)
+    overlay, colormap = makeOverlay(image, mask)
+    return overlay
+
+
+title = "Face Parser"
+description = r"""
+## Image face parser for research
+
+This is an implementation of <a href='https://github.com/zllrunning/face-parsing.PyTorch' target='_blank'>face-parsing.PyTorch</a>.
+It has no any particular purpose than start research on AI models.
+
+"""
+
+article = r"""
+Questions, doubts, comments, please email 📧 `leonelhs@gmail.com`
+
+This demo is running on a CPU, if you like this project please make us a donation to run on a GPU or just give us a <a href='https://github.com/leonelhs/zeroscratches/' target='_blank'>Github ⭐</a>
+
+<a href="https://www.buymeacoffee.com/leonelhs"><img src="https://img.buymeacoffee.com/button-api/?text=Buy me a coffee&emoji=&slug=leonelhs&button_colour=FFDD00&font_colour=000000&font_family=Cookie&outline_colour=000000&coffee_colour=ffffff" /></a>
+
+<center><img src='https://visitor-badge.glitch.me/badge?page_id=zeroscratches.visitor-badge' alt='visitor badge'></center>
+"""
+
+demo = gr.Interface(
+    inference, [
+        gr.Image(type="pil", label="Input"),
+    ], [
+        gr.Image(type="numpy", label="Image face parsed")
+    ],
+    title=title,
+    description=description,
+    article=article)
+
+demo.queue().launch()

requirements.txt

@@ -0,0 +1,5 @@
+torch>=2.0.1
+torchvision~=0.15.2
+pillow~=9.5.0
+bisnet~=1.0.1
+opencv-python

utils.py

@@ -0,0 +1,50 @@
+import cv2
+import numpy as np
+import torchvision.transforms as transforms
+
+# Colors for all 20 parts
+part_colors = [[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 0, 85], [255, 0, 170],
+               [0, 255, 0], [85, 255, 0], [170, 255, 0], [0, 255, 85], [0, 255, 170],
+               [0, 0, 255], [85, 0, 255], [170, 0, 255], [0, 85, 255], [0, 170, 255],
+               [255, 255, 0], [255, 255, 85], [255, 255, 170], [255, 0, 255], [255, 85, 255],
+               [255, 170, 255], [0, 255, 255], [85, 255, 255], [170, 255, 255]]
+
+colormap = np.array(part_colors, dtype=np.uint8)
+
+
+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 decode_segmentation_masks(mask, n_classes=20):
+    red = np.zeros_like(mask).astype(np.uint8)
+    green = np.zeros_like(mask).astype(np.uint8)
+    blue = np.zeros_like(mask).astype(np.uint8)
+    for chanel in range(0, n_classes):
+        idx = mask == chanel
+        red[idx] = colormap[chanel, 0]
+        green[idx] = colormap[chanel, 1]
+        blue[idx] = colormap[chanel, 2]
+    return np.stack([red, green, blue], axis=2)
+
+
+def vis_parsing_maps(image: np.array, parsing_anno, stride=1):
+    image = np.array(image)
+    vis_im = image.copy().astype(np.uint8)
+    vis_parsing_anno = parsing_anno.copy().astype(np.uint8)
+    vis_parsing_anno = cv2.resize(vis_parsing_anno, None, fx=stride, fy=stride, interpolation=cv2.INTER_NEAREST)
+    vis_parsing_anno_color = np.zeros((vis_parsing_anno.shape[0], vis_parsing_anno.shape[1], 3)) + 255
+
+    num_of_class = np.max(vis_parsing_anno)
+
+    for pi in range(1, num_of_class + 1):
+        index = np.where(vis_parsing_anno == pi)
+        vis_parsing_anno_color[index[0], index[1], :] = part_colors[pi]
+
+    vis_parsing_anno_color = vis_parsing_anno_color.astype(np.uint8)
+    vis_im = cv2.addWeighted(cv2.cvtColor(vis_im, cv2.COLOR_RGB2BGR), 0.4, vis_parsing_anno_color, 0.6, 0)
+
+    return vis_parsing_anno, vis_im