parce whit annotations

.gitignore

@@ -2,4 +2,8 @@
 __pycache__/
 .gradio
 playground.py
+makeup.py
+test.py
+parsing_map_on_im.jpg
+parsing_map_on_im.png
 resnet18-5c106cde.pth

app.py

@@ -22,13 +22,13 @@
 # - [BiSeNet] [https://github.com/CoinCheung/BiSeNet]
 
 import gradio as gr
-import numpy as np
+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
-from utils import vis_parsing_maps, decode_segmentation_masks, image_to_tensor
 
 REPO_ID = "leonelhs/faceparser"
 MODEL_NAME = "79999_iter.pth"
@@ -40,30 +40,64 @@ 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 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 makeMask(image):
+def predict(image):
     with torch.no_grad():
-        image = image.resize((512, 512), Image.BILINEAR)
+        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()
-        output = model(input_tensor)[0]
-        return output.squeeze(0).cpu().numpy().argmax(0)
-
-
-def predict(image):
-    mask = makeMask(image)
-    overlay, colormap = makeOverlay(image, mask)
-    return overlay
+        mask = model(input_tensor)[0]
+        mask = mask.squeeze(0).cpu().numpy().argmax(0)
+        sections = parse_face(mask)
+        return image, sections
 
 
 aboutme = r"""
@@ -93,7 +127,7 @@ with gr.Blocks(title="Face Parser") as app:
             inp = gr.Image(type="pil", label="Upload Image")
             btn_predict = gr.Button("Parse")
         with gr.Column(scale=2):
-            out = gr.Image(type="pil", label="Output image")
+            out = gr.AnnotatedImage(label="Face parsed annotated")
 
     btn_predict.click(predict, inputs=[inp], outputs=[out])
 
@@ -101,4 +135,5 @@ with gr.Blocks(title="Face Parser") as app:
 with app.route("About this", "/about"):
     gr.Markdown(aboutme)
 
-app.launch()
+app.launch(share=False, debug=True, show_error=True, mcp_server=True, pwa=True)
+app.queue()

utils.py

@@ -1,50 +0,0 @@
-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